HyukjinKwon opened a new pull request, #57952:
URL: https://github.com/apache/spark/pull/57952
### What changes were proposed in this pull request?
This is a **draft / WIP** that adds a Python analog of the Scala typed
`org.apache.spark.sql.expressions.Aggregator[IN, BUF, OUT]` with **true
incremental (partial)
aggregation** — i.e. map-side combine, not whole-group materialization.
Users subclass a new `Aggregator` base class (`zero` / `reduce` / `merge` /
`finish` +
`bufferSchema`) and wrap it with `arrow_udaf(...)` for use in
`groupBy().agg(...)`:
```python
from pyspark.sql.pandas.aggregator import Aggregator, arrow_udaf
from pyspark.sql.types import StructType, StructField, DoubleType, LongType
class Mean(Aggregator):
@property
def bufferSchema(self):
return StructType([StructField("sum", DoubleType()),
StructField("count", LongType())])
@property
def outputType(self):
return DoubleType()
def zero(self): return (0.0, 0)
def reduce(self, buf, v): return (buf[0] + v[0], buf[1] + 1)
def merge(self, a, b): return (a[0] + b[0], a[1] + b[1])
def finish(self, buf): return buf[0] / buf[1] if buf[1] else None
df.groupBy("k").agg(arrow_udaf(Mean())(df.v))
```
Unlike grouped-agg pandas/arrow UDFs (`PythonUDAF` +
`ArrowAggregatePythonExec`), which collect the
whole group and call Python once, this is planned as a **two-stage
aggregation**:
- a map-side **PARTIAL** stage folds each group's input rows into a
per-group buffer via `reduce`;
- the buffers are shuffled by the grouping key (as an Arrow struct column);
- a **FINAL** stage merges the partial buffers via `merge` and produces the
output via `finish`.
Because `merge` is associative/commutative, the result is independent of
partition count.
Class hierarchy / trace:
- `PythonEvalType`: new `SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF`
(255) and `..._FINAL_UDF`
(256), added on both the Python (`pyspark.util`) and JVM
(`api.python.PythonEvalType`) sides.
- Catalyst: new `PythonAggregate` expression (an `UnevaluableAggregateFunc`,
like `PythonUDAF`)
carrying the intermediate `bufferSchema`.
- Planning: `SparkStrategies.Aggregation` routes an all-`PythonAggregate`
aggregate to
`PythonIncrementalAggregateExec.plan(...)`, which builds
`PythonIncrementalAggregatePartialExec` -> (Exchange, inserted by
`EnsureRequirements`) ->
`PythonIncrementalAggregateFinalExec`. Both operators reuse
`ArrowPythonWithNamedArgumentRunner`
+ `GroupedPythonArrowInput`.
- Worker (`worker.py`): a PARTIAL handler that folds input batches into a
buffer via `reduce`, and
a FINAL handler that merges partial-buffer rows via `merge` then `finish`.
- The buffer schema is threaded to the JVM via a new nullable `bufferType` on
`UserDefinedPythonFunction` (an auxiliary constructor preserves the
existing Py4J arity).
Out of scope for this draft (planned follow-ups): `DISTINCT`, mixing with
SQL aggregate functions in
one `Aggregate`, window/streaming, Spark Connect parity, SQL
`spark.udf.register`, real disk spill
(currently the map side bounds memory by per-partition grouping;
associativity makes early partial
emission safe), and a typed-columnar vs. pickled buffer performance variant.
### Why are the changes needed?
PySpark has no incremental user-defined aggregator: every custom-aggregation
path (grouped-agg
`pandas_udf`/`arrow_udf`, `applyInPandas`) materializes the whole group and
invokes Python once,
with no map-side combine or partial/merge across the shuffle. This adds the
missing
`Aggregator`-style abstraction with genuine partial aggregation, matching
the Scala typed
`Aggregator`.
### Does this PR introduce _any_ user-facing change?
Yes — a new public API: `pyspark.sql.pandas.aggregator.Aggregator` and
`arrow_udaf(...)`, usable in
`groupBy().agg(...)`. No existing behavior changes.
### How was this patch tested?
- Compilation verified: `sql/compile` (catalyst + core + sql) builds cleanly
with the new
expression, operators, and planner routing.
- Added `python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py`
(`ArrowPythonAggregatorTests`): checks the incremental aggregator matches
built-in `avg`/`sum`,
a no-group case, a custom buffer, and that results are independent of
partition count (exercising
partial + merge).
NOTE (draft): a full end-to-end test run in the author's local environment
is still pending due to
an offline-build constraint (fresh-`master` dependency
`at.yawk.lz4:lz4-java:1.11.2` is not in the
local mirror, and the only prebuilt assembly available is from a different
branch). CI on this PR is
the intended validation.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
--
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]
