This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon 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 15af8057fe66 [SPARK-58695][SQL][PYTHON] Support scalar pandas and
Arrow UDFs inside higher-order function lambdas
15af8057fe66 is described below
commit 15af8057fe663662a3b53ad6536b5496c416b5b6
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Tue Aug 11 12:37:42 2026 +0900
[SPARK-58695][SQL][PYTHON] Support scalar pandas and Arrow UDFs inside
higher-order function lambdas
### What changes were proposed in this pull request?
This is a follow-up to SPARK-27052, which let a plain (row-at-a-time)
Python UDF be used inside a higher-order function's lambda by lifting it out of
the lambda and applying it once to the whole array. That work handled only
`SQL_BATCHED_UDF` / `SQL_ARROW_BATCHED_UDF`, lifting them to a single
element-wise eval type `SQL_ARROW_ELEMENTWISE_UDF` (=102).
This PR extends the same rewrite to the vectorized scalar UDFs, so that e.g.
```python
pandas_udf("int")
def plus_one(s: pd.Series) -> pd.Series:
return s + 1
df.select(F.transform("values", lambda x: plus_one(x)))
```
works, instead of failing with
`[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]`. All four vectorized
scalar flavors are covered: scalar pandas (`pandas_udf`), scalar Arrow
(`arrow_udf`), and their iterator variants.
Each flavor lifts to its own dedicated element-wise eval type so the Python
worker keeps that flavor's native batching contract, mirroring the existing 102:
- `SQL_SCALAR_PANDAS_ELEMENTWISE_UDF` (=103)
- `SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF` (=104)
- `SQL_SCALAR_ARROW_ELEMENTWISE_UDF` (=105)
- `SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF` (=106)
`PythonUDF.liftedElementwiseEvalType` maps each base eval type to its
lifted counterpart, and `isElementwiseRewritableUDF` now accepts the vectorized
types (still excluding zero-argument, named-argument, and UDT shapes, which the
lift cannot preserve). In the Python worker, each argument arrives as
`array<T>` aligned with the iterated array; the worker flattens each list
column to its element column, runs the vectorized function once over that flat
column (a `pandas.Series` / `pyarrow [...]
### Why are the changes needed?
SPARK-27052 supported only plain Python UDFs inside HOF lambdas. Vectorized
scalar UDFs (pandas / Arrow) are the more performant and more common form, and
users expect them to compose with higher-order functions the same way.
### Does this PR introduce _any_ user-facing change?
Yes. A scalar pandas or Arrow UDF (and their iterator variants) can now be
used inside a higher-order function's lambda, e.g.
`df.select(F.transform("values", lambda x: pandas_plus_one(x)))`. Before this
change such a query failed at analysis with
`[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]`; now it evaluates. No
existing successful query changes behavior. This extends an unreleased feature
(SPARK-27052), so it is not a change relative to any released version.
### How was this patch tested?
- `pyspark.sql.tests.test_udf_in_higher_order_function` (classic and
Connect parity): new tests for all four vectorized flavors across `transform` /
`filter`, covering null arrays, null elements, empty arrays, and multi-batch
iterator regrouping.
- `ExtractPythonUDFFromLambdaSuite`: the previous "a pandas UDF inside a
lambda still fails" test is replaced by one asserting each vectorized flavor is
lifted to its element-wise eval type; the rewritable-predicate test now covers
all rewritable eval types and the base -> lifted mapping.
- `PythonUDFSuite` and `BatchEvalPythonExecSuite` pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
This pull request and its description were written by Isaac.
Closes #57899 from HyukjinKwon/SPARK-58695-pandas-arrow-udf-in-hof.
Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
.../org/apache/spark/api/python/PythonRunner.scala | 22 ++
.../docs/source/tutorial/sql/arrow_python_udf.rst | 5 +
python/pyspark/sql/_typing.pyi | 4 +
.../sql/tests/test_udf_in_higher_order_function.py | 317 +++++++++++++++++-
python/pyspark/util.py | 8 +
python/pyspark/worker.py | 362 +++++++++++++++++++++
.../spark/sql/catalyst/expressions/PythonUDF.scala | 54 ++-
.../sql/execution/python/ArrowEvalPythonExec.scala | 10 +
.../sql/execution/python/ArrowPythonRunner.scala | 6 +-
.../python/ExtractPythonUDFFromLambda.scala | 20 +-
.../sql/execution/python/ExtractPythonUDFs.scala | 11 +-
.../python/BatchEvalPythonExecSuite.scala | 21 ++
.../python/ExtractPythonUDFFromLambdaSuite.scala | 41 ++-
13 files changed, 854 insertions(+), 27 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 19a482850311..d7800b1147e1 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -53,7 +53,13 @@ private[spark] object PythonEvalType {
val SQL_ARROW_BATCHED_UDF = 101
// A scalar Python UDF applied element-wise over the elements of an array
column, used to
// support Python UDFs inside higher-order function lambdas. See
ExtractPythonUDFFromLambda.
+ // 102 lifts a row-at-a-time UDF (SQL_BATCHED_UDF / SQL_ARROW_BATCHED_UDF);
103-106 lift the
+ // vectorized scalar UDFs, preserving pandas- vs. Arrow-shaped batches and
the iterator contract.
val SQL_ARROW_ELEMENTWISE_UDF = 102
+ val SQL_SCALAR_PANDAS_ELEMENTWISE_UDF = 103
+ val SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF = 104
+ val SQL_SCALAR_ARROW_ELEMENTWISE_UDF = 105
+ val SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF = 106
val SQL_SCALAR_PANDAS_UDF = 200
val SQL_GROUPED_MAP_PANDAS_UDF = 201
@@ -90,6 +96,10 @@ private[spark] object PythonEvalType {
case SQL_BATCHED_UDF => "SQL_BATCHED_UDF"
case SQL_ARROW_BATCHED_UDF => "SQL_ARROW_BATCHED_UDF"
case SQL_ARROW_ELEMENTWISE_UDF => "SQL_ARROW_ELEMENTWISE_UDF"
+ case SQL_SCALAR_PANDAS_ELEMENTWISE_UDF =>
"SQL_SCALAR_PANDAS_ELEMENTWISE_UDF"
+ case SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF =>
"SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF"
+ case SQL_SCALAR_ARROW_ELEMENTWISE_UDF => "SQL_SCALAR_ARROW_ELEMENTWISE_UDF"
+ case SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF =>
"SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF"
case SQL_SCALAR_PANDAS_UDF => "SQL_SCALAR_PANDAS_UDF"
case SQL_GROUPED_MAP_PANDAS_UDF => "SQL_GROUPED_MAP_PANDAS_UDF"
case SQL_GROUPED_AGG_PANDAS_UDF => "SQL_GROUPED_AGG_PANDAS_UDF"
@@ -121,6 +131,18 @@ private[spark] object PythonEvalType {
case SQL_WINDOW_AGG_ARROW_UDF => "SQL_WINDOW_AGG_ARROW_UDF"
case SQL_GROUPED_AGG_ARROW_ITER_UDF => "SQL_GROUPED_AGG_ARROW_ITER_UDF"
}
+
+ // The eval types produced by ExtractPythonUDFFromLambda: a scalar UDF
lifted out of a
+ // higher-order function's lambda, which receives each argument as an
`array<T>` column and is
+ // applied element-wise inside the Python worker. See
ExtractPythonUDFFromLambda.
+ def isElementwiseUDF(evalType: Int): Boolean = evalType match {
+ case SQL_ARROW_ELEMENTWISE_UDF |
+ SQL_SCALAR_PANDAS_ELEMENTWISE_UDF |
+ SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF |
+ SQL_SCALAR_ARROW_ELEMENTWISE_UDF |
+ SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF => true
+ case _ => false
+ }
}
private[spark] object BasePythonRunner extends Logging {
diff --git a/python/docs/source/tutorial/sql/arrow_python_udf.rst
b/python/docs/source/tutorial/sql/arrow_python_udf.rst
index 0a0509f05003..fb1ff7515218 100644
--- a/python/docs/source/tutorial/sql/arrow_python_udf.rst
+++ b/python/docs/source/tutorial/sql/arrow_python_udf.rst
@@ -490,6 +490,11 @@ handle those cases inside the function itself. This path
is enabled by
``spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled`` (default
``true``); set it to
``false`` to reject such queries at analysis instead.
+This works for a plain Python UDF as well as for a vectorized scalar UDF - a
scalar pandas UDF
+(``pandas_udf``), a scalar Arrow UDF (``arrow_udf``), and their iterator
variants. A vectorized
+UDF still receives its native batch (a ``pandas.Series`` / ``pyarrow.Array``,
or an iterator of
+them) over the flattened array elements, and one output value is produced per
input element.
+
The Arrow data type of the returned ``pyarrow.Array`` should match the
declared ``returnType``.
When there is a mismatch, Spark will attempt to convert the returned data to
the expected type
using Arrow's safe casting, which raises an error on overflow or precision
loss.
diff --git a/python/pyspark/sql/_typing.pyi b/python/pyspark/sql/_typing.pyi
index 95bf27662ff1..a4713370aaf7 100644
--- a/python/pyspark/sql/_typing.pyi
+++ b/python/pyspark/sql/_typing.pyi
@@ -65,6 +65,10 @@ RowLike = TypeVar("RowLike", List[Any], Tuple[Any, ...],
pyspark.sql.types.Row)
SQLBatchedUDFType = Literal[100]
SQLArrowBatchedUDFType = Literal[101]
SQLArrowElementwiseUDFType = Literal[102]
+SQLScalarPandasElementwiseUDFType = Literal[103]
+SQLScalarPandasIterElementwiseUDFType = Literal[104]
+SQLScalarArrowElementwiseUDFType = Literal[105]
+SQLScalarArrowIterElementwiseUDFType = Literal[106]
SQLTableUDFType = Literal[300]
SQLArrowTableUDFType = Literal[301]
SQLArrowUDTFType = Literal[302]
diff --git a/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
b/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
index c027c4779fea..9d0beb024abe 100644
--- a/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
+++ b/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
@@ -643,20 +643,325 @@ class UDFInHigherOrderFunctionTestsMixin:
df.select(agg).collect()
self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF",
str(ctx.exception))
- def test_pandas_udf_in_lambda_still_fails(self):
- # A vectorized (pandas) UDF in a lambda is not supported.
+ def test_scalar_pandas_udf_in_lambda(self):
+ # A vectorized scalar pandas UDF is lifted and applied over the
flattened elements, so it
+ # still receives a pandas Series (its native contract) once per batch,
not per element.
import pandas as pd
from pyspark.sql.functions import pandas_udf
- df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, None],)], "values array<int>"
+ )
@pandas_udf(IntegerType())
def plus_one_pandas(s: pd.Series) -> pd.Series:
return s + 1
- with self.assertRaises(AnalysisException) as ctx:
- df.select(sf.transform("values", lambda x:
plus_one_pandas(x))).collect()
- self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_pandas(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, None],)],
+ )
+ # Arithmetic composition around the UDF result is ordinary JVM work.
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: plus_one_pandas(x) *
2).alias("r")),
+ [([4, 6, 8],), ([],), (None,), ([22, None],)],
+ )
+
+ def test_scalar_arrow_udf_in_lambda(self):
+ # A vectorized scalar Arrow UDF is lifted the same way; it takes and
returns a pyarrow
+ # Array over the flattened elements.
+ import pyarrow as pa
+ from pyspark.sql.functions import arrow_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, None],)], "values array<int>"
+ )
+
+ @arrow_udf(IntegerType())
+ def plus_one_arrow(a: pa.Array) -> pa.Array:
+ return pa.compute.add(a, 1)
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_arrow(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, None],)],
+ )
+ assertDataFrameEqual(
+ df.select(sf.filter("values", lambda x: plus_one_arrow(x) >
2).alias("r")),
+ [([2, 3],), ([],), (None,), ([10],)],
+ )
+
+ def test_scalar_pandas_iter_udf_in_lambda(self):
+ # A scalar iterator pandas UDF keeps its iterator contract: it
consumes and produces an
+ # iterator of Series. The worker feeds it the flattened elements and
re-groups the streamed
+ # results back into arrays positionally, so output batch boundaries
need not match input.
+ from typing import Iterator
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, 20],), ([5],)], "values
array<int>"
+ )
+
+ @pandas_udf(IntegerType())
+ def plus_one_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield s + 1
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_iter(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, 21],), ([6],)],
+ )
+
+ def test_scalar_arrow_iter_udf_in_lambda(self):
+ # A scalar iterator Arrow UDF, lifted the same way as the pandas
iterator variant.
+ from typing import Iterator
+ import pyarrow as pa
+ from pyspark.sql.functions import arrow_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, 20],), ([5],)], "values
array<int>"
+ )
+
+ @arrow_udf(IntegerType())
+ def plus_one_arrow_iter(it: Iterator[pa.Array]) -> Iterator[pa.Array]:
+ for a in it:
+ yield pa.compute.add(a, 1)
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_arrow_iter(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, 21],), ([6],)],
+ )
+
+ def test_scalar_pandas_iter_udf_multiple_arguments_differ_in_type(self):
+ # A two-argument iterator UDF whose arguments have different element
types: the array
+ # element (int) and an outer column (string) that the rewrite repeats
into an aligned
+ # array. Each argument must be flattened with its own element type,
not the first's.
+ from typing import Iterator, Tuple
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3], "a"), ([], "b"), (None, "c"), ([10], "d")],
+ "values array<int>, tag string",
+ )
+
+ @pandas_udf(StringType())
+ def tag_each(it: Iterator[Tuple[pd.Series, pd.Series]]) ->
Iterator[pd.Series]:
+ for x, t in it:
+ yield t + x.astype("string")
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: tag_each(x,
sf.col("tag"))).alias("r")),
+ [(["a1", "a2", "a3"],), ([],), (None,), (["d10"],)],
+ )
+
+ def test_scalar_pandas_iter_udf_timestamp_return_type(self):
+ # A timestamp-returning pandas iterator UDF with a non-UTC session
timezone: the result
+ # chunks are typed with the session timezone, so the streamed buffer
must take its type
+ # from the first chunk rather than assuming UTC, or pa.concat_arrays
would fail. Assert both
+ # against the equivalent non-iterator pandas UDF (isolates the concat
fix) and against a
+ # native Spark expression computing the same instants (so a timezone
bug common to both UDF
+ # paths would still be caught, while going through identical
driver-collection semantics).
+ from typing import Iterator
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.sql.types import TimestampType
+
+ with self.sql_conf({"spark.sql.session.timeZone":
"America/Los_Angeles"}):
+ df = self.spark.createDataFrame([([1, 2],), (None,), ([3],)],
"values array<int>")
+
+ def compute(x):
+ return pd.to_datetime(x, unit="D", origin="2020-01-01")
+
+ @pandas_udf(TimestampType())
+ def to_ts(s: pd.Series) -> pd.Series:
+ return compute(s)
+
+ @pandas_udf(TimestampType())
+ def to_ts_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield compute(s)
+
+ iter_df = df.select(sf.transform("values", lambda x:
to_ts_iter(x)).alias("r"))
+ # Native equivalent: pandas interprets the tz-naive origin in the
session timezone, so
+ # `timestamp_add(DAY, x, TIMESTAMP '2020-01-01 00:00:00')` (also
session-local) matches.
+ native_df = df.select(
+ sf.transform(
+ "values",
+ lambda x: sf.timestamp_add("DAY", x,
sf.lit("2020-01-01").cast("timestamp")),
+ ).alias("r")
+ )
+ # Consistent with the non-iterator pandas UDF, and with the native
instants.
+ assertDataFrameEqual(
+ iter_df,
+ df.select(sf.transform("values", lambda x:
to_ts(x)).alias("r")),
+ )
+ assertDataFrameEqual(iter_df, native_df)
+
+ def test_scalar_iter_udf_over_all_empty_and_null_partition(self):
+ # SPARK-58695: when a whole partition holds only empty/null arrays,
the flattened inputs are
+ # all zero-length and a skip-empty iterator UDF yields no chunks.
Those rows still need one
+ # (empty / null) output row each, or the positional JVM join drops
them silently. Cover both
+ # the pandas and Arrow iterator flavors.
+ from typing import Iterator
+ import pandas as pd
+ import pyarrow as pa
+ from pyspark.sql.functions import pandas_udf, arrow_udf
+
+ # Single partition so the whole batch is empty/null arrays.
+ df = self.spark.createDataFrame(
+ [([],), (None,), ([],), (None,)], "values array<int>"
+ ).coalesce(1)
+
+ @pandas_udf(IntegerType())
+ def skip_empty_pandas(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ if len(s) == 0:
+ continue
+ yield s + 1
+
+ @arrow_udf(IntegerType())
+ def skip_empty_arrow(it: Iterator[pa.Array]) -> Iterator[pa.Array]:
+ for a in it:
+ if len(a) == 0:
+ continue
+ yield pa.compute.add(a, 1)
+
+ for f in (skip_empty_pandas, skip_empty_arrow):
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: f(x)).alias("r")),
+ [([],), (None,), ([],), (None,)],
+ )
+
+ def test_scalar_pandas_iter_udf_timestamp_after_empty_batch(self):
+ # A zero-length result chunk (from an input batch holding only
empty/null arrays) must not
+ # pin the output stream's timestamp type to the UTC-typed default: the
rows it emits and
+ # the rows emitted from a later real chunk (typed with the session
timezone) would then
+ # disagree, and the Arrow stream writer would reject the second output
batch. Assert against
+ # the equivalent non-iterator pandas UDF (identical driver-collection
semantics), so the
+ # check proves the schema fix without depending on absolute timezone
offsets.
+ from typing import Iterator
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.sql.types import TimestampType
+
+ with self.sql_conf(
+ {
+ "spark.sql.session.timeZone": "America/Los_Angeles",
+ # One row per Arrow batch so the empty-array row forms its own
(first) batch.
+ "spark.sql.execution.arrow.maxRecordsPerBatch": "1",
+ }
+ ):
+ df = self.spark.createDataFrame([([],), ([1],)], "values
array<int>").coalesce(1)
+
+ def compute(x):
+ return pd.to_datetime(x, unit="D", origin="2020-01-01")
+
+ @pandas_udf(TimestampType())
+ def to_ts(s: pd.Series) -> pd.Series:
+ return compute(s)
+
+ @pandas_udf(TimestampType())
+ def to_ts_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield compute(s)
+
+ # Without the fix this raises ArrowInvalid ("different schema")
writing the second
+ # output batch; with it the iterator result matches the
non-iterator pandas UDF.
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
to_ts_iter(x)).alias("r")),
+ df.select(sf.transform("values", lambda x:
to_ts(x)).alias("r")),
+ )
+
+ def test_scalar_pandas_udf_struct_element_return_type(self):
+ # A vectorized pandas UDF returning a struct element (a
pandas.DataFrame per batch) inside a
+ # lambda. Covers the struct-DataFrame result path of the element-wise
pandas branch.
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.sql.types import StructType, StructField
+
+ df = self.spark.createDataFrame([([1, 2, 3],), ([],), (None,)],
"values array<int>")
+ ret = StructType([StructField("v", IntegerType()), StructField("neg",
IntegerType())])
+
+ @pandas_udf(ret)
+ def to_struct(s: pd.Series) -> pd.DataFrame:
+ return pd.DataFrame({"v": s, "neg": -s})
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
to_struct(x)).alias("r")),
+ [
+ ([(1, -1), (2, -2), (3, -3)],),
+ ([],),
+ (None,),
+ ],
+ )
+
+ def test_chained_vectorized_udfs_in_lambda(self):
+ # Nested calls f(g(x)) inside a lambda: g is lifted first, then f
consumes g's array result.
+ # Cover both the non-iterator and iterator vectorized flavors.
+ from typing import Iterator
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10],)], "values array<int>"
+ )
+
+ @pandas_udf(IntegerType())
+ def plus_one(s: pd.Series) -> pd.Series:
+ return s + 1
+
+ @pandas_udf(IntegerType())
+ def times_two(s: pd.Series) -> pd.Series:
+ return s * 2
+
+ @pandas_udf(IntegerType())
+ def plus_one_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield s + 1
+
+ @pandas_udf(IntegerType())
+ def times_two_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield s * 2
+
+ # (x + 1) * 2, verified against the equivalent native expression.
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
times_two(plus_one(x))).alias("r")),
+ df.select(sf.transform("values", lambda x: (x + 1) *
2).alias("r")),
+ )
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("values", lambda x:
times_two_iter(plus_one_iter(x))).alias("r")
+ ),
+ df.select(sf.transform("values", lambda x: (x + 1) *
2).alias("r")),
+ )
+
+ def test_scalar_iter_udf_struct_element_return_type(self):
+ # A scalar iterator pandas UDF returning a struct element (a
pandas.DataFrame per batch)
+ # inside a lambda. Covers the struct-DataFrame result path of the
iterator element-wise
+ # branch (Iterator[pd.DataFrame] contract).
+ from typing import Iterator
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.sql.types import StructType, StructField
+
+ df = self.spark.createDataFrame([([1, 2, 3],), ([],), (None,)],
"values array<int>")
+ ret = StructType([StructField("v", IntegerType()), StructField("neg",
IntegerType())])
+
+ @pandas_udf(ret)
+ def to_struct_iter(it: Iterator[pd.Series]) -> Iterator[pd.DataFrame]:
+ for s in it:
+ yield pd.DataFrame({"v": s, "neg": -s})
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
to_struct_iter(x)).alias("r")),
+ [
+ ([(1, -1), (2, -2), (3, -3)],),
+ ([],),
+ (None,),
+ ],
+ )
def test_non_arrow_udf_is_also_supported(self):
# A UDF created with useArrow=False is still rewritable; the generated
array wrapper
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index 640728725ee3..76be9bcc9998 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -91,6 +91,10 @@ if typing.TYPE_CHECKING:
from pyspark.sql._typing import (
SQLArrowBatchedUDFType,
SQLArrowElementwiseUDFType,
+ SQLScalarPandasElementwiseUDFType,
+ SQLScalarPandasIterElementwiseUDFType,
+ SQLScalarArrowElementwiseUDFType,
+ SQLScalarArrowIterElementwiseUDFType,
SQLArrowTableUDFType,
SQLBatchedUDFType,
SQLTableUDFType,
@@ -664,6 +668,10 @@ class PythonEvalType:
SQL_BATCHED_UDF: "SQLBatchedUDFType" = 100
SQL_ARROW_BATCHED_UDF: "SQLArrowBatchedUDFType" = 101
SQL_ARROW_ELEMENTWISE_UDF: "SQLArrowElementwiseUDFType" = 102
+ SQL_SCALAR_PANDAS_ELEMENTWISE_UDF: "SQLScalarPandasElementwiseUDFType" =
103
+ SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF:
"SQLScalarPandasIterElementwiseUDFType" = 104
+ SQL_SCALAR_ARROW_ELEMENTWISE_UDF: "SQLScalarArrowElementwiseUDFType" = 105
+ SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF:
"SQLScalarArrowIterElementwiseUDFType" = 106
SQL_SCALAR_PANDAS_UDF: "PandasScalarUDFType" = 200
SQL_GROUPED_MAP_PANDAS_UDF: "PandasGroupedMapUDFType" = 201
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index c579c719c6e2..4e77c630777d 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -498,6 +498,11 @@ def _is_iter_based(eval_type: int) -> bool:
return eval_type in (
PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
+ # Iterator UDFs lifted out of a higher-order function lambda keep the
iterator contract:
+ # the user function still consumes and produces an iterator of
batches; the worker only
+ # feeds it the flattened elements and re-nests the results. See
ExtractPythonUDFFromLambda.
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
PythonEvalType.SQL_MAP_ARROW_ITER_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
@@ -607,6 +612,7 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
if eval_type in (
PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_ARROW_BATCHED_UDF,
):
func = profiling_func
@@ -624,6 +630,10 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
if eval_type in (
PythonEvalType.SQL_ARROW_BATCHED_UDF,
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
@@ -1819,10 +1829,100 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
return mapper, None, ser, ser
+def _elementwise_renest(flat_values, shape_lengths, is_large):
+ """Re-nest a flat Array of per-element results into an ``array<R>`` column.
+
+ ``flat_values`` holds the results for every non-null element in order;
``shape_lengths`` is
+ the per-array element count of the iterated argument (``None`` for a null
array, which stays
+ null and consumes no elements). ``is_large`` preserves the input's list
width (``ListArray``
+ with int32 offsets vs. ``LargeListArray`` with int64).
+
+ Shared by the vectorized element-wise worker paths (scalar pandas / Arrow
and their iterator
+ variants) that back Python UDFs inside higher-order function lambdas. See
+ ``ExtractPythonUDFFromLambda``.
+ """
+ import pyarrow as pa
+
+ offsets = [0]
+ running = 0
+ mask = []
+ for n in shape_lengths:
+ mask.append(n is None)
+ if n is not None:
+ running += n
+ offsets.append(running)
+ list_cls = pa.LargeListArray if is_large else pa.ListArray
+ offsets_arr = pa.array(offsets, type=pa.int64() if is_large else
pa.int32())
+ null_mask = pa.array(mask, type=pa.bool_())
+ return list_cls.from_arrays(offsets_arr, flat_values, mask=null_mask)
+
+
+def _elementwise_flatten_column(flat, element_type, is_pandas, runner_conf):
+ """Adapt one already-flattened ``array<T>`` element column to the
vectorized fn's input.
+
+ ``flat`` is the flattened element ``pa.Array`` (the caller flattens once
per batch and shares it
+ across fused UDFs). Returns it unchanged for the Arrow flavor, or
converted to a pandas Series /
+ DataFrame with the element type ``T`` for the pandas flavor. Shared by the
vectorized
+ element-wise worker paths that back Python UDFs inside higher-order
function lambdas. See
+ ``ExtractPythonUDFFromLambda``.
+ """
+ if not is_pandas:
+ return flat
+ from pyspark.sql.conversion import ArrowArrayToPandasConversion
+
+ return ArrowArrayToPandasConversion.convert(
+ flat,
+ element_type,
+ timezone=runner_conf.timezone,
+ struct_in_pandas="dict",
+ ndarray_as_list=False,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ df_for_struct=True,
+ )
+
+
+def _elementwise_result_to_arrow(result, return_type, arrow_element_type,
is_pandas, runner_conf):
+ """Convert one vectorized UDF result over the flat elements to a single
flat Arrow Array.
+
+ ``result`` is a pandas Series / DataFrame (pandas flavor) or a
``pa.Array`` (Arrow flavor); the
+ returned array holds one element per input element. The Arrow flavor is
coerced to
+ ``arrow_element_type`` (UTC-typed); the pandas flavor is typed by
``PandasToArrowConversion``
+ using the session timezone, so its timestamp type may differ from
``arrow_element_type`` -
+ callers that concatenate results must take the type from the returned
array, not assume UTC.
+ Shared by the vectorized element-wise worker paths. See
``ExtractPythonUDFFromLambda``.
+ """
+ import pyarrow as pa
+
+ if is_pandas:
+ batch = PandasToArrowConversion.convert(
+ [result],
+ StructType([StructField("_0", return_type)]),
+ timezone=runner_conf.timezone,
+ safecheck=runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=runner_conf.use_large_var_types,
+ assign_cols_by_name=runner_conf.assign_cols_by_name,
+
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ )
+ else:
+ batch = ArrowBatchTransformer.enforce_schema(
+ pa.RecordBatch.from_arrays([result], ["_0"]),
+ pa.schema([pa.field("_0", arrow_element_type)]),
+ safecheck=runner_conf.safecheck,
+ )
+ # PandasToArrowConversion / enforce_schema both return a pa.RecordBatch,
so column(0) is a
+ # single pa.Array (never a ChunkedArray).
+ return batch.column(0)
+
+
def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
if eval_type in (
PythonEvalType.SQL_ARROW_BATCHED_UDF,
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF,
PythonEvalType.SQL_SCALAR_ARROW_UDF,
PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
@@ -3122,6 +3222,268 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# profiling is not supported for UDF
return func, None, ser, ser
+ if eval_type in (
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ ):
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+
+ require_minimum_pyarrow_version()
+
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ # A scalar pandas or Arrow UDF lifted out of a higher-order function's
lambda by
+ # ExtractPythonUDFFromLambda. Each argument arrives as ``array<T>``
aligned with the
+ # iterated array. We flatten each list column to its element column,
run the *vectorized*
+ # function once over that flat column (so it still receives a pandas
Series / DataFrame or a
+ # pa.Array, its native contract), then re-nest the flat result to
``array<R>`` using the
+ # input's offsets - one row in, one row out, one Python round trip per
batch.
+ is_pandas = eval_type ==
PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF
+
+ udf_infos = []
+ for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in
udfs:
+ wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
+ udf_func, udf_args_offsets, udf_kwargs_offsets
+ )
+ # The UDF returns one value per element, so its declared return
type is the element
+ # type of the ``array<R>`` this operator produces.
+ arrow_element_type = to_arrow_type(
+ udf_return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ )
+ udf_infos.append(
+ (wrapped_func, args_kwargs_offsets, udf_return_type,
arrow_element_type)
+ )
+ col_names = [f"_{i}" for i in range(len(udfs))]
+
+ if is_pandas:
+ import pandas as pd
+
+ # Each argument is ``array<T>``; the vectorized function must see the
element type ``T``.
+ element_types = [f.dataType.elementType for f in eval_conf.input_type]
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ for input_batch in data:
+ # Flatten each list column to its element array once per batch
and share it across
+ # fused UDFs (the 102 path does the same), rather than
re-flattening per UDF.
+ flat_columns = [
+ _elementwise_flatten_column(
+ col.flatten(), element_types[i], is_pandas, runner_conf
+ )
+ for i, col in enumerate(input_batch.itercolumns())
+ ]
+
+ output_arrays = []
+ for wrapped_func, offsets, return_type, arrow_element_type in
udf_infos:
+ # Re-nest by this UDF's first argument's shape. Different
UDFs in one operator
+ # may iterate differently shaped arrays, so each re-nests
by its own argument.
+ shape = input_batch.column(offsets[0])
+ shape_lengths = pc.list_value_length(shape).to_pylist()
+ total_elements = sum(n for n in shape_lengths if n is not
None)
+
+ result = wrapped_func(*[flat_columns[o] for o in offsets])
+ if is_pandas:
+ if not hasattr(result, "__len__"):
+ pd_type = (
+ "pandas.DataFrame"
+ if isinstance(return_type, StructType)
+ else "pandas.Series"
+ )
+ raise PySparkTypeError(
+ errorClass="UDF_RETURN_TYPE",
+ messageParameters={
+ "expected": pd_type,
+ "actual": type(result).__name__,
+ },
+ )
+ # struct return type must be a DataFrame (matches the
base pandas path).
+ if isinstance(return_type, StructType) and not
isinstance(
+ result, pd.DataFrame
+ ):
+ raise PySparkValueError(
+ "Invalid return type. Please make sure that
the UDF returns a "
+ "pandas.DataFrame when the specified return
type is StructType."
+ )
+ # Verify the flat length before re-nesting so a
wrong-length result raises
+ # the friendly RESULT_ROWS_MISMATCH rather than an
opaque pyarrow error.
+ verify_result_row_count(len(result), total_elements)
+ else:
+ # Arrow flavor: a non-array-like result (e.g. a bare
int) raises the
+ # friendly UDF_RETURN_TYPE rather than a bare
TypeError from len(), matching
+ # the base SQL_SCALAR_ARROW_UDF path, and also checks
the flat length.
+ verify_scalar_result(result, total_elements)
+
+ flat_arr = _elementwise_result_to_arrow(
+ result, return_type, arrow_element_type, is_pandas,
runner_conf
+ )
+ nested = _elementwise_renest(
+ flat_arr, shape_lengths,
pa.types.is_large_list(shape.type)
+ )
+ output_arrays.append(nested)
+
+ yield pa.RecordBatch.from_arrays(output_arrays, col_names)
+
+ # profiling is not supported for UDF
+ return func, None, ser, ser
+
+ if eval_type in (
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
+ ):
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+
+ require_minimum_pyarrow_version()
+
+ import collections
+
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ is_pandas = eval_type ==
PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF
+ if is_pandas:
+ import pandas as pd
+
+ assert num_udfs == 1, "One SCALAR_*_ITER_ELEMENTWISE UDF expected
here."
+ udf_func, args_offsets, kwargs_offsets, return_type = udfs[0]
+ assert not kwargs_offsets, "Iterator UDFs do not take keyword
arguments."
+
+ # A scalar iterator UDF (pandas or Arrow) lifted out of a higher-order
function's lambda.
+ # The user function keeps its iterator contract: it consumes an
iterator of batches and
+ # yields an iterator of batches, one output value per input value. We
preserve that by
+ # feeding it the *flattened* elements of each input batch and, since
the JVM joins UDF
+ # output to input positionally by row (one ``array<R>`` per input
``array<T>`` row, in
+ # order), buffering a FIFO of the per-row element counts to re-group
the streamed flat
+ # results back into arrays. Output batch boundaries need not match
input ones.
+ # Each argument is ``array<T>``; the pandas function must see each
argument's own element
+ # type ``T`` (arguments may differ, e.g. an outer column repeated into
an aligned array).
+ element_types = [f.dataType.elementType for f in eval_conf.input_type]
+ arrow_element_type = to_arrow_type(
+ return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ )
+ is_large = None # set from the first input batch; the list width is
uniform per column.
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ # FIFO of per-row element counts (None for a null array) awaiting
their flat results,
+ # and the flat elements produced so far but not yet enough to
complete the head shapes.
+ pending_shapes: "collections.deque" = collections.deque()
+ num_input_elements = 0
+
+ def extract_flat(batch: pa.RecordBatch):
+ nonlocal is_large, num_input_elements
+ shape = batch.column(args_offsets[0])
+ if is_large is None:
+ is_large = pa.types.is_large_list(shape.type)
+ pending_shapes.append(pc.list_value_length(shape).to_pylist())
+ # Flatten each argument to its element column; the user
function sees the flat
+ # elements as a pandas Series / DataFrame (pandas) or a
pa.Array (Arrow). Each
+ # argument is converted with its own element type.
+ flat_cols = [
+ _elementwise_flatten_column(
+ batch.column(o).flatten(), element_types[o],
is_pandas, runner_conf
+ )
+ for o in args_offsets
+ ]
+ num_input_elements += len(flat_cols[0])
+ return flat_cols[0] if len(flat_cols) == 1 else
tuple(flat_cols)
+
+ flat_args_iter = map(extract_flat, data)
+
+ if not is_pandas:
+ verified_iter = verify_return_type(
+ udf_func(flat_args_iter),
+ Iterator[pa.Array], # type: ignore[type-abstract]
+ )
+ else:
+ pandas_iter_type = (
+ Iterator[pd.DataFrame]
+ if isinstance(return_type, StructType)
+ else Iterator[pd.Series]
+ )
+ verified_iter = verify_return_type(udf_func(flat_args_iter),
pandas_iter_type)
+
+ # Buffer the streamed flat element results and emit an
``array<R>`` row as soon as the
+ # shape at the head of the FIFO is fully covered. A row whose
length is 0 (an empty
+ # array) or None (a null array) needs no elements, so it is
emitted immediately even
+ # before any chunk arrives - this matters when a whole partition
is empty/null arrays
+ # and the UDF yields nothing, otherwise those rows would be
dropped by the positional
+ # JVM join. Chunks are held in a list and concatenated only when a
shape spans more than
+ # one, so a UDF that yields once per input batch (the common case)
never re-copies the
+ # buffer. ``empty_type`` supplies the element type for a
zero-length emit; it tracks the
+ # most recent chunk's type (even a zero-length chunk carries the
flavor's type - the
+ # pandas flavor types timestamps with the session timezone),
falling back to the
+ # UTC-typed ``arrow_element_type`` only before any chunk arrives,
so all emitted batches
+ # share one schema.
+ pending_chunks: "list" = []
+ pending_len = 0
+ empty_type = arrow_element_type
+ num_output_elements = 0
+
+ def emit_ready():
+ nonlocal pending_chunks, pending_len
+ while pending_shapes:
+ lengths = pending_shapes[0]
+ needed = sum(n for n in lengths if n is not None)
+ if needed > pending_len:
+ break
+ pending_shapes.popleft()
+ if needed == 0:
+ flat = pa.nulls(0, type=empty_type)
+ else:
+ combined = (
+ pending_chunks[0]
+ if len(pending_chunks) == 1
+ else pa.concat_arrays(pending_chunks)
+ )
+ flat = combined.slice(0, needed)
+ remainder = combined.slice(needed)
+ pending_chunks = [remainder] if len(remainder) else []
+ pending_len -= needed
+ nested = _elementwise_renest(flat, lengths, bool(is_large))
+ yield pa.RecordBatch.from_arrays([nested], ["_0"])
+
+ def process_results():
+ nonlocal pending_chunks, pending_len, empty_type,
num_output_elements
+ for result in verified_iter:
+ if is_pandas:
+ verify_pandas_result(
+ result,
+ return_type,
+ assign_cols_by_name=True,
+ truncate_return_schema=True,
+ )
+ chunk = _elementwise_result_to_arrow(
+ result, return_type, arrow_element_type, is_pandas,
runner_conf
+ )
+ num_output_elements += len(chunk)
+ # Fail fast if the UDF over-produces, before the buffer
grows unbounded (the
+ # base iterator paths do the same via
verify_output_row_limit).
+ if num_output_elements > num_input_elements:
+ raise PySparkRuntimeError(
+ errorClass="OUTPUT_EXCEEDS_INPUT_ROWS",
messageParameters={}
+ )
+ # Even a zero-length chunk carries the flavor's element
type (the pandas flavor
+ # types timestamps with the session timezone), so always
take it: otherwise
+ # rows emitted for an all-empty batch before the first
non-empty chunk would use
+ # the UTC-typed default and disagree with later batches,
breaking the output
+ # stream's single-schema contract.
+ empty_type = chunk.type
+ if len(chunk):
+ pending_chunks.append(chunk)
+ pending_len += len(chunk)
+ yield from emit_ready()
+
+ # The iterator is exhausted: every input row's flat elements
must have arrived.
+ verify_result_row_count(num_output_elements,
num_input_elements)
+ # Flush any residual all-empty / all-null rows (they consume
no elements).
+ if pending_shapes:
+ yield from emit_ready()
+ verify_iterator_exhausted(flat_args_iter)
+
+ yield from process_results()
+
+ # profiling is not supported for UDF
+ return func, None, ser, ser
+
if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF:
import pandas as pd
import pyarrow as pa
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
index ec7e7832bb65..e421d946a669 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
@@ -38,8 +38,13 @@ object PythonUDF {
PythonEvalType.SQL_ARROW_BATCHED_UDF,
// Element-wise UDFs are row-shaped from the plan's point of view: one
array column in, one
// array column out per row. They are extracted by `ExtractPythonUDFs`
like any other scalar
- // UDF; only the Python worker treats them element-wise.
+ // UDF; only the Python worker treats them element-wise. One eval type per
lifted flavor keeps
+ // the worker's pandas- vs. Arrow-shaped batching and the iterator
contract distinct.
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
PythonEvalType.SQL_SCALAR_ARROW_UDF,
@@ -54,15 +59,20 @@ object PythonUDF {
* Whether `e` is a Python UDF that can be lifted out of a higher-order
function's lambda by
* `ExtractPythonUDFFromLambda`, which applies it to the whole array outside
the lambda.
*
- * Only the row-at-a-time eval types qualify, since the rule lifts a UDF
that takes one value per
- * call. Otherwise-eligible shapes are excluded because the rewrite cannot
preserve them:
+ * Both the row-at-a-time eval types (plain and Arrow batched) and the
vectorized scalar eval
+ * types (scalar pandas / Arrow and their iterator variants) qualify: the
rule lifts the UDF
+ * structurally over `array<T>` arguments, and the Python worker flattens
each array, invokes the
+ * function on the flat element column with its own batching contract, and
re-nests. See
+ * [[liftedElementwiseEvalType]] for the mapping to the eval type the lifted
UDF runs under.
+ *
+ * Otherwise-eligible shapes are excluded because the rewrite cannot
preserve them:
* - a zero-argument call, `f()`: the lift turns each argument into an
aligned array, so with no
* argument there is no array to carry the iterated shape, and the
element-wise UDF would
* reach the worker with no input column and crash there instead of
failing analysis;
* - a call with named arguments: its `NamedArgumentExpression` children
would be buried inside
* the generated `ArrayTransform`, losing the kwargs mapping the runner
derives from the
* direct children;
- * - a UDF whose argument or return type involves a UDT: the lift forces
the Arrow element-wise
+ * - a UDF whose argument or return type involves a UDT: the lift forces
an Arrow element-wise
* eval type, which has no UDT fallback (unlike `correctEvalType`'s
Arrow -> pickle path), so
* it would fail at runtime instead of at analysis.
* All keep the previous behavior (an analysis error) rather than being
rewritten.
@@ -72,8 +82,7 @@ object PythonUDF {
*/
def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
case udf: PythonUDF =>
- (udf.evalType == PythonEvalType.SQL_BATCHED_UDF ||
- udf.evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) &&
+ isElementwiseRewritableEvalType(udf.evalType) &&
udf.children.nonEmpty &&
!udf.children.exists(_.isInstanceOf[NamedArgumentExpression]) &&
!containsUDT(udf.dataType) &&
@@ -81,6 +90,38 @@ object PythonUDF {
case _ => false
}
+ private def isElementwiseRewritableEvalType(evalType: Int): Boolean =
evalType match {
+ case PythonEvalType.SQL_BATCHED_UDF |
+ PythonEvalType.SQL_ARROW_BATCHED_UDF |
+ PythonEvalType.SQL_SCALAR_PANDAS_UDF |
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF |
+ PythonEvalType.SQL_SCALAR_ARROW_UDF |
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF => true
+ case _ => false
+ }
+
+ /**
+ * The eval type a rewritable UDF runs under once lifted out of the lambda.
Each maps to the
+ * element-wise flavor that preserves its worker contract: the row-at-a-time
types share the one
+ * pickle-based element-wise path, while each vectorized scalar type keeps
its own pandas- vs.
+ * Arrow-shaped batching and iterator behavior. `evalType` must satisfy
+ * [[isElementwiseRewritableEvalType]].
+ */
+ def liftedElementwiseEvalType(evalType: Int): Int = evalType match {
+ case PythonEvalType.SQL_BATCHED_UDF | PythonEvalType.SQL_ARROW_BATCHED_UDF
=>
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF
+ case PythonEvalType.SQL_SCALAR_PANDAS_UDF =>
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF
+ case PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF =>
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF
+ case PythonEvalType.SQL_SCALAR_ARROW_UDF =>
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF
+ case PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF =>
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF
+ case other =>
+ throw internalError(s"Not a rewritable elementwise UDF eval type:
$other")
+ }
+
/**
* Whether every Python UDF in `hof`'s lambdas can be lifted out by
`ExtractPythonUDFFromLambda`.
* Used by `CheckAnalysis` to decide whether to reject the plan. These
shapes cannot be rewritten:
@@ -88,7 +129,6 @@ object PythonUDF {
* `i` is not a real column. (A UDF in a nested *argument*,
`transform(arr, x ->
* transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
* - a UDF in `aggregate` / `reduce`: the fold is sequential, so it sees
earlier steps' outputs;
- * - a vectorized (scalar pandas / arrow) UDF, which is not supported;
* - a *nondeterministic iterated argument*, `filter(shuffle(arr), x ->
f(x))`: the rewrite
* references that argument several times (the carrier's `c0`, each
lifted UDF's argument, the
* `map_keys`/`map_values` desugar, the pairwise `array_sort` path), and
nondeterministic
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
index 7c7ed9c50ab6..5206b15877c5 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
@@ -71,6 +71,12 @@ private[spark] class BatchIterator[T](iter: Iterator[T],
batchSize: Int)
* <li> SQL_SCALAR_ARROW_ITER_UDF for Scalar Iterator Arrow UDF
* <li> SQL_SCALAR_PANDAS_UDF for Scalar Pandas UDF
* <li> SQL_SCALAR_PANDAS_ITER_UDF for Scalar Iterator Pandas UDF
+ * <li> SQL_ARROW_ELEMENTWISE_UDF for a row-at-a-time UDF lifted out of a
higher-order function's
+ * lambda (see ExtractPythonUDFFromLambda)
+ * <li> SQL_SCALAR_PANDAS_ELEMENTWISE_UDF for a Scalar Pandas UDF lifted out
of such a lambda
+ * <li> SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF for a Scalar Iterator Pandas
UDF lifted out of one
+ * <li> SQL_SCALAR_ARROW_ELEMENTWISE_UDF for a Scalar Arrow UDF lifted out
of such a lambda
+ * <li> SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF for a Scalar Iterator Arrow
UDF lifted out of one
* </ul>
*
*/
@@ -163,6 +169,10 @@ case class ArrowEvalPythonExec(
Array(
PythonEvalType.SQL_ARROW_BATCHED_UDF,
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_ARROW_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF,
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
index 4c1c472cb10d..d0383c32a889 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
@@ -139,8 +139,10 @@ class ArrowPythonWithNamedArgumentRunner(
override protected def evalConf: Map[String, String] = {
// An element-wise UDF receives each argument as an `array<T>` column and
flattens it, so like
- // the Arrow batched UDF it needs the input schema to convert the incoming
Arrow types.
- if (evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF ||
+ // the Arrow batched UDF it needs the input schema to convert the incoming
Arrow types. This
+ // covers the row-at-a-time lift (SQL_ARROW_ELEMENTWISE_UDF) as well as
the vectorized lifts
+ // (scalar pandas / Arrow, and their iterator variants), which flatten the
same way.
+ if (PythonEvalType.isElementwiseUDF(evalType) ||
evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
super.evalConf ++ Map(
"input_type" -> schema.json
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
index 637ec259c651..cab6c9895ca9 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
@@ -17,7 +17,6 @@
package org.apache.spark.sql.execution.python
-import org.apache.spark.api.python.PythonEvalType
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
@@ -40,10 +39,15 @@ import org.apache.spark.sql.types.{ArrayType, IntegerType,
MapType}
* transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s
-> s.u0)
* }}}
*
- * `plus_one_over_array` is the same function re-typed as `array<T> =>
array<R>` and run with
- * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour
lives in the Python
- * worker: it flattens each list column once, calls the function over all
elements of the batch, and
- * re-nests by the input's offsets - one row in, one row out, one Python round
trip per batch.
+ * `plus_one_over_array` is the same function re-typed as `array<T> =>
array<R>` and run with an
+ * element-wise eval type chosen from the UDF's own flavor (see
+ * [[PythonUDF.liftedElementwiseEvalType]]): the row-at-a-time UDFs share the
pickle-based
+ * [[org.apache.spark.api.python.PythonEvalType]]'s SQL_ARROW_ELEMENTWISE_UDF,
while a scalar pandas
+ * / Arrow UDF (and its iterator variant) lifts to its own element-wise type
so the worker keeps
+ * that flavor's batching contract.
+ * The array-at-a-time behaviour lives in the Python worker: it flattens each
list column once,
+ * calls the function over all elements of the batch, and re-nests by the
input's offsets - one row
+ * in, one row out, one Python round trip per batch.
*
* Every lifted argument is a single-level `array<T>` aligned with the
iterated array (an
* element-independent value is repeated into one with a native `transform`),
so the worker flattens
@@ -61,7 +65,6 @@ import org.apache.spark.sql.types.{ArrayType, IntegerType,
MapType}
* transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
* - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees
earlier steps'
* outputs, not array elements.
- * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
*/
object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
@@ -397,7 +400,10 @@ object ExtractPythonUDFFromLambda extends
Rule[LogicalPlan] {
// containsNull = true.
ArrayType(udf.dataType, containsNull = true),
arrayArgs,
- PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ // Each rewritable flavor lifts to its own element-wise eval type so
the worker keeps that
+ // flavor's batching contract (pickle row-at-a-time, pandas Series,
Arrow Array, or an
+ // iterator of batches). See `PythonUDF.liftedElementwiseEvalType`.
+ PythonUDF.liftedElementwiseEvalType(udf.evalType),
udf.udfDeterministic)
val signature: Expression = if (udf.udfDeterministic)
lifted.canonicalized else lifted
val ordinal = ordinalBySignature.getOrElseUpdate(signature, {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
index 7a9ec702fea8..c290ec04deda 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
@@ -241,7 +241,12 @@ object ExtractPythonUDFs extends Rule[LogicalPlan] with
Logging {
def canChainWithParallelUDFs(evalType: Int): Boolean = {
if (evalType == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF ||
- evalType == PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF) {
+ evalType == PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF ||
+ // The lifted iterator flavors carry the same one-UDF-per-operator
constraint as their base
+ // iterator types: the worker feeds a single iterator UDF, so they
must not be
+ // parallel-fused with siblings (e.g. `transform(arr, x -> f(x) +
g(x))` for iterator UDFs).
+ evalType == PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF ||
+ evalType == PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF) {
false
} else {
evalType == firstVisitedScalarUDFEvalType.get
@@ -350,6 +355,10 @@ object ExtractPythonUDFs extends Rule[LogicalPlan] with
Logging {
| PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
| PythonEvalType.SQL_ARROW_BATCHED_UDF
| PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF
+ | PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF
+ | PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF
+ | PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF
+ | PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF
| PythonEvalType.SQL_SCALAR_ARROW_UDF
| PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF =>
ArrowEvalPython(validUdfs, resultAttrs, child, evalType)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
index d0228f9341c1..34d63df710cb 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
@@ -368,6 +368,27 @@ class MyDummyScalarPandasUDF extends
UserDefinedPythonFunction(
pythonEvalType = PythonEvalType.SQL_SCALAR_PANDAS_UDF,
udfDeterministic = true)
+class MyDummyScalarArrowUDF extends UserDefinedPythonFunction(
+ name = "dummyScalarArrowUDF",
+ func = new DummyUDF,
+ dataType = BooleanType,
+ pythonEvalType = PythonEvalType.SQL_SCALAR_ARROW_UDF,
+ udfDeterministic = true)
+
+class MyDummyScalarPandasIterUDF extends UserDefinedPythonFunction(
+ name = "dummyScalarPandasIterUDF",
+ func = new DummyUDF,
+ dataType = BooleanType,
+ pythonEvalType = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
+ udfDeterministic = true)
+
+class MyDummyScalarArrowIterUDF extends UserDefinedPythonFunction(
+ name = "dummyScalarArrowIterUDF",
+ func = new DummyUDF,
+ dataType = BooleanType,
+ pythonEvalType = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
+ udfDeterministic = true)
+
class MyDummyUDTPythonUDF extends UserDefinedPythonFunction(
name = "dummyUDTUDF",
func = new DummyUDF,
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
index b764b5c9cf9d..6ec31e67ca69 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
@@ -41,6 +41,9 @@ class ExtractPythonUDFFromLambdaSuite extends QueryTest with
SharedSparkSession
private val pythonUDF = new MyDummyPythonUDF
private val scalarPandasUDF = new MyDummyScalarPandasUDF
+ private val scalarArrowUDF = new MyDummyScalarArrowUDF
+ private val scalarPandasIterUDF = new MyDummyScalarPandasIterUDF
+ private val scalarArrowIterUDF = new MyDummyScalarArrowIterUDF
// Used where a UDF call must receive two arguments, e.g. a pairwise
comparator.
private val pythonUDF2 = new MyDummyPythonUDF
private val nondeterministicUDF = new MyDummyNondeterministicPythonUDF
@@ -262,11 +265,24 @@ class ExtractPythonUDFFromLambdaSuite extends QueryTest
with SharedSparkSession
assert(liftedUDFs(df).size == 1)
}
- test("a pandas UDF inside a lambda still fails analysis") {
- val e = intercept[AnalysisException] {
- arrayDF.select(transform(col("values"), x =>
scalarPandasUDF(x))).collect()
+ test("a vectorized scalar UDF inside a lambda is lifted to its element-wise
eval type") {
+ // Each vectorized scalar flavor lifts to its own element-wise eval type
so the worker keeps
+ // that flavor's batching contract (pandas Series, Arrow Array, or an
iterator of batches).
+ val cases = Seq(
+ scalarPandasUDF -> PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ scalarArrowUDF -> PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ scalarPandasIterUDF ->
PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ scalarArrowIterUDF ->
PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF)
+ cases.foreach { case (udf, expectedEvalType) =>
+ val df = arrayDF.select(transform(col("values"), x => udf(x)).as("r"))
+ val evalTypeName = PythonEvalType.toString(expectedEvalType)
+ assert(udfsInsideLambda(df).isEmpty, s"UDF still inside a lambda for
$evalTypeName")
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ assert(lifted.head.evalType == expectedEvalType)
+ assert(lifted.head.dataType == ArrayType(udf.dataType, containsNull =
true))
+ assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType])
}
- assert(e.getCondition ==
"UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
}
test("a UDF with a UDT type inside a lambda still fails analysis") {
@@ -288,6 +304,23 @@ class ExtractPythonUDFFromLambdaSuite extends QueryTest
with SharedSparkSession
PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true)
assert(PythonUDF.isElementwiseRewritableUDF(plain))
+ // Every row-at-a-time and vectorized scalar eval type is rewritable; each
maps to its own
+ // element-wise lifted eval type.
+ Seq(
+ PythonEvalType.SQL_ARROW_BATCHED_UDF ->
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_UDF ->
PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF ->
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_UDF ->
PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF ->
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF).foreach {
+ case (base, lifted) =>
+ assert(PythonUDF.isElementwiseRewritableUDF(plain.copy(evalType =
base)))
+ assert(PythonUDF.liftedElementwiseEvalType(base) == lifted)
+ }
+ assert(PythonUDF.liftedElementwiseEvalType(PythonEvalType.SQL_BATCHED_UDF)
==
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF)
+
val zeroArg = plain.copy(children = Seq.empty)
assert(!PythonUDF.isElementwiseRewritableUDF(zeroArg))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
