[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

@@ -21,11 +21,17 @@
  * Expression information, will be used to describe a expression.
  */
 public class ExpressionInfo {
+
+    public enum FunctionType {
+        BUILTIN, PERSISTENT, TEMPORARY;
+    }
+
     private String className;
     private String usage;
     private String name;
     private String extended;
     private String db;
+    private FunctionType functionType;
 
     public String getClassName() {
         return className;
@@ -47,19 +53,32 @@ public String getDb() {
         return db;
     }
 
-    public ExpressionInfo(String className, String db, String name, String usage, String extended) {
+    public FunctionType getFunctionType() {
+        return functionType;
+    }
+
+    public ExpressionInfo(String className, String db, String name, String usage, String extended, FunctionType functionType) {
         this.className = className;
         this.db = db;
         this.name = name;
         this.usage = usage;
         this.extended = extended;
+        this.functionType = functionType;
     }
 
     public ExpressionInfo(String className, String name) {
-        this(className, null, name, null, null);
+        this(className, null, name, null, null, FunctionType.TEMPORARY);
+    }
+
+    public ExpressionInfo(String className, String name, FunctionType functionType) {
+        this(className, null, name, null, null, functionType);
     }
 
     public ExpressionInfo(String className, String db, String name) {
-        this(className, db, name, null, null);
+        this(className, db, name, null, null, FunctionType.TEMPORARY);
+    }
+
+    public ExpressionInfo(String className, String db, String name, FunctionType functionType) {
+        this(className, db, name, null, null, functionType);
     }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/AnalysisException.scala

@@ -55,3 +55,12 @@ class AnalysisException protected[sql] (
     s"$message;$lineAnnotation$positionAnnotation"
   }
 }
+
+object AnalysisException {
+  /**
+   * Create a no such temporary macro exception.
+   */
+  def noSuchTempMacroException(func: String): AnalysisException = {
+    new AnalysisException(s"Temporary macro '$func' not found")
+  }
+}

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

@@ -26,6 +26,7 @@ import scala.util.{Failure, Success, Try}
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.ExpressionInfo.FunctionType
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.expressions.xml._
 import org.apache.spark.sql.catalyst.util.StringKeyHashMap
@@ -120,6 +121,58 @@ class SimpleFunctionRegistry extends FunctionRegistry {
   }
 }
 
+class SystemFunctionRegistry(builtin: SimpleFunctionRegistry) extends SimpleFunctionRegistry {
+
+  override def registerFunction(
+    name: String,
+    info: ExpressionInfo,
+    builder: FunctionBuilder): Unit = synchronized {
+    if (info.getFunctionType.equals(FunctionType.BUILTIN)) {
+      builtin.registerFunction(name, info, builder)
+    } else {
+      functionBuilders.put(name, (info, builder))
+    }
+  }
+
+  override def lookupFunction(name: String, children: Seq[Expression]): Expression = {
+    val func = synchronized {
+      functionBuilders.get(name).map(_._2).orElse(builtin.lookupFunctionBuilder(name)).getOrElse {
+        throw new AnalysisException(s"undefined function $name")
+      }
+    }
+    func(children)
+  }
+
+  override def listFunction(): Seq[String] = synchronized {
+    (functionBuilders.iterator.map(_._1).toList ++ builtin.listFunction()).distinct.sorted
+  }
+
+  override def lookupFunction(name: String): Option[ExpressionInfo] = synchronized {
+    functionBuilders.get(name).map(_._1).orElse(builtin.lookupFunction(name))
+  }
+
+  override def lookupFunctionBuilder(name: String): Option[FunctionBuilder] = synchronized {
+    functionBuilders.get(name).map(_._2).orElse(builtin.lookupFunctionBuilder(name))
+  }
+
+  override def dropFunction(name: String): Boolean = synchronized {
+    functionBuilders.remove(name).isDefined
+  }
+
+  override def clear(): Unit = synchronized {
+    builtin.clear()
+    functionBuilders.clear()
+  }
+
+  override def clone(): SimpleFunctionRegistry = synchronized {
+    val registry = new SystemFunctionRegistry(builtin.clone())
+    functionBuilders.iterator.foreach { case (name, (info, builder)) =>
+      registry.registerFunction(name, info, builder)
+    }
+    registry
+  }
+}
+
 /**
  * A trivial catalog that returns an error when a function is requested. Used for testing when all
  * functions are already filled in and the analyzer needs only to resolve attribute references.
@@ -456,6 +509,8 @@ object FunctionRegistry {
     fr
   }
 
+  val systemRegistry = new SystemFunctionRegistry(builtin)
+
   val functionSet: Set[String] = builtin.listFunction().toSet
 
   /** See usage above. */
@@ -519,7 +574,8 @@ object FunctionRegistry {
     }
     val clazz = scala.reflect.classTag[Cast].runtimeClass
     val usage = "_FUNC_(expr) - Casts the value `expr` to the target data type `_FUNC_`."
-    (name, (new ExpressionInfo(clazz.getCanonicalName, null, name, usage, null), builder))
+    (name, (new ExpressionInfo(clazz.getCanonicalName, null, name, usage, null,
+      FunctionType.BUILTIN), builder))
   }
 
   /**
@@ -529,9 +585,10 @@ object FunctionRegistry {
     val clazz = scala.reflect.classTag[T].runtimeClass
     val df = clazz.getAnnotation(classOf[ExpressionDescription])
     if (df != null) {
-      new ExpressionInfo(clazz.getCanonicalName, null, name, df.usage(), df.extended())
+      new ExpressionInfo(clazz.getCanonicalName, null, name, df.usage(), df.extended(),
+        ExpressionInfo.FunctionType.BUILTIN)
     } else {
-      new ExpressionInfo(clazz.getCanonicalName, name)
+      new ExpressionInfo(clazz.getCanonicalName, name, ExpressionInfo.FunctionType.BUILTIN)
     }
   }
 

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

@@ -34,6 +34,7 @@ import org.apache.spark.sql.catalyst._
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
 import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo}
+import org.apache.spark.sql.catalyst.expressions.ExpressionInfo.FunctionType
 import org.apache.spark.sql.catalyst.parser.{CatalystSqlParser, ParserInterface}
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, SubqueryAlias, View}
 import org.apache.spark.sql.catalyst.util.StringUtils
@@ -1090,6 +1091,21 @@ class SessionCatalog(
     }
   }
 
+  /** Create a temporary macro. */
+  def createTempMacro(
+      name: String,
+      info: ExpressionInfo,
+      functionBuilder: FunctionBuilder): Unit = {
+    functionRegistry.registerFunction(name, info, functionBuilder)
+  }
+
+  /** Drop a temporary macro. */
+  def dropTempMacro(name: String, ignoreIfNotExists: Boolean): Unit = {
+    if (!functionRegistry.dropFunction(name) && !ignoreIfNotExists) {
+      throw AnalysisException.noSuchTempMacroException(name)
+    }
+  }
+
   /**
    * Returns whether it is a temporary function. If not existed, returns false.
    */
@@ -1126,7 +1142,8 @@ class SessionCatalog(
           new ExpressionInfo(
             metadata.className,
             qualifiedName.database.orNull,
-            qualifiedName.identifier)
+            qualifiedName.identifier,
+            FunctionType.PERSISTENT)
         } else {
           failFunctionLookup(name.funcName)
         }
@@ -1248,7 +1265,11 @@ class SessionCatalog(
       if (func.database.isDefined) {
         dropFunction(func, ignoreIfNotExists = false)
       } else {
-        dropTempFunction(func.funcName, ignoreIfNotExists = false)
+        val functionType = functionRegistry.lookupFunction(func.funcName)
+          .map(_.getFunctionType).getOrElse(FunctionType.TEMPORARY)
+        if (!functionType.equals(FunctionType.BUILTIN)) {
+          dropTempFunction(func.funcName, ignoreIfNotExists = false)
+        }
       }
     }
     clearTempTables()

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala

@@ -260,7 +260,8 @@ object CreateStruct extends FunctionBuilder {
       null,
       "struct",
       "_FUNC_(col1, col2, col3, ...) - Creates a struct with the given field values.",
-      "")
+      "",
+      ExpressionInfo.FunctionType.BUILTIN)
     ("struct", (info, this))
   }
 }

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,36 @@ 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 columns = createSchema(ctx.colTypeList)
+    val e = expression(ctx.expression)
+    CreateMacroCommand(
+      ctx.macroName.getText,
+      MacroFunctionWrapper(columns, 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.
    */