Add tests

imback82 · imback82 · commit a68b9774a277 · 2021-02-06T14:02:40.000-08:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala b/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
@@ -114,10 +114,10 @@ class UDFRegistration private[sql] (functionRegistry: FunctionRegistry) extends
         def builder(children: Seq[Expression]) = udaf.scalaAggregator(children)
         functionRegistry.createOrReplaceTempFunction(name, builder)
         udf
-      case _ =>
-        def builder(children: Seq[Expression]) = udf.apply(children.map(Column.apply) : _*).expr
+      case other =>
+        def builder(children: Seq[Expression]) = other.apply(children.map(Column.apply) : _*).expr
         functionRegistry.createOrReplaceTempFunction(name, builder)
-        udf
+        other
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
@@ -26,15 +26,18 @@ import scala.collection.mutable.{ArrayBuffer, WrappedArray}
 
 import org.apache.spark.SparkException
 import org.apache.spark.sql.api.java._
-import org.apache.spark.sql.catalyst.encoders.OuterScopes
+import org.apache.spark.sql.catalyst.FunctionIdentifier
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, OuterScopes}
+import org.apache.spark.sql.catalyst.expressions.{Literal, ScalaUDF}
 import org.apache.spark.sql.catalyst.plans.logical.Project
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.execution.{QueryExecution, SimpleMode}
+import org.apache.spark.sql.execution.aggregate.{ScalaAggregator, ScalaUDAF}
 import org.apache.spark.sql.execution.columnar.InMemoryRelation
 import org.apache.spark.sql.execution.command.{CreateDataSourceTableAsSelectCommand, ExplainCommand}
 import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand
-import org.apache.spark.sql.expressions.SparkUserDefinedFunction
-import org.apache.spark.sql.functions.{lit, struct, udf}
+import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, SparkUserDefinedFunction, UserDefinedAggregateFunction}
+import org.apache.spark.sql.functions.{lit, struct, udaf, udf}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.test.SQLTestData._
@@ -798,4 +801,47 @@ class UDFSuite extends QueryTest with SharedSparkSession {
       .select(myUdf(Column("col"))),
       Row(ArrayBuffer(100)))
   }
+
+  test("SPARK-34388: UDF name is propagated with registration for ScalaUDF") {
+    spark.udf.register("udf34388", udf((value: Int) => value > 2))
+    spark.sessionState.catalog.lookupFunction(
+      FunctionIdentifier("udf34388"), Seq(Literal(1))) match {
+      case udf: ScalaUDF => assert(udf.name === "udf34388")
+    }
+  }
+
+  test("SPARK-34388: UDF name is propagated with registration for ScalaAggregator") {
+    val agg = new Aggregator[Long, Long, Long] {
+      override def zero: Long = 0L
+      override def reduce(b: Long, a: Long): Long = a + b
+      override def merge(b1: Long, b2: Long): Long = b1 + b2
+      override def finish(reduction: Long): Long = reduction
+      override def bufferEncoder: Encoder[Long] = ExpressionEncoder[Long]()
+      override def outputEncoder: Encoder[Long] = ExpressionEncoder[Long]()
+    }
+
+    spark.udf.register("agg34388", udaf(agg))
+    spark.sessionState.catalog.lookupFunction(
+      FunctionIdentifier("agg34388"), Seq(Literal(1))) match {
+      case agg: ScalaAggregator[_, _, _] => assert(agg.name === "agg34388")
+    }
+  }
+
+  test("SPARK-34388: UDF name is propagated with registration for ScalaUDAF") {
+    val udaf = new UserDefinedAggregateFunction {
+      def inputSchema: StructType = new StructType().add("a", LongType)
+      def bufferSchema: StructType = new StructType().add("product", LongType)
+      def dataType: DataType = LongType
+      def deterministic: Boolean = true
+      def initialize(buffer: MutableAggregationBuffer): Unit = {}
+      def update(buffer: MutableAggregationBuffer, input: Row): Unit = {}
+      def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {}
+      def evaluate(buffer: Row): Any = buffer.getLong(0)
+    }
+    spark.udf.register("udaf34388", udaf)
+    spark.sessionState.catalog.lookupFunction(
+      FunctionIdentifier("udaf34388"), Seq(Literal(1))) match {
+      case udaf: ScalaUDAF => assert(udaf.name === "udaf34388")
+    }
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -114,10 +114,10 @@ class UDFRegistration private[sql] (functionRegistry: FunctionRegistry) extends`
`114`	`114`	`def builder(children: Seq[Expression]) = udaf.scalaAggregator(children)`
`115`	`115`	`functionRegistry.createOrReplaceTempFunction(name, builder)`
`116`	`116`	`udf`
`117`		`- case _ =>`
`118`		`- def builder(children: Seq[Expression]) = udf.apply(children.map(Column.apply) : _*).expr`
	`117`	`+ case other =>`
	`118`	`+ def builder(children: Seq[Expression]) = other.apply(children.map(Column.apply) : _*).expr`
`119`	`119`	`functionRegistry.createOrReplaceTempFunction(name, builder)`
`120`		`- udf`
	`120`	`+ other`
`121`	`121`	`}`
`122`	`122`	`}`
`123`	`123`