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

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

@@ -129,6 +129,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 | COST)?
         statement                                                      #explain
     | SHOW TABLES ((FROM | IN) db=identifier)?

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionInfo.java

@@ -26,6 +26,7 @@ public class ExpressionInfo {
     private String name;
     private String extended;
     private String db;
+    private boolean macro;
 
     public String getClassName() {
         return className;
@@ -47,19 +48,32 @@ public String getDb() {
         return db;
     }
 
-    public ExpressionInfo(String className, String db, String name, String usage, String extended) {
+    public boolean isMacro() {
+        return macro;
+    }
+
+    public ExpressionInfo(String className, String db, String name, String usage, String extended, boolean macro) {
         this.className = className;
         this.db = db;
         this.name = name;
         this.usage = usage;
         this.extended = extended;
+        this.macro = macro;
+    }
+
+    public ExpressionInfo(String className, String db, String name, String usage, String extended) {
+        this(className, db, name, usage, extended, false);
     }
 
     public ExpressionInfo(String className, String name) {
-        this(className, null, name, null, null);
+        this(className, null, name, null, null, false);
+    }
+
+    public ExpressionInfo(String className, String name, boolean macro) {
+        this(className, null, name, null, null, macro);
     }
 
     public ExpressionInfo(String className, String db, String name) {
-        this(className, db, name, null, null);
+        this(className, db, name, null, null, false);
     }
 }

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

@@ -60,6 +60,9 @@ trait FunctionRegistry {
   /** Drop a function and return whether the function existed. */
   def dropFunction(name: String): Boolean
 
+  /** Drop a macro and return whether the macro existed. */
+  def dropMacro(name: String): Boolean
+
   /** Checks if a function with a given name exists. */
   def functionExists(name: String): Boolean = lookupFunction(name).isDefined
 
@@ -107,6 +110,14 @@ class SimpleFunctionRegistry extends FunctionRegistry {
     functionBuilders.remove(name).isDefined
   }
 
+  override def dropMacro(name: String): Boolean = synchronized {
+    if (functionBuilders.get(name).map(_._1).filter(_.isMacro).isDefined) {
+      functionBuilders.remove(name).isDefined
+    } else {
+      false
+    }
+  }
+
   override def clear(): Unit = synchronized {
     functionBuilders.clear()
   }
@@ -146,6 +157,10 @@ object EmptyFunctionRegistry extends FunctionRegistry {
     throw new UnsupportedOperationException
   }
 
+  override def dropMacro(name: String): Boolean = {
+    throw new UnsupportedOperationException
+  }
+
   override def dropFunction(name: String): Boolean = {
     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

@@ -1090,6 +1090,24 @@ class SessionCatalog(
     }
   }
 
+  /** Create a temporary macro. */
+  def createTempMacro(
+      name: String,
+      info: ExpressionInfo,
+      functionBuilder: FunctionBuilder): Unit = {
+    if (functionRegistry.functionExists(name)) {
+      throw new AnalysisException(s"Function $name already exists")
+    }
+    functionRegistry.registerFunction(name, info, functionBuilder)
+  }
+
+  /** Drop a temporary macro. */
+  def dropTempMacro(name: String, ignoreIfNotExists: Boolean): Unit = {
+    if (!functionRegistry.dropMacro(name) && !ignoreIfNotExists) {
+      throw new NoSuchTempMacroException(name)
+    }
+  }
+
   /**
    * Returns whether it is a temporary function. If not existed, returns false.
    */

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

@@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.command._
 import org.apache.spark.sql.execution.datasources.{CreateTable, _}
 import org.apache.spark.sql.internal.{HiveSerDe, SQLConf, VariableSubstitution}
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{StructField, StructType}
 
 /**
  * Concrete parser for Spark SQL statements.
@@ -715,6 +715,37 @@ class SparkSqlAstBuilder(conf: SQLConf) extends AstBuilder(conf) {
       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])
+    val e = expression(ctx.expression)
+    CreateMacroCommand(
+      ctx.macroName.getText,
+      MacroFunctionWrapper(arguments, e))
+  }
+
+  /**
+   * 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,99 @@
+/*
+ * 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.analysis._
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types.StructField
+
+/**
+ * This class provides arguments and body expression of the macro function.
+ */
+case class MacroFunctionWrapper(columns: Seq[StructField], macroFunction: Expression)
+
+/**
+ * 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,
+    funcWrapper: MacroFunctionWrapper)
+  extends RunnableCommand {
+
+  override def run(sparkSession: SparkSession): Seq[Row] = {
+    val catalog = sparkSession.sessionState.catalog
+    val columns = funcWrapper.columns.map { col =>
+      AttributeReference(col.name, col.dataType, col.nullable, col.metadata)() }
+    val colToIndex: Map[String, Int] = columns.map(_.name).zipWithIndex.toMap
+    if (colToIndex.size != columns.size) {
+      throw new AnalysisException(s"Cannot support duplicate colNames " +
+        s"for CREATE TEMPORARY MACRO $macroName, actual columns: ${columns.mkString(",")}")
+    }
+    val macroFunction = funcWrapper.macroFunction.transform {
+      case u: UnresolvedAttribute =>
+        val index = colToIndex.get(u.name).getOrElse(
+          throw new AnalysisException(s"Cannot find colName: ${u} " +
+            s"for CREATE TEMPORARY MACRO $macroName, actual columns: ${columns.mkString(",")}"))
+        BoundReference(index, columns(index).dataType, columns(index).nullable)
+      case u: UnresolvedFunction =>
+        sparkSession.sessionState.catalog.lookupFunction(u.name, u.children)
+      case s: SubqueryExpression =>
+        throw new AnalysisException(s"Cannot support Subquery: ${s} " +
+          s"for CREATE TEMPORARY MACRO $macroName")
+      case u: UnresolvedGenerator =>
+        throw new AnalysisException(s"Cannot support Generator: ${u} " +
+          s"for CREATE TEMPORARY MACRO $macroName")
+    }
+
+    val macroInfo = columns.mkString(",") + " -> " + funcWrapper.macroFunction.toString
+    val info = new ExpressionInfo(macroInfo, macroName, true)
+    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.transform {
+        // Skip to validate the input type because check it at runtime.
+        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

@@ -1516,6 +1516,35 @@ abstract class DDLSuite extends QueryTest with SQLTestUtils {
     )
   }
 
+  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")
+  }
+
   test("create a data source table without schema") {
     import testImplicits._
     withTempPath { tempDir =>

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

@@ -1153,13 +1153,6 @@ class HiveQuerySuite extends HiveComparisonTest with SQLTestUtils with BeforeAnd
       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") }
-  }
-
   test("dynamic partitioning is allowed when hive.exec.dynamic.partition.mode is nonstrict") {
     val modeConfKey = "hive.exec.dynamic.partition.mode"
     withTable("with_parts") {