fix tests

imback82 · imback82 · commit 56c3001ea558 · 2021-02-06T21:11: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
@@ -113,7 +113,7 @@ class UDFRegistration private[sql] (functionRegistry: FunctionRegistry) extends
       case udaf: UserDefinedAggregator[_, _, _] =>
         def builder(children: Seq[Expression]) = udaf.scalaAggregator(children)
         functionRegistry.createOrReplaceTempFunction(name, builder)
-        udf
+        udaf
       case other =>
         def builder(children: Seq[Expression]) = other.apply(children.map(Column.apply) : _*).expr
         functionRegistry.createOrReplaceTempFunction(name, builder)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IntegratedUDFTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/IntegratedUDFTestUtils.scala
@@ -321,26 +321,34 @@ object IntegratedUDFTestUtils extends SQLHelper {
    *   casted_col.cast(df.schema("col").dataType)
    * }}}
    */
-  case class TestScalaUDF(name: String) extends TestUDF {
-    private[IntegratedUDFTestUtils] lazy val udf = new SparkUserDefinedFunction(
-      (input: Any) => if (input == null) {
-        null
-      } else {
-        input.toString
-      },
-      StringType,
-      inputEncoders = Seq.fill(1)(None),
-      name = Some(name)) {
-
-      override def apply(exprs: Column*): Column = {
-        assert(exprs.length == 1, "Defined UDF only has one column")
-        val expr = exprs.head.expr
-        assert(expr.resolved, "column should be resolved to use the same type " +
-          "as input. Try df(name) or df.col(name)")
-        Column(Cast(createScalaUDF(Cast(expr, StringType) :: Nil), expr.dataType))
-      }
+  class TestInternalScalaUDF(name: String) extends SparkUserDefinedFunction(
+    (input: Any) => if (input == null) {
+      null
+    } else {
+      input.toString
+    },
+    StringType,
+    inputEncoders = Seq.fill(1)(None),
+    name = Some(name)) {
+
+    override def apply(exprs: Column*): Column = {
+      assert(exprs.length == 1, "Defined UDF only has one column")
+      val expr = exprs.head.expr
+      assert(expr.resolved, "column should be resolved to use the same type " +
+        "as input. Try df(name) or df.col(name)")
+      Column(Cast(createScalaUDF(Cast(expr, StringType) :: Nil), expr.dataType))
     }
 
+    override def withName(name: String): TestInternalScalaUDF = {
+      // "withName" should overridden to return TestInternalScalaUDF. Otherwise, the current object
+      // is sliced and the overridden "apply" is not invoked.
+      new TestInternalScalaUDF(name)
+    }
+  }
+
+  case class TestScalaUDF(name: String) extends TestUDF {
+    private[IntegratedUDFTestUtils] lazy val udf = new TestInternalScalaUDF(name)
+
     def apply(exprs: Column*): Column = udf(exprs: _*)
 
     val prettyName: String = "Scala UDF"
