Skip to content

Conversation

@ueshin
Copy link
Member

@ueshin ueshin commented Aug 22, 2023

What changes were proposed in this pull request?

Supports named arguments in scalar Python/Pandas UDF.

For example:

>>> @udf("int")
... def test_udf(a, b):
...     return a + 10 * b
...
>>> spark.udf.register("test_udf", test_udf)

>>> spark.range(2).select(test_udf(b=col("id") * 10, a=col("id"))).show()
+---------------------------------+
|test_udf(b => (id * 10), a => id)|
+---------------------------------+
|                                0|
|                              101|
+---------------------------------+

>>> spark.sql("SELECT test_udf(b => id * 10, a => id) FROM range(2)").show()
+---------------------------------+
|test_udf(b => (id * 10), a => id)|
+---------------------------------+
|                                0|
|                              101|
+---------------------------------+

or:

>>> @pandas_udf("int")
... def test_udf(a, b):
...     return a + 10 * b
...
>>> spark.udf.register("test_udf", test_udf)

>>> spark.range(2).select(test_udf(b=col("id") * 10, a=col("id"))).show()
+---------------------------------+
|test_udf(b => (id * 10), a => id)|
+---------------------------------+
|                                0|
|                              101|
+---------------------------------+

>>> spark.sql("SELECT test_udf(b => id * 10, a => id) FROM range(2)").show()
+---------------------------------+
|test_udf(b => (id * 10), a => id)|
+---------------------------------+
|                                0|
|                              101|
+---------------------------------+

Why are the changes needed?

Now that named arguments support was added (#41796, #42020).

Scalar Python/Pandas UDFs can support it.

Does this PR introduce any user-facing change?

Yes, named arguments will be available for scalar Python/Pandas UDFs.

How was this patch tested?

Added related tests.

Was this patch authored or co-authored using generative AI tooling?

No.

@ueshin
Copy link
Member Author

ueshin commented Aug 22, 2023

cc @dtenedor @HyukjinKwon @xinrong-meng

Copy link
Contributor

@dtenedor dtenedor left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generally LGTM. I left one suggestion for a testing idea. Thanks a lot for adding this feature!!

/*
* Check if the named arguments:
* - don't have duplicated names
* - don't contain positional arguments
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
* - don't contain positional arguments
* - don't contain positional arguments after named arguments
* - all map to valid argument names from the function declaration

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The third item is not right for Python UDF/UDTF. Currently it relies on Python.


inputRDD.mapPartitions { iter =>
val context = TaskContext.get()
val contextAwareIterator = new ContextAwareIterator(context, iter)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch!

@ueshin
Copy link
Member Author

ueshin commented Aug 24, 2023

Thanks! merging to master.

@ueshin ueshin closed this in ce6b5f3 Aug 24, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants