Skip to content

Commit 61ae5b7

Browse files
committed
add Java UDF test
1 parent 1111c6f commit 61ae5b7

1 file changed

Lines changed: 20 additions & 0 deletions

File tree

spark/src/test/scala/org/apache/comet/CometCodegenDispatchSmokeSuite.scala

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ package org.apache.comet
2222
import org.apache.arrow.vector._
2323
import org.apache.spark.{SparkConf, TaskContext}
2424
import org.apache.spark.sql.CometTestBase
25+
import org.apache.spark.sql.api.java.UDF1
2526
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
2627
import org.apache.spark.sql.types._
2728

@@ -269,6 +270,25 @@ class CometCodegenDispatchSmokeSuite extends CometTestBase with AdaptiveSparkPla
269270
}
270271
}
271272

273+
test("registered Java UDF1 routes through dispatcher") {
274+
// Java API path: `spark.udf.register(name, UDF1<...>, returnType)`. Spark wraps the Java
275+
// functional interface in a Scala function and produces a `ScalaUDF` expression at plan
276+
// time, so the dispatcher handles it the same as a Scala-registered UDF. Sanity check that
277+
// both registration paths land on the same routing code.
278+
spark.udf.register(
279+
"javaLen",
280+
new UDF1[String, Integer] {
281+
override def call(s: String): Integer = if (s == null) -1 else s.length
282+
},
283+
IntegerType)
284+
withSubjects("abc", "hello", null, "x") {
285+
assertCodegenDidWork {
286+
checkSparkAnswerAndOperator(sql("SELECT javaLen(s) FROM t"))
287+
}
288+
assertKernelSignaturePresent(Seq(classOf[VarCharVector]), IntegerType)
289+
}
290+
}
291+
272292
test("multi-arg ScalaUDF over string + literal routes through dispatcher") {
273293
spark.udf.register(
274294
"prepend",

0 commit comments

Comments
 (0)