[SPARK-15988] [SQL] Implement DDL commands: Create/Drop temporary macro

sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4

@@ -97,6 +97,9 @@ statement
     | CREATE TEMPORARY? FUNCTION qualifiedName AS className=STRING
         (USING resource (',' resource)*)?                              #createFunction
     | DROP TEMPORARY? FUNCTION (IF EXISTS)? qualifiedName              #dropFunction
+    | CREATE TEMPORARY MACRO macroName=identifier
+        '(' colTypeList? ')' expression                        #createMacro
+    | DROP TEMPORARY MACRO (IF EXISTS)? macroName=identifier           #dropMacro
     | EXPLAIN (LOGICAL | FORMATTED | EXTENDED | CODEGEN)? statement    #explain
     | SHOW TABLES ((FROM | IN) db=identifier)?
         (LIKE? pattern=STRING)?                                        #showTables

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
+import scala.collection.mutable.HashSet
 import scala.language.existentials
 import scala.reflect.ClassTag
 import scala.util.{Failure, Success, Try}
@@ -59,6 +60,10 @@ trait FunctionRegistry {
   /** Checks if a function with a given name exists. */
   def functionExists(name: String): Boolean = lookupFunction(name).isDefined
 
+  def registerMacro(name: String, info: ExpressionInfo, builder: FunctionBuilder): Unit
+
+  def dropMacro(name: String): Boolean
+
   /** Clear all registered functions. */
   def clear(): Unit
 
@@ -69,6 +74,8 @@ class SimpleFunctionRegistry extends FunctionRegistry {
   protected val functionBuilders =
     StringKeyHashMap[(ExpressionInfo, FunctionBuilder)](caseSensitive = false)
 
+  val macros = new HashSet[String]
+
   override def registerFunction(
       name: String,
       info: ExpressionInfo,
@@ -101,8 +108,26 @@ class SimpleFunctionRegistry extends FunctionRegistry {
     functionBuilders.remove(name).isDefined
   }
 
+  override def registerMacro(
+      name: String,
+      info: ExpressionInfo,
+      builder: FunctionBuilder): Unit = synchronized {
+    functionBuilders.put(name, (info, builder))
+    macros += name.toLowerCase()
+  }
+
+  override  def dropMacro(name: String): Boolean = synchronized {
+    if (macros.contains(name.toLowerCase)) {
+      macros -= name.toLowerCase
+      functionBuilders.remove(name).isDefined
+    } else {
+      false
+    }
+  }
+
   override def clear(): Unit = synchronized {
     functionBuilders.clear()
+    macros.clear()
   }
 
   def copy(): SimpleFunctionRegistry = synchronized {
@@ -144,6 +169,14 @@ object EmptyFunctionRegistry extends FunctionRegistry {
     throw new UnsupportedOperationException
   }
 
+  override def registerMacro(name: String, info: ExpressionInfo, builder: FunctionBuilder): Unit = {
+    throw new UnsupportedOperationException
+  }
+
+  override def dropMacro(name: String): Boolean = {
+    throw new UnsupportedOperationException
+  }
+
   override def clear(): Unit = {
     throw new UnsupportedOperationException
   }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/NoSuchItemException.scala

@@ -52,3 +52,6 @@ class NoSuchPartitionsException(db: String, table: String, specs: Seq[TableParti
 
 class NoSuchTempFunctionException(func: String)
   extends AnalysisException(s"Temporary function '$func' not found")
+
+class NoSuchTempMacroException(func: String)
+  extends AnalysisException(s"Temporary macro '$func' not found")

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala

@@ -758,6 +758,24 @@ class SessionCatalog(
     }
   }
 
+  /** Create a temporary macro. */
+  def createTempMacro(
+      name: String,
+      info: ExpressionInfo,
+      funcDefinition: FunctionBuilder): Unit = {
+    if (functionRegistry.functionExists(name)) {
+      throw new TempFunctionAlreadyExistsException(name)
+    }
+    functionRegistry.registerMacro(name, info, funcDefinition)
+  }
+
+  /** Drop a temporary macro. */
+  def dropTempMacro(name: String, ignoreIfNotExists: Boolean): Unit = {
+    if (!functionRegistry.dropMacro(name) && !ignoreIfNotExists) {
+      throw new NoSuchTempMacroException(name)
+    }
+  }
+
   protected def failFunctionLookup(name: String): Nothing = {
     throw new NoSuchFunctionException(db = currentDb, func = name)
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala

@@ -25,14 +25,16 @@ import org.antlr.v4.runtime.tree.TerminalNode
 
 import org.apache.spark.sql.SaveMode
 import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier}
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.catalog._
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.parser._
 import org.apache.spark.sql.catalyst.parser.SqlBaseParser._
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, OneRowRelation, ScriptInputOutputSchema}
 import org.apache.spark.sql.execution.command._
 import org.apache.spark.sql.execution.datasources.{CreateTempViewUsing, _}
 import org.apache.spark.sql.internal.{HiveSerDe, SQLConf, VariableSubstitution}
-import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.types.{DataType, StructField}
 
 /**
  * Concrete parser for Spark SQL statements.
@@ -589,6 +591,53 @@ class SparkSqlAstBuilder(conf: SQLConf) extends AstBuilder {
       ctx.TEMPORARY != null)
   }
 
+  /**
+   * Create a [[CreateMacroCommand]] command.
+   *
+   * For example:
+   * {{{
+   *   CREATE TEMPORARY MACRO macro_name([col_name col_type, ...]) expression;
+   * }}}
+   */
+  override def visitCreateMacro(ctx: CreateMacroContext): LogicalPlan = withOrigin(ctx) {
+    val arguments = Option(ctx.colTypeList).map(visitColTypeList(_))
+      .getOrElse(Seq.empty[StructField]).map { col =>
+      AttributeReference(col.name, col.dataType, col.nullable, col.metadata)() }
+    val colToIndex: Map[String, Int] = arguments.map(_.name).zipWithIndex.toMap
+    if (colToIndex.size != arguments.size) {
+      throw operationNotAllowed(
+        s"Cannot support duplicate colNames for CREATE TEMPORARY MACRO ", ctx)
+    }
+    val macroFunction = expression(ctx.expression).transformUp {
+      case u: UnresolvedAttribute =>
+        val index = colToIndex.get(u.name).getOrElse(
+          throw new ParseException(
+            s"Cannot find colName: [${u}] for CREATE TEMPORARY MACRO", ctx))
+        BoundReference(index, arguments(index).dataType, arguments(index).nullable)
+      case _: SubqueryExpression =>
+        throw operationNotAllowed(s"Cannot support Subquery for CREATE TEMPORARY MACRO", ctx)
+    }
+
+    CreateMacroCommand(
+      ctx.macroName.getText,
+      arguments,
+      macroFunction)
+  }
+
+  /**
+   * Create a [[DropMacroCommand]] command.
+   *
+   * For example:
+   * {{{
+   *   DROP TEMPORARY MACRO [IF EXISTS] macro_name;
+   * }}}
+   */
+  override def visitDropMacro(ctx: DropMacroContext): LogicalPlan = withOrigin(ctx) {
+    DropMacroCommand(
+      ctx.macroName.getText,
+      ctx.EXISTS != null)
+  }
+
   /**
    * Create a [[DropTableCommand]] command.
    */

sql/core/src/main/scala/org/apache/spark/sql/execution/command/macros.scala

@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.command
+
+import org.apache.spark.sql.{AnalysisException, Row, SparkSession}
+import org.apache.spark.sql.catalyst.expressions._
+
+/**
+ * The DDL command that creates a macro.
+ * To create a temporary macro, the syntax of using this command in SQL is:
+ * {{{
+ *    CREATE TEMPORARY MACRO macro_name([col_name col_type, ...]) expression;
+ * }}}
+ */
+case class CreateMacroCommand(
+    macroName: String,
+    columns: Seq[AttributeReference],
+    macroFunction: Expression)
+  extends RunnableCommand {
+
+  override def run(sparkSession: SparkSession): Seq[Row] = {
+    val catalog = sparkSession.sessionState.catalog
+    val macroInfo = columns.mkString(",") + " -> " + macroFunction.toString
+    val info = new ExpressionInfo(macroInfo, macroName)
+    val builder = (children: Seq[Expression]) => {
+      if (children.size != columns.size) {
+        throw new AnalysisException(s"Actual number of columns: ${children.size} != " +
+          s"expected number of columns: ${columns.size} for Macro $macroName")
+      }
+      macroFunction.transformUp {
+        case b: BoundReference => children(b.ordinal)
+      }
+    }
+    catalog.createTempMacro(macroName, info, builder)
+    Seq.empty[Row]
+  }
+}
+
+/**
+ * The DDL command that drops a macro.
+ * ifExists: returns an error if the macro doesn't exist, unless this is true.
+ * {{{
+ *    DROP TEMPORARY MACRO [IF EXISTS] macro_name;
+ * }}}
+ */
+case class DropMacroCommand(macroName: String, ifExists: Boolean)
+  extends RunnableCommand {
+
+  override def run(sparkSession: SparkSession): Seq[Row] = {
+    val catalog = sparkSession.sessionState.catalog
+    catalog.dropTempMacro(macroName, ifExists)
+    Seq.empty[Row]
+  }
+}

sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala

@@ -1309,4 +1309,32 @@ class DDLSuite extends QueryTest with SharedSQLContext with BeforeAndAfterEach {
     assertUnsupported("TRUNCATE TABLE my_tab PARTITION (age=10)")
   }
 
+  test("create/drop temporary macro") {
+    intercept[AnalysisException] {
+      sql(s"CREATE TEMPORARY MACRO simple_add_error(x int) x + y")
+    }
+    intercept[AnalysisException] {
+      sql(s"CREATE TEMPORARY MACRO simple_add_error(x int, x int) x + y")
+    }
+    intercept[AnalysisException] {
+      sql(s"CREATE TEMPORARY MACRO simple_add_error(x int) x NOT IN (select c2 from t2) ")
+    }
+    sql("CREATE TEMPORARY MACRO fixed_number() 42")
+    checkAnswer(sql("SELECT fixed_number()"), Row(42))
+    sql("CREATE TEMPORARY MACRO string_len_plus_two(x string) length(x) + 2")
+    checkAnswer(sql("SELECT string_len_plus_two('abc')"), Row(5))
+    sql("CREATE TEMPORARY MACRO simple_add(x int, y int) x + y")
+    checkAnswer(sql("SELECT simple_add(1, 2)"), Row(3))
+    intercept[AnalysisException] {
+      sql(s"SELECT simple_add(1)")
+    }
+    sql("DROP TEMPORARY MACRO fixed_number")
+    intercept[AnalysisException] {
+      sql(s"DROP TEMPORARY MACRO abs")
+    }
+    intercept[AnalysisException] {
+      sql("DROP TEMPORARY MACRO SOME_MACRO")
+    }
+    sql("DROP TEMPORARY MACRO IF EXISTS SOME_MACRO")
+  }
 }

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala

@@ -1205,13 +1205,6 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
     assertUnsupportedFeature {
       sql("ALTER INDEX my_index ON my_table set IDXPROPERTIES (\"prop1\"=\"val1_new\")")}
   }
-
-  test("create/drop macro commands are not supported") {
-    assertUnsupportedFeature {
-      sql("CREATE TEMPORARY MACRO SIGMOID (x DOUBLE) 1.0 / (1.0 + EXP(-x))")
-    }
-    assertUnsupportedFeature { sql("DROP TEMPORARY MACRO SIGMOID") }
-  }
 }
 
 // for SPARK-2180 test