[SPARK-48516][PYTHON][CONNECT] Turn on Arrow optimization for Python UDFs by default #49482
+38
−8
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
| @@ -71,6 +71,15 @@ | |||
| pass | |||
|
|
|||
|
|
|||
| has_arrow: bool = False | |||
There was a problem hiding this comment.
I think we can use from pyspark.testing.utils import have_pyarrow
There was a problem hiding this comment.
Can we address this comment @xinrong-meng
There was a problem hiding this comment.
Good point! Resolved, thank you
There was a problem hiding this comment.
There will be a circular import if we do that. Let me follow up with a separate PR instead.
There was a problem hiding this comment.
It seems that there are many unit test failures unfortunately, @xinrong-meng .
Feel free to have your time. I believe that the community agree that this is worthy for Apache Spark 4.0.0. We can backport this when this PR is ready.
Let's make CIs pass first.
|
Gentle ping, @xinrong-meng . If this is targeting Apache Spark 4.0, we had better have this before February 1st. |
|
Thanks @dongjoon-hyun I just got some free cycles for that and will resolve it ASAP. |
|
A quick update: the PR is blocked by UDT support in Arrow Python UDF, which I’m currently working on |
|
Thank you for the updated context. |
xinrong-meng
changed the title
|
Thanks @dongjoon-hyun for attention! The current proposal is to fall back to the existing (non-Arrow-optimized) Python UDF when UDT is involved. My understanding is that no further testing is needed and the code change is minimal (just an if-else), but I respect the community’s decision. |
|
I marked it as WIP because I wanted to file a separate PR for the fallback mechanism with tests. Once that PR is in, this PR will be unblocked immediately. |
|
Let's make this ready ASAP for 4.0 |
xinrong-meng
added a commit
that referenced
this pull request
Feb 5, 2025
…t and output types ### What changes were proposed in this pull request? Introduce a fallback mechanism for Arrow-optimized Python UDFs when either the input or return types contain User-Defined Types (UDTs). If UDTs are detected, the system logs a warning and switches to currently default, non-Arrow-optimized UDF. ### Why are the changes needed? To unblock enabling Arrow-optimized Python UDFs by default, see [pr](#49482) ### Does this PR introduce _any_ user-facing change? Yes. UDT input and output types will not fail Arrow Python UDF anymore, as shown below: ```py >>> import pyspark.sql.functions as F >>> from pyspark.sql import Row >>> from pyspark.testing.sqlutils import ExamplePoint, ExamplePointUDT # UDT intput >>> from pyspark.sql.types import * >>> row = Row(label=1.0, point=ExamplePoint(1.0, 2.0)) >>> df = spark.createDataFrame([row]) >>> >>> udf1 = F.udf(lambda p: p.y, DoubleType(), useArrow=True) >>> df.select(udf1(df.point)).show() 25/02/03 17:49:57 WARN ExtractPythonUDFs: Arrow optimization disabled due to UDT input or return type. Falling back to non-Arrow-optimized UDF execution. +---------------+ |<lambda>(point)| +---------------+ | 2.0| +---------------+ # UDT output >>> row = Row(value=3.0) >>> df = spark.createDataFrame([row]) >>> udf_with_udt_output = F.udf(lambda v: ExamplePoint(v, v + 1), ExamplePointUDT(), useArrow=True) >>> df.select(udf_with_udt_output(df.value)).show() 25/02/03 17:51:43 WARN ExtractPythonUDFs: Arrow optimization disabled due to UDT input or return type. Falling back to non-Arrow-optimized UDF execution. +---------------+ |<lambda>(value)| +---------------+ | (3.0, 4.0)| +---------------+ ``` ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #49786 from xinrong-meng/udt_arrow_udf. Authored-by: Xinrong Meng <[email protected]> Signed-off-by: Xinrong Meng <[email protected]>
xinrong-meng
added a commit
that referenced
this pull request
Feb 5, 2025
…t and output types ### What changes were proposed in this pull request? Introduce a fallback mechanism for Arrow-optimized Python UDFs when either the input or return types contain User-Defined Types (UDTs). If UDTs are detected, the system logs a warning and switches to currently default, non-Arrow-optimized UDF. ### Why are the changes needed? To unblock enabling Arrow-optimized Python UDFs by default, see [pr](#49482) ### Does this PR introduce _any_ user-facing change? Yes. UDT input and output types will not fail Arrow Python UDF anymore, as shown below: ```py >>> import pyspark.sql.functions as F >>> from pyspark.sql import Row >>> from pyspark.testing.sqlutils import ExamplePoint, ExamplePointUDT # UDT intput >>> from pyspark.sql.types import * >>> row = Row(label=1.0, point=ExamplePoint(1.0, 2.0)) >>> df = spark.createDataFrame([row]) >>> >>> udf1 = F.udf(lambda p: p.y, DoubleType(), useArrow=True) >>> df.select(udf1(df.point)).show() 25/02/03 17:49:57 WARN ExtractPythonUDFs: Arrow optimization disabled due to UDT input or return type. Falling back to non-Arrow-optimized UDF execution. +---------------+ |<lambda>(point)| +---------------+ | 2.0| +---------------+ # UDT output >>> row = Row(value=3.0) >>> df = spark.createDataFrame([row]) >>> udf_with_udt_output = F.udf(lambda v: ExamplePoint(v, v + 1), ExamplePointUDT(), useArrow=True) >>> df.select(udf_with_udt_output(df.value)).show() 25/02/03 17:51:43 WARN ExtractPythonUDFs: Arrow optimization disabled due to UDT input or return type. Falling back to non-Arrow-optimized UDF execution. +---------------+ |<lambda>(value)| +---------------+ | (3.0, 4.0)| +---------------+ ``` ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #49786 from xinrong-meng/udt_arrow_udf. Authored-by: Xinrong Meng <[email protected]> Signed-off-by: Xinrong Meng <[email protected]> (cherry picked from commit ac2f3a4) Signed-off-by: Xinrong Meng <[email protected]>
xinrong-meng
force-pushed
the
arrow_on
branch
from
00d4e03 to
e25b0eb
Compare
xinrong-meng
changed the title
|
The Arrow fallback PR is in so this PR should be unblocked. I’ll keep an eye on testing and make it ready ASAP! Thank you @HyukjinKwon @dongjoon-hyun |
HyukjinKwon
approved these changes
Feb 5, 2025
|
can we address #49482 (comment) @xinrong-meng ? |
xinrong-meng
force-pushed
the
arrow_on
branch
from
ac09220 to
e25b0eb
Compare
|
Can you take a look at the test failure to make sure? I think those failures look related. |
zecookiez
pushed a commit
to zecookiez/spark
that referenced
this pull request
Feb 6, 2025
…t and output types ### What changes were proposed in this pull request? Introduce a fallback mechanism for Arrow-optimized Python UDFs when either the input or return types contain User-Defined Types (UDTs). If UDTs are detected, the system logs a warning and switches to currently default, non-Arrow-optimized UDF. ### Why are the changes needed? To unblock enabling Arrow-optimized Python UDFs by default, see [pr](apache#49482) ### Does this PR introduce _any_ user-facing change? Yes. UDT input and output types will not fail Arrow Python UDF anymore, as shown below: ```py >>> import pyspark.sql.functions as F >>> from pyspark.sql import Row >>> from pyspark.testing.sqlutils import ExamplePoint, ExamplePointUDT # UDT intput >>> from pyspark.sql.types import * >>> row = Row(label=1.0, point=ExamplePoint(1.0, 2.0)) >>> df = spark.createDataFrame([row]) >>> >>> udf1 = F.udf(lambda p: p.y, DoubleType(), useArrow=True) >>> df.select(udf1(df.point)).show() 25/02/03 17:49:57 WARN ExtractPythonUDFs: Arrow optimization disabled due to UDT input or return type. Falling back to non-Arrow-optimized UDF execution. +---------------+ |<lambda>(point)| +---------------+ | 2.0| +---------------+ # UDT output >>> row = Row(value=3.0) >>> df = spark.createDataFrame([row]) >>> udf_with_udt_output = F.udf(lambda v: ExamplePoint(v, v + 1), ExamplePointUDT(), useArrow=True) >>> df.select(udf_with_udt_output(df.value)).show() 25/02/03 17:51:43 WARN ExtractPythonUDFs: Arrow optimization disabled due to UDT input or return type. Falling back to non-Arrow-optimized UDF execution. +---------------+ |<lambda>(value)| +---------------+ | (3.0, 4.0)| +---------------+ ``` ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#49786 from xinrong-meng/udt_arrow_udf. Authored-by: Xinrong Meng <[email protected]> Signed-off-by: Xinrong Meng <[email protected]>
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?
Turn on Arrow optimization for Python UDFs by default
Why are the changes needed?
Arrow optimization was introduced in 3.4.0. See SPARK-40307 for more context.
Arrow-optimized Python UDF is approximately 1.6 times faster than the original pickled Python UDF. More details can be found in this blog post.
In version 4.0.0, we propose enabling the optimization by default. If PyArrow is not installed, it will fall back to the original pickled Python UDF.
Does this PR introduce any user-facing change?
Yes
How was this patch tested?
Existing tests
Was this patch authored or co-authored using generative AI tooling?
No