HyukjinKwon opened a new pull request, #58148:
URL: https://github.com/apache/spark/pull/58148
### What changes were proposed in this pull request?
This PR lets a scalar Python UDF be used as the `merge` (and optional
`finish`) lambda of the SQL
`aggregate` / `reduce` higher-order function, by running the fold **inside
the Python worker**.
Element-wise HOF UDFs (`transform`, `filter`, ...) are supported today by
*lifting*: the UDF is
applied once to the whole array outside the lambda
(`ExtractPythonUDFFromLambda`). That does not
work for `aggregate` / `reduce`: the fold is a strict sequential left fold,
so the `merge` UDF reads
the running accumulator (the output of the previous step), not array
elements, and cannot be applied
once to the whole array. Such UDFs are therefore rejected today with
`UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF`.
Instead of lifting, this PR runs the fold in the worker. The optimizer
rewrites
```
aggregate(arr, zero, (acc, x) -> merge(acc, x), (acc) -> finish(acc))
```
into
```
finish( FOLD(arr, zero) )
```
where `FOLD` is a single `PythonUDF` carrying the `merge` function under a
new eval type
`SQL_SCALAR_FOLD_UDF`, with the array and zero value as its children. The
worker folds each row's
array sequentially (`acc = zero`; `acc = merge(acc, element)` per element),
and `finish` is left as
an ordinary outer scalar UDF over the fold result. This reuses the existing
pickle-based
`BatchEvalPython` operator and worker transport unchanged (the operator is
only parameterized with
the eval type) -- no new physical operator, planner strategy, or serializer.
A null array folds to
null, matching native `ArrayAggregate`.
Main changes:
- **New eval type** `SQL_SCALAR_FOLD_UDF` (`PythonRunner.scala`, `util.py`,
`_typing.pyi`).
- **Analysis** (`PythonUDF.isFoldableAggregate`, `CheckAnalysis`): accept a
foldable `aggregate` /
`reduce` under
`spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled`.
- **Rewrite** (`ExtractPythonUDFFromLambda.rewriteFoldableAggregate`) and
extraction routing
(`ExtractPythonUDFs`, `BatchEvalPython`/`BatchEvalPythonExec` gain an
`evalType`).
- **Worker** (`worker.py`): a fold branch that reuses the batched
row-at-a-time mapper.
This is an initial (WIP) version. It recognizes only the natural shape where
the `merge` lambda is a
single UDF call `(acc, x) -> udf(acc, x)` (row-at-a-time `SQL_BATCHED_UDF` /
`SQL_ARROW_BATCHED_UDF`)
and `finish` is the identity or a single UDF `(acc) -> udf(acc)`. Shapes
that mix a UDF with other
expressions in a lambda, use only a UDF in `finish` with a native `merge`,
or use a vectorized
pandas/Arrow merge remain rejected as before, and are candidate follow-ups.
### Why are the changes needed?
`aggregate` / `reduce` is the natural way to express a custom sequential
fold over an array, and
Python users reach for a Python UDF to write the accumulation logic. This
was the one common
higher-order function still unable to host a Python UDF, so users had to
fall back to
`explode` + group-by aggregation or a `mapInPandas`/UDF over the whole
array. This closes that gap
using the same config and user-facing surface as the existing element-wise
HOF UDF support.
### Does this PR introduce _any_ user-facing change?
Yes. With `spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled`
(default `true`), a Python
UDF may now be used as the `merge`/`finish` lambda of `aggregate` /
`reduce`. Previously this failed
analysis with `UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF`; now it
computes the fold. For
example:
```python
from pyspark.sql import functions as sf
from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType
add = udf(lambda acc, x: (acc or 0) + x, IntegerType())
df.select(sf.aggregate("values", sf.lit(0), lambda acc, x: add(acc, x)))
```
No change when the config is `false` (still rejected as before).
### How was this patch tested?
New unit tests in
`python/pyspark/sql/tests/test_udf_in_higher_order_function.py` (run for both
classic and Spark Connect via the parity mixin):
`test_udf_in_aggregate_fold`,
`test_udf_in_aggregate_fold_with_finish`,
`test_udf_in_aggregate_fold_is_sequential` (an
order-dependent fold, to confirm it is a genuine left fold), and
`test_udf_in_aggregate_fold_disabled_by_conf`;
`test_udf_in_aggregate_fails` was updated to reflect that the foldable shape
now succeeds while the
mixed/native-merge shapes still fail. The full suite passes locally (`Ran 63
tests ... OK`), and
each positive test asserts the result against the equivalent plain-Python
fold, including
null-array (-> null) and empty-array (-> zero) cases.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
