[SPARK-58736][SQL][PYTHON][CONNECT] Support incremental Python aggregators via Arrow

[HyukjinKwon]

What changes were proposed in this pull request?

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 udaf(...) for use in groupBy().agg(...):

from pyspark.sql.aggregator import Aggregator, 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(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).

Spark Connect: also supported. A new optional buffer_type field on the PythonUDF proto
message carries the buffer schema to the server; the Connect client (connect/udf.py,
connect/expressions.py) serializes it and udaf dispatches on is_remote(); the server
SparkConnectPlanner threads buffer_type into UserDefinedPythonFunction and builds
PythonAggregate, after which execution reuses the same operators/worker code as classic.

SQL registration: spark.udf.register("my_agg", udaf(agg)) works in both classic and Connect,
so the aggregator is usable from SQL text (SELECT my_agg(v) FROM t GROUP BY k) — the counterpart
of Scala's spark.udf.register(name, functions.udaf(agg)).

Out of scope (planned follow-ups): DISTINCT, mixing with SQL aggregate functions in
one Aggregate, window/streaming, 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.aggregator.Aggregator and udaf(...), usable in
groupBy().agg(...) and registrable via spark.udf.register(...) for use in SQL. 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), plus test_sql_registration (register via spark.udf.register, invoke from
  SQL text). A Connect parity suite (ArrowPythonAggregatorParityTests) runs the same mixin
  under ReusedConnectTestCase.
• Compilation verified for both classic and Connect: sql/compile and connect/compile build
  cleanly (including proto regeneration). Full test execution runs in this PR's CI.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 4.8)

[HyukjinKwon]

cc @zhengruifeng @cloud-fan @Yicong-Huang Seems like this way it can do the actual partial aggregation.

[Yicong-Huang]

why can't this reuse PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF?

[zhengruifeng]

I think it is a different iterator? it is for element-iterator inside a row, not a row-iterator

[HyukjinKwon]

yeah my take is that PythonEvalType.* decides internal computation type

[HyukjinKwon]

GROUPED-AGG pandas UDF            single stage · groupByKey
───────────────────────────────────────────────────────────
     P1[a b a]     P2[b a b]     P3[a a b]     3 partitions
         └──────────────┼──────────────┘
                        ▼
     ═══ SHUFFLE ═══ all 9 raw rows move ═══
                ┌───────┴───────┐
                ▼               ▼
       ┌────────────────┐  ┌────────────────┐
       │ key a          │  │ key b          │
       │ a a a a a      │  │ b b b b        │  ← whole group,
       │ → udf(Series)  │  │ → udf(Series)  │    one worker
       └────────────────┘  └────────────────┘
                a → r               b → r


INCREMENTAL Aggregator (udaf)     two stages
───────────────────────────────────────────────────────────
     P1[a b a]     P2[b a b]     P3[a a b]
       │ reduce       │ reduce       │ reduce  ┐
       ▼              ▼              ▼
    [Σa Σb]        [Σa Σb]        [Σa Σb]      ┘ (map-side combine)
         └──────────────┼──────────────┘
                        ▼
     ═══ SHUFFLE ═══ only 6 buffers move ═══
                ┌───────┴───────┐
                ▼               ▼
       ┌────────────────┐  ┌────────────────┐
       │ key a          │  │ key b
       │ Σa Σa Σa       │  │ Σb Σb Σb       │  ← only a few
       │ → merge → fin  │  │ → merge → fin
       └────────────────┘  └────────────────┘
                a → r               b → r

[cloud-fan]

2 blocking, 0 non-blocking, 0 nits.
The two-stage integration is coherent, but the worker lifecycle still has two blocking semantic gaps in incremental memory use and empty-input aggregation.

Design / architecture (1)

• Blocking: python/pyspark/worker.py:2234: Stream partial input batches into the aggregation buffers instead of retaining and concatenating the complete group first. -- see inline

Correctness (1)

• Blocking: sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala:126: Emit the identity buffer for empty global input so finish(zero) produces the required single aggregate row. -- see inline

Verification

Traced udaf through classic and Connect UDF construction, Catalyst planning, both physical stages, and the Python worker handlers. Confirmed that the partial handler calls list(group) before reduce, and that the physical operator returns an empty iterator before the identity buffer can be emitted for empty global input. Tests were not run as part of this review.

[cloud-fan]

Blocking:

The partial stage still materializes every Arrow batch for the group before reduce runs, so a skewed group retains the same whole-group peak memory that this API is meant to avoid. Please fold each incoming batch directly into the per-aggregator buffers and retain only those buffers.

[HyukjinKwon]

Fixed in a12f825. The PARTIAL (and FINAL) worker handlers now stream Arrow batches and fold each one directly into the per-aggregator buffers — no more list(group)/concat_batches. Map-side peak memory is now bounded by a single batch plus the buffers, not the whole group.

[cloud-fan]

Blocking:

This shortcut drops empty global aggregations before Python can apply zero and finish, so df.limit(0).agg(udaf(...)) returns no row instead of the aggregate's identity result. Please emit an identity partial buffer for the no-group empty-input case and cover it with a focused test.

[HyukjinKwon]

Fixed in a12f825. A global aggregation over empty input now returns the identity row finish(zero). Since GroupedPythonArrowInput can't transmit an empty group, the FINAL stage (single AllTuples partition) injects one all-null buffer row; the worker skips null partial buffers, so it merges nothing and emits finish(zero). Added a focused test: df.limit(0).agg(udaf(...)).

[HyukjinKwon]

Thanks for taking a look — happy to hold. Answering the three points:

1. Physical operator — it already is one implementation. Both stages run through a single
PythonIncrementalAggregateExecBase.doExecute; PythonIncrementalAggregate{Partial,Final}Exec are
thin parameterizations of it (which columns are sent to Python, the eval type, the required child
distribution, the output attributes). That mirrors the normal path, where the planner emits two
HashAggregateExec nodes (partial + final) around an Exchange, both backed by one operator
class. If you'd prefer the shape where it's literally a single case class parameterized by an
aggregate mode (rather than a shared base + two thin subclasses), that's a mechanical change since
all the logic already lives in the base — glad to do it.

2. What "blocks" collapsing the two eval types. They encode two genuinely different worker
computations with different I/O, not just a mode bit:

stage	worker input	worker computation	worker output
PARTIAL	raw argument columns	zero then fold via reduce	one buffer struct
FINAL	one buffer struct column	fold via merge, then finish	the result value

On the JVM the mode lives on AggregateExpression.mode and the operator just reads it — there's no
boundary. But the worker sits across the serialization boundary, and PythonEvalType is the
discriminator Spark uses there (as came up in the earlier thread — "PythonEvalType decides internal
computation type"). Collapsing to one eval type doesn't remove the branch (the worker still has to
know partial vs final); it just moves the discriminator into a new per-call mode field in the
Python UDF wire protocol, which is more surface than two eval-type constants. So it's a
protocol-surface tradeoff, not a fundamental blocker. If you'd rather have one eval type + a mode
field, I'm happy to switch — just wanted to flag the cost.

3. Multiple UDFs and mixing.

• Multiple udafs in one aggregate: supported. Each aggregator carries its own bufferSchema;
  PARTIAL emits one buffer struct column per aggregator (_0.._n), FINAL reads each aggregator's own
  buffer column positionally and runs merge+finish independently. Added a test with two
  aggregators (different buffer schemas) over the same input.
• Mixing udaf with grouped-agg (iter) pandas/arrow UDFs, or with SQL aggregates: not supported
  in this PR — it hits the same INVALID_PANDAS_UDF_PLACEMENT analysis error that already forbids
  mixing a grouped-agg pandas UDF with other aggregate functions (the planner requires forall one
  kind). I extended that error so it also names the incremental aggregators. True cross-kind mixing
  (incremental + whole-group in one Aggregate) would need a combined physical operator and is
  called out as a follow-up in the description.

Latest push (3917089) has the multi-UDF test and the error-message fix. Let me know which shape you
prefer for (1)/(2) and I'll adjust.