You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2025/09/17 02:41:27 UTC
(spark) branch master updated: [SPARK-53594][PYTHON] Make arrow UDF respect user-specified eval type
This is an automated email from the ASF dual-hosted git repository.
ruifengz 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 3d87de318b62 [SPARK-53594][PYTHON] Make arrow UDF respect user-specified eval type
3d87de318b62 is described below
commit 3d87de318b62703785a07845899a0d30e32dfd29
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Sep 17 10:41:13 2025 +0800
[SPARK-53594][PYTHON] Make arrow UDF respect user-specified eval type
### What changes were proposed in this pull request?
Make arrow UDF respect user-specified eval type
### Why are the changes needed?
to be consistent with pandas UDF, which allow type hints be partially set:
```
In [4]: arrow_udf("long", ArrowUDFType.SCALAR_ITER)
...: def add_one(it: Iterator[pa.Array]):
...: for v in it:
...: yield pa.compute.add(v, 1)
```
before:
```
PySparkValueError: [TYPE_HINT_SHOULD_BE_SPECIFIED] Type hints for the return type should be specified; however, got (it: Iterator[pyarrow.lib.Array]).
```
after:
```
In [4]: arrow_udf("long", ArrowUDFType.SCALAR_ITER)
...: def add_one(it: Iterator[pa.Array]):
...: for v in it:
...: yield pa.compute.add(v, 1)
...:
/Users/ruifeng.zheng/Dev/spark/python/pyspark/sql/pandas/functions.py:766: UserWarning: It is preferred to specify type hints for arrow UDF instead of specifying arrow UDF type.
warnings.warn(
```
### Does this PR introduce _any_ user-facing change?
yes
### How was this patch tested?
new tests
### Was this patch authored or co-authored using generative AI tooling?
no
Closes #52351 from zhengruifeng/make_arrow_udf_respect_type.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
python/pyspark/sql/pandas/functions.py | 10 +++++++
python/pyspark/sql/tests/arrow/test_arrow_udf.py | 35 ++++++++++++++++++++++++
2 files changed, 45 insertions(+)
diff --git a/python/pyspark/sql/pandas/functions.py b/python/pyspark/sql/pandas/functions.py
index 79f3f75bd573..0c571ba0acfb 100644
--- a/python/pyspark/sql/pandas/functions.py
+++ b/python/pyspark/sql/pandas/functions.py
@@ -758,6 +758,16 @@ def _validate_vectorized_udf(f, evalType, kind: str = "pandas") -> int:
"in the future releases. See SPARK-28264 for more details.",
UserWarning,
)
+ elif evalType in [
+ PythonEvalType.SQL_SCALAR_ARROW_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
+ ]:
+ warnings.warn(
+ "It is preferred to specify type hints for "
+ "arrow UDF instead of specifying arrow UDF type.",
+ UserWarning,
+ )
elif evalType in [
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf.py b/python/pyspark/sql/tests/arrow/test_arrow_udf.py
index 693d7ba484d8..b56ed75acad9 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_udf.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_udf.py
@@ -19,6 +19,7 @@ import os
import time
import unittest
import datetime
+from typing import Iterator
from pyspark.sql.functions import arrow_udf, ArrowUDFType, PandasUDFType
from pyspark.sql import functions as F, Row
@@ -272,6 +273,40 @@ class ArrowUDFTestsMixin:
self.assertEqual(df.schema[0].dataType.simpleString(), "interval day to second")
self.assertEqual(df.first()[0], datetime.timedelta(microseconds=123))
+ def test_scalar_arrow_udf_with_specified_eval_type(self):
+ import pyarrow as pa
+
+ df = self.spark.range(10).selectExpr("id", "id as v")
+ expected = df.selectExpr("(v + 1) as plus_one").collect()
+
+ @arrow_udf("long", ArrowUDFType.SCALAR)
+ def plus_one(v: pa.Array):
+ return pa.compute.add(v, 1)
+
+ result1 = df.select(plus_one("v").alias("plus_one")).collect()
+ self.assertEqual(expected, result1)
+
+ @arrow_udf("long", ArrowUDFType.SCALAR_ITER)
+ def plus_one_iter(it: Iterator[pa.Array]):
+ for v in it:
+ yield pa.compute.add(v, 1)
+
+ result2 = df.select(plus_one_iter("v").alias("plus_one")).collect()
+ self.assertEqual(expected, result2)
+
+ def test_agg_arrow_udf_with_specified_eval_type(self):
+ import pyarrow as pa
+
+ df = self.spark.range(10).selectExpr("id", "id as v")
+ expected = df.selectExpr("max(v) as m").collect()
+
+ @arrow_udf("long", ArrowUDFType.GROUPED_AGG)
+ def calc_max(v: pa.Array):
+ return pa.compute.max(v)
+
+ result = df.select(calc_max("v").alias("m")).collect()
+ self.assertEqual(expected, result)
+
class ArrowUDFTests(ArrowUDFTestsMixin, ReusedSQLTestCase):
def setUp(self):
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org