(spark) branch master updated: [SPARK-50265][CONNECT] Support spark.udf.registerJavaUdf in Connect

[hv...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hvanhovell pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new bf442c5f1914 [SPARK-50265][CONNECT] Support spark.udf.registerJavaUdf in Connect
bf442c5f1914 is described below

commit bf442c5f1914f7c90971406bced87f13d8dee622
Author: Herman van Hovell <he...@databricks.com>
AuthorDate: Sat Feb 1 08:53:24 2025 -0400

    [SPARK-50265][CONNECT] Support spark.udf.registerJavaUdf in Connect
    
    ### What changes were proposed in this pull request?
    This PR adds support for `UDFRegistration.registerJavaUdf(..)` to the Connect Scala Client.
    
    ### Why are the changes needed?
    This increase compatibility between the Classic and Connect implementations.
    
    ### Does this PR introduce _any_ user-facing change?
    Yes. It enables Scala client users to use `UDFRegistration.registerJavaUdf(..)`.
    
    ### How was this patch tested?
    I have added a Unit tests to the `UserDefinedFunctionE2ETestSuite`.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Closes #49759 from hvanhovell/SPARK-50265.
    
    Authored-by: Herman van Hovell <he...@databricks.com>
    Signed-off-by: Herman van Hovell <he...@databricks.com>
---
 .../sql/connect/UserDefinedFunctionE2ETestSuite.scala | 19 ++++++++++++++++---
 .../apache/spark/sql/connect/UDFRegistration.scala    | 10 ++++++++--
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/UserDefinedFunctionE2ETestSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/UserDefinedFunctionE2ETestSuite.scala
index a4ef11554d30..71734b907b9f 100644
--- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/UserDefinedFunctionE2ETestSuite.scala
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/UserDefinedFunctionE2ETestSuite.scala
@@ -25,11 +25,11 @@ import scala.jdk.CollectionConverters._
 import org.apache.spark.api.java.function._
 import org.apache.spark.sql.{AnalysisException, Encoder, Encoders, Row}
 import org.apache.spark.sql.api.java.UDF2
-import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{PrimitiveIntEncoder, PrimitiveLongEncoder}
+import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{PrimitiveIntEncoder, PrimitiveLongEncoder, StringEncoder}
 import org.apache.spark.sql.connect.test.{QueryTest, RemoteSparkSession}
 import org.apache.spark.sql.expressions.Aggregator
-import org.apache.spark.sql.functions.{col, struct, udaf, udf}
-import org.apache.spark.sql.types.IntegerType
+import org.apache.spark.sql.functions.{call_function, col, struct, udaf, udf}
+import org.apache.spark.sql.types.{IntegerType, StringType}
 
 /**
  * All tests in this class requires client UDF defined in this test class synced with the server.
@@ -433,6 +433,15 @@ class UserDefinedFunctionE2ETestSuite extends QueryTest with RemoteSparkSession
     assert(ds.select(RowAggregator.toColumn).head() == 405)
     assert(ds.agg(RowAggregator.toColumn).head().getLong(0) == 405)
   }
+
+  test("registerJavaUdf") {
+    spark.udf.registerJava("sconcat", classOf[StringConcat].getName, StringType)
+    val ds = spark
+      .range(2)
+      .select(call_function("sconcat", col("id").cast("string"), (col("id") + 1).cast("string")))
+      .as(StringEncoder)
+    checkDataset(ds, "01", "12")
+  }
 }
 
 case class UdafTestInput(id: Long, extra: Long)
@@ -490,3 +499,7 @@ object RowAggregator extends Aggregator[Row, (Long, Long), Long] {
     Encoders.tuple(Encoders.scalaLong, Encoders.scalaLong)
   override def outputEncoder: Encoder[Long] = Encoders.scalaLong
 }
+
+class StringConcat extends UDF2[String, String, String] {
+  override def call(t1: String, t2: String): String = t1 + t2
+}
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/UDFRegistration.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/UDFRegistration.scala
index 67471f9cf523..68a216a70802 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/UDFRegistration.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/UDFRegistration.scala
@@ -17,7 +17,9 @@
 
 package org.apache.spark.sql.connect
 
+import org.apache.spark.connect.proto
 import org.apache.spark.sql
+import org.apache.spark.sql.connect.common.DataTypeProtoConverter
 import org.apache.spark.sql.expressions.{UserDefinedAggregateFunction, UserDefinedFunction}
 import org.apache.spark.sql.types.DataType
 
@@ -42,8 +44,12 @@ class UDFRegistration(session: SparkSession) extends sql.UDFRegistration {
   }
 
   override def registerJava(name: String, className: String, returnDataType: DataType): Unit = {
-    throw new UnsupportedOperationException(
-      "registerJava is currently not supported in Spark Connect.")
+    val builder = proto.CommonInlineUserDefinedFunction.newBuilder().setFunctionName(name)
+    builder.getJavaUdfBuilder
+      .setClassName(className)
+      .setOutputType(DataTypeProtoConverter.toConnectProtoType(returnDataType))
+      .setAggregate(false)
+    session.registerUdf(builder.build())
   }
 
   /** @inheritdoc */


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org