This is an automated email from the ASF dual-hosted git repository.
gurwls223 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 939129ec01af [SPARK-50674][PYTHON] Fix check for ‘terminate’ method
existence in UDTF evaluation
939129ec01af is described below
commit 939129ec01af7f7b6dbec737f2d4149d2fc0d9a3
Author: Xinrong Meng <[email protected]>
AuthorDate: Fri Dec 27 19:11:40 2024 +0900
[SPARK-50674][PYTHON] Fix check for ‘terminate’ method existence in UDTF
evaluation
### What changes were proposed in this pull request?
Fix check for ‘terminate’ method existence in UDTF evaluation
### Why are the changes needed?
To ensure that UDTFs without a terminate method can still be used with
partitioning without causing an AttributeError.
Previously, udtf with partitioning will raise an AttributeError if the
terminate method is not defined, as shown below
```py
>>> from pyspark.sql.functions import udtf
>>> from pyspark.sql import Row
>>>
>>> udtf(returnType="a: int")
... class TestUDTF:
... def eval(self, row: Row):
... if row[0] > 5:
... yield row[0],
...
>>> spark.udtf.register("test_udtf", TestUDTF)
<pyspark.sql.udtf.UserDefinedTableFunction object at 0x10298a1d0>
>>> spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8))
PARTITION BY id)").show()
org.apache.spark.api.python.PythonException: Traceback (most recent call
last):
...
File "...pyspark/worker.py", line 1052, in eval
if self._udtf.terminate is not None:
AttributeError: 'TestUDTF' object has no attribute 'terminate'
```
However, the terminate method is not required in such cases.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #49299 from xinrong-meng/udtf_terminate.
Authored-by: Xinrong Meng <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
python/pyspark/sql/tests/test_udtf.py | 11 +++++++++++
python/pyspark/worker.py | 4 ++--
2 files changed, 13 insertions(+), 2 deletions(-)
diff --git a/python/pyspark/sql/tests/test_udtf.py
b/python/pyspark/sql/tests/test_udtf.py
index 31cd4c80370e..2424f74f93d3 100644
--- a/python/pyspark/sql/tests/test_udtf.py
+++ b/python/pyspark/sql/tests/test_udtf.py
@@ -2198,6 +2198,17 @@ class BaseUDTFTestsMixin:
],
)
+ def test_udtf_with_table_argument_and_partition_by_no_terminate(self):
+ func = self.udtf_for_table_argument() # a udtf with no terminate
method defined
+ self.spark.udtf.register("test_udtf", func)
+
+ assertDataFrameEqual(
+ self.spark.sql(
+ "SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8))
PARTITION BY id)"
+ ),
+ [Row(a=6), Row(a=7)],
+ )
+
def test_udtf_with_table_argument_and_partition_by_and_order_by(self):
class TestUDTF:
def __init__(self):
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index a11465e7a323..712f71d3861a 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -1049,7 +1049,7 @@ def read_udtf(pickleSer, infile, eval_type):
list(args) + list(kwargs.values())
)
if changed_partitions:
- if self._udtf.terminate is not None:
+ if hasattr(self._udtf, "terminate"):
result = self._udtf.terminate()
if result is not None:
for row in result:
@@ -1075,7 +1075,7 @@ def read_udtf(pickleSer, infile, eval_type):
self._eval_raised_skip_rest_of_input_table = True
def terminate(self) -> Iterator:
- if self._udtf.terminate is not None:
+ if hasattr(self._udtf, "terminate"):
return self._udtf.terminate()
return iter(())
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
