[SPARK-48008][1/2] Support UDAFs in Spark Connect #46245
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
hvanhovell
reviewed
May 8, 2024
sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
Show resolved
Hide resolved
hvanhovell
reviewed
May 8, 2024
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
Show resolved
Hide resolved
hvanhovell
reviewed
May 8, 2024
...connect/client/jvm/src/main/scala/org/apache/spark/sql/expressions/UserDefinedFunction.scala
Outdated
Show resolved
Hide resolved
hvanhovell
reviewed
May 8, 2024
connector/connect/common/src/main/protobuf/spark/connect/expressions.proto
Outdated
Show resolved
Hide resolved
hvanhovell
reviewed
May 8, 2024
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
Show resolved
Hide resolved
hvanhovell
reviewed
May 21, 2024
...connect/client/jvm/src/main/scala/org/apache/spark/sql/expressions/UserDefinedFunction.scala
Show resolved
Hide resolved
hvanhovell
reviewed
May 21, 2024
| @@ -346,4 +347,42 @@ class UserDefinedFunctionE2ETestSuite extends QueryTest { | |||
| val result = df.select(f($"id")).as[Long].head() | |||
| assert(result == 1L) | |||
| } | |||
|
|
|||
| test("UDAF custom Aggregator - primitive types") { | |||
There was a problem hiding this comment.
Can you add a test for a UDAF with a custom toColumn implementation?
There was a problem hiding this comment.
Let's do this in the next PR where we add support of toColumn.
# Conflicts: # connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/UserDefinedFunctionE2ETestSuite.scala
xupefei
changed the title
# Conflicts: # python/pyspark/sql/connect/proto/commands_pb2.py
|
@HyukjinKwon I think the current protobuf style checks are a bit too strict. The changes made by @xupefei are wire compatible. Can we make this a bit more lenient? |
|
@xupefei can you fix style, and perhaps revert the renaming of the messages? |
xupefei
changed the title
Done! I've reverted the Protobuf change but kept the naming changes in the Scala code. |
Just saw this. seems tests passing fine (?). |
This is the Proto test that is failing: https://github.com/xupefei/spark/actions/runs/9255404729/job/25459140837, for commit 82b802d. I reverted the Proto naming change. |
|
cc @grundprinzip ^^^ |
|
Merging! |
riyaverm-db
pushed a commit
to riyaverm-db/spark
that referenced
this pull request
Jun 7, 2024
### What changes were proposed in this pull request?
This PR changes Spark Connect to support defining and registering `Aggregator[IN, BUF, OUT]` UDAFs.
The mechanism is similar to supporting Scaler UDFs. On the client side, we serialize and send the `Aggregator` instance to the server, where the data is deserialized into an `Aggregator` instance recognized by Spark Core.
With this PR we now have two `Aggregator` interfaces defined, one in Connect API and one in Core. They define exactly the same abstract methods and share the same `SerialVersionUID`, so the Java serialization engine could map one to another. It is very important to keep these two definitions always in sync.
Second part of this effort will be adding `Aggregator.toColumn` API (now NotImplemented due to deps to Spark Core).
### Why are the changes needed?
Spark Connect does not have UDAF support. We need to fix that.
### Does this PR introduce _any_ user-facing change?
Yes, Connect users could now define an Aggregator and register it:
```scala
val agg = new Aggregator[INT, INT, INT] { ... }
spark.udf.register("agg", udaf(agg))
val ds: Dataset[Data] = ...
val aggregated = ds.selectExpr("agg(i)")
```
### How was this patch tested?
Added new tests.
### Was this patch authored or co-authored using generative AI tooling?
Nope.
Closes apache#46245 from xupefei/connect-udaf.
Authored-by: Paddy Xu <xupaddy@gmail.com>
Signed-off-by: Herman van Hovell <herman@databricks.com>
HyukjinKwon
pushed a commit
that referenced
this pull request
Jul 12, 2024
### What changes were proposed in this pull request? This PR follows #46245 to add support `udaf.toColumn` API in Spark Connect. Here we introduce a new Protobuf message, `proto.TypedAggregateExpression`, that includes a serialized UDF packet. On the server, we unpack it into an `Aggregator` object and generate a real `TypedAggregateExpression` instance with the encoder information passed along with the UDF. ### Why are the changes needed? Because the `toColumn` API is not supported in the previous PR. ### Does this PR introduce _any_ user-facing change? Yes, from now on users could create typed UDAF using `udaf.toColumn` API/. ### How was this patch tested? New tests. ### Was this patch authored or co-authored using generative AI tooling? Nope. Closes #46849 from xupefei/connect-udaf-tocolumn. Authored-by: Paddy Xu <xupaddy@gmail.com> Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
jingz-db
pushed a commit
to jingz-db/spark
that referenced
this pull request
Jul 22, 2024
### What changes were proposed in this pull request? This PR follows apache#46245 to add support `udaf.toColumn` API in Spark Connect. Here we introduce a new Protobuf message, `proto.TypedAggregateExpression`, that includes a serialized UDF packet. On the server, we unpack it into an `Aggregator` object and generate a real `TypedAggregateExpression` instance with the encoder information passed along with the UDF. ### Why are the changes needed? Because the `toColumn` API is not supported in the previous PR. ### Does this PR introduce _any_ user-facing change? Yes, from now on users could create typed UDAF using `udaf.toColumn` API/. ### How was this patch tested? New tests. ### Was this patch authored or co-authored using generative AI tooling? Nope. Closes apache#46849 from xupefei/connect-udaf-tocolumn. Authored-by: Paddy Xu <xupaddy@gmail.com> Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
attilapiros
pushed a commit
to attilapiros/spark
that referenced
this pull request
Oct 4, 2024
### What changes were proposed in this pull request? This PR follows apache#46245 to add support `udaf.toColumn` API in Spark Connect. Here we introduce a new Protobuf message, `proto.TypedAggregateExpression`, that includes a serialized UDF packet. On the server, we unpack it into an `Aggregator` object and generate a real `TypedAggregateExpression` instance with the encoder information passed along with the UDF. ### Why are the changes needed? Because the `toColumn` API is not supported in the previous PR. ### Does this PR introduce _any_ user-facing change? Yes, from now on users could create typed UDAF using `udaf.toColumn` API/. ### How was this patch tested? New tests. ### Was this patch authored or co-authored using generative AI tooling? Nope. Closes apache#46849 from xupefei/connect-udaf-tocolumn. Authored-by: Paddy Xu <xupaddy@gmail.com> Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
himadripal
pushed a commit
to himadripal/spark
that referenced
this pull request
Oct 19, 2024
### What changes were proposed in this pull request? This PR follows apache#46245 to add support `udaf.toColumn` API in Spark Connect. Here we introduce a new Protobuf message, `proto.TypedAggregateExpression`, that includes a serialized UDF packet. On the server, we unpack it into an `Aggregator` object and generate a real `TypedAggregateExpression` instance with the encoder information passed along with the UDF. ### Why are the changes needed? Because the `toColumn` API is not supported in the previous PR. ### Does this PR introduce _any_ user-facing change? Yes, from now on users could create typed UDAF using `udaf.toColumn` API/. ### How was this patch tested? New tests. ### Was this patch authored or co-authored using generative AI tooling? Nope. Closes apache#46849 from xupefei/connect-udaf-tocolumn. Authored-by: Paddy Xu <xupaddy@gmail.com> Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR changes Spark Connect to support defining and registering
Aggregator[IN, BUF, OUT]UDAFs.The mechanism is similar to supporting Scaler UDFs. On the client side, we serialize and send the
Aggregatorinstance to the server, where the data is deserialized into anAggregatorinstance recognized by Spark Core.With this PR we now have two
Aggregatorinterfaces defined, one in Connect API and one in Core. They define exactly the same abstract methods and share the sameSerialVersionUID, so the Java serialization engine could map one to another. It is very important to keep these two definitions always in sync.Second part of this effort will be adding
Aggregator.toColumnAPI (now NotImplemented due to deps to Spark Core).Why are the changes needed?
Spark Connect does not have UDAF support. We need to fix that.
Does this PR introduce any user-facing change?
Yes, Connect users could now define an Aggregator and register it:
How was this patch tested?
Added new tests.
Was this patch authored or co-authored using generative AI tooling?
Nope.