[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(...):

frompyspark.sql.aggregatorimportAggregator, udaffrompyspark.sql.typesimportStructType, StructField, DoubleType, LongTypeclassMean(Aggregator):
@propertydefbufferSchema(self):
returnStructType([StructField("sum", DoubleType()), StructField("count", LongType())])
@propertydefoutputType(self):
returnDoubleType()
defzero(self): return (0.0, 0)
defreduce(self, buf, v): return (buf[0] +v[0], buf[1] +1)
defmerge(self, a, b): return (a[0] +b[0], a[1] +b[1])
deffinish(self, buf): returnbuf[0] /buf[1] ifbuf[1] elseNonedf.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.

[HyukjinKwon]

Code review (self, head 1d6c9391)

Verdict: approve with minor cleanups. No correctness or design blockers. The two-stage design is coherent, both earlier blocking comments (stream batches into buffers; emit finish(zero) for empty global input) are genuinely fixed, and Connect parity + SQL registration are wired through. Findings below are all minor.

Findings

1. FINAL worker handler builds output arrays without an explicit Arrow type — robustness

In python/pyspark/worker.py, the FINAL handler does:

result_arrays= [pa.array([r]) forrinresults]

whereas the PARTIAL handler correctly passes type=return_schema.field(i).type. Relying purely on enforce_schema to coerce an inferred type is fragile for non-trivial outputTypes (decimal, timestamp, nested struct) or an all-None column. Suggest aligning FINAL with PARTIAL:

result_arrays= [pa.array([r], type=return_schema.field(i).type) fori, rinenumerate(results)]

2. Loop-invariant field_names recomputation in both worker handlers — minor perf

field_names = [f.name for f in agg.bufferSchema.fields] is recomputed per (group × batch × aggregator) in the FINAL handler and per (group × aggregator) in the PARTIAL handler, though it depends only on the fixed i-th aggregator. Precompute a field_names_by_udf list once where grouped_func is defined and index by i.

3. Unqualified Scaladoc link — doc nit

In sql/catalyst/.../expressions/PythonUDF.scala, [[PythonIncrementalAggregateExec]] cannot resolve from sql/catalyst (the class lives in sql/core). Fully-qualify it as [[org.apache.spark.sql.execution.python.PythonIncrementalAggregateExec]] (as the reverse-direction reference already does) or use plain text.

4. aggregator.py docstring polish — optional

• The reduce example (buffer[0] + v, buffer[1] + 1) raises on a null input value; consider showing null handling so the example isn't copied as a fragile pattern.
• udaf's Raises: lists only PySparkImportError, but it also raises PySparkTypeError for a non-Aggregator arg or a non-StructTypebufferSchema.

Test portfolio

Strong coverage: builtin mean/sum equivalence, no-group, empty-global input, custom buffer, multiple aggregators, partition-count independence, SQL registration, and a Connect parity mirror. Two gaps worth a follow-up: no test with a non-trivial output type (decimal/timestamp) — which would exercise finding #1 — and no test of reduce receiving a null input value.

[dongjoon-hyun]

Thank you for working on this, @HyukjinKwon. The two-stage partial aggregation design looks promising. I reviewed the change and found several issues — the common pattern is that the new PythonAggregate expression bypasses every Catalyst guard that checks isInstanceOf[PythonUDAF], so paths that cleanly reject pandas UDAFs now either return wrong results or fail with internal errors.

Correctness (silent wrong results)

1. DISTINCT / FILTER are silently dropped.checkUnsupportedAggregateClause (FunctionResolution.scala#L534) guards only PythonUDAF, and PythonIncrementalAggregateExec.plan never reads AggregateExpression.isDistinct/.filter. After spark.udf.register("my_mean", udaf(Mean())), SELECT my_mean(DISTINCT v) or my_mean(v) FILTER (WHERE v > 0) runs and silently returns the non-distinct/unfiltered result.

2. Pivot produces wrong results.ResolvePivot's checkValidAggregateExpression (Analyzer.scala#L961) rejects PythonUDAF but lets PythonAggregate into the fallback If(pivotCol == value, v, null) rewrite, which relies on the aggregate ignoring nulls. An aggregator whose reduce does not skip None counts every input row for every pivot column.

3. Named arguments are lost in the worker. The builder admits named arguments for the FINAL eval type and the exec forwards ArgumentMetadata keys, but the PARTIAL/FINAL handlers in worker.py read only args_offsets and ignore kwargs_offsets. udaf(Mean())(v=df.v) (or SQL my_agg(v => x)) hands reduce an empty tuple — the doc example crashes on (v,) = value, and arity-tolerant aggregators silently mis-aggregate.

4. Floating-point grouping keys are not normalized. The SQL aggregate branch normalizes grouping keys via NormalizeFloatingNumbers during planning, but the new branch (SparkStrategies.scala#L806) passes them raw, so 0.0/-0.0 (and NaN bit patterns) split one logical group into two output rows. (The existing PythonUDAF branch shares this gap, but this PR adds a second operator replicating it.)

5. Duplicate buffer field names corrupt then crash.udaf() validates only isinstance(bufferSchema, StructType). With duplicate field names (legal in StructType), the PARTIAL stage's name-keyed dict silently collapses fields (pyarrow fills both struct children without error), then the FINAL stage's to_pylist() fails post-shuffle with an opaque ValueError: ... duplicate field names .... Validating at udaf() creation would be much friendlier.

Missing guards (internal errors instead of clear messages)

6. Window:isWindowPandasUDF matches only PythonUDAF, so udaf(...).over(window) is classified WindowFunctionType.SQL, planned into WindowExec, and dies in AggregateProcessor with SparkException.internalError("Unsupported aggregate function ...").

7. Streaming: the guard at SparkStrategies.scala#L583 checks only PythonUDAF, so a streaming groupBy().agg(udaf(...)) reaches AggUtils.planStreamingAggregation and crashes on UnevaluableAggregateFunc.aggBufferAttributes.

8. Mixed aggregates get a misleading error.df.groupBy(k).agg(udaf_mean(v), count("*")) — likely the most common first thing users try — falls through to INVALID_PANDAS_UDF_PLACEMENT, whose message blames "group aggregate pandas UDF" and "non-pandas aggregate functions"; both wrong for an Arrow-based Aggregator. A dedicated error (or support via AggUtils planning) would help.

Performance / design

9. The map-side stage forces a full pre-shuffle sort.requiredChildOrdering (PythonIncrementalAggregateExec.scala#L84) inserts a per-partition SortExec over full-width input rows before the PARTIAL stage — a cost HashAggregateExec's partial mode avoids and ArrowAggregatePythonExec pays only after the shuffle. Combined with GroupedPythonArrowInput opening one Arrow IPC stream per group (and one single-row batch back per group), high-cardinality keys degenerate to an O(n log n) sort plus per-row protocol overhead, making the partial stage worse than shuffling raw rows. A hash-based map-side buffer and packing many groups per batch would avoid this.

10. bufferSchema is threaded as a monkey-patched attribute. It is set on both the UDF object and its wrapper, re-attached in register, and read back via getattr(..., None). Any reconstruction path drops it silently — e.g. wrapper.asNondeterministic() already loses it today, and registering that wrapper fails later at plan time with a raw IllegalArgumentException from require(bufferType != null) (no error class). Making it a first-class UserDefinedFunction field would make this impossible.

Items 1–3 seem blocking since they silently return wrong answers for documented usage; 4–8 need either support or clean unsupported-errors before this ships.

[dongjoon-hyun]

Could you check the above comment?

[HyukjinKwon]

Thanks a lot for the thorough review, @dongjoon-hyun -- these were all real issues. Addressed in 5daa627 (pushed). Point by point:

Correctness

1. DISTINCT / FILTER -- checkUnsupportedAggregateClause now also runs for PythonAggregate in FunctionResolution, so both are rejected at analysis instead of being silently dropped.
2. Pivot -- ResolvePivot.checkValidAggregateExpression now rejects PythonAggregate as well, so it no longer falls into the null-ignoring If(...) rewrite.
3. Named arguments -- rather than reject them, I threaded them through: the PARTIAL worker now appends the kwargs_offsets columns (in call order) after the positional ones into the reduce value tuple. Verified with both agg(udaf(...)(v=col)) and SQL my_agg(v => x).
4. Floating-point grouping keys -- the PythonAggregate planning branch now normalizes grouping keys via NormalizeFloatingNumbers, mirroring the SQL branch, so 0.0/-0.0 and NaN patterns collapse into one group.

Guards (clean errors instead of internal failures)

6. Window -- WindowResolution now rejects PythonAggregate in a window with UNSUPPORTED_EXPR_FOR_WINDOW instead of misclassifying it as a SQL aggregate.
7. Streaming -- the streaming-aggregation guard in SparkStrategies now also matches PythonAggregate.
8. Mixed aggregates -- added a dedicated INVALID_PYTHON_UDF_PLACEMENT error so the message no longer wrongly blames "group aggregate pandas UDF".

Validation / design

5. Duplicate buffer field names -- validated at udaf() creation now, raising DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT up front.
6. bufferSchema -- promoted to a first-class UserDefinedFunction field (classic + Connect), so it survives _wrapped(), asNondeterministic(), and spark.udf.register without being a monkey-patched attribute.
7. Map-side sort -- the PARTIAL stage is now a hash-based combine: no requiredChildOrdering/pre-shuffle SortExec, and many groups per Arrow batch instead of one IPC stream per group. The worker keeps one running buffer per key and re-emits the keys; the FINAL stage re-groups authoritatively after the shuffle (kept sort-based), so the map-side combine only needs to be best-effort -- any keys it can't collapse (e.g. NaN) are merged downstream.

All the arrow aggregator tests (classic + Connect parity) pass locally, and I added coverage for each of the above (DISTINCT/FILTER, pivot, window, mixed placement, named args, duplicate buffer names, float/NaN and struct grouping keys). One gap: I didn't add an automated streaming-rejection test (flaky in local runs) -- the guard is symmetric with the existing PythonUDAF one. Happy to add it if you'd prefer.

PTAL, thanks again!

[dongjoon-hyun]

Review findings

Overall this looks well-designed and well-tested. I found a few issues worth addressing before merge — two functional, two error-handling, and some duplication/efficiency notes.

Functional

1. Profiler breaks incremental aggregators — python/pyspark/sql/udf.py (~L605): the profiler exclusion list in UserDefinedFunction.__call__ was not extended with SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF. With spark.python.profile=true, self.func (the Aggregator object) is replaced by a plain wrapper function before pickling, so the worker fails with AttributeError: 'function' object has no attribute 'bufferSchema'. With spark.python.profile.memory=true it fails even earlier on the driver at inspect.getsourcelines(f.__code__), since an Aggregator instance has no __code__.

2. NormalizePlan.normalizeExprIds not extended — sql/catalyst/.../plans/NormalizePlan.scala (~L112): resultId is normalized for PythonUDF/PythonUDAF but not for the new PythonAggregate. Any comparePlans-based test on plans containing a PythonAggregate will fail spuriously, and HybridAnalyzer's dual-run comparison would report a mismatch for semantically identical plans.

Error handling

3. Error message asserts a restriction the implementation doesn't have — error-conditions.json (INVALID_PYTHON_UDF_PLACEMENT): "Each such function must be the only aggregate expression in its aggregation" contradicts the supported (and tested, test_multiple_incremental_aggregators) case of multiple incremental aggregators in one aggregation. Only mixing with other kinds of aggregate functions is invalid.

4. Bare require instead of a classed error — UserDefinedPythonFunction.builder (~L113): require(bufferType != null, ...) + unchecked asInstanceOf[StructType]. A Connect proto with eval_type=256 and a missing (or non-struct) buffer_type surfaces as a raw IllegalArgumentException/ClassCastException (INTERNAL_ERROR) rather than following the planner's InvalidPlanInput convention. Also reachable classically via UserDefinedFunction(f, evalType=256) with no bufferSchema.

5. Mixed pandas + incremental diagnostic drops names — SparkStrategies.scala (~L825): when both a grouped-agg pandas UDAF and an incremental aggregator are mixed with other aggregates, the new fallthrough names only the PythonAggregate functions; the co-offending pandas UDAF names (which the replaced code listed) are omitted. This mixed case has no test.

6. (Minor, pre-existing pattern)Connect is_distinct silently ignored — SparkConnectPlanner.scala (~L2204): agg.toAggregateExpression() hardcodes isDistinct = false, so a client sending is_distinct=true silently gets non-distinct results — while this PR's own FunctionResolution change explicitly rejects DISTINCT on the SQL path ("reject rather than silently drop the clause"). Same hole exists for PythonUDAF, but this PR is where the reject-don't-drop rule was introduced.

Efficiency / duplication

7. Unbounded map-side combine — worker.py PARTIAL handler (~L2294): all per-group buffers are held in one dict for the whole partition and emitted as a singleRecordBatch (no maxRecordsPerBatch chunking, no cap/flush). High-cardinality keys — exactly where partial aggregation degenerates — will OOM the Python worker where JVM hash aggregation would spill. Since the FINAL stage re-merges duplicate keys authoritatively, capping the map and flushing early (and emitting the end-of-partition output in bounded chunks) is safe and cheap. I saw spill is listed as a follow-up, but a simple size cap + chunked emission seems worth doing in this PR.

8. doExecute duplication has already drifted — PythonIncrementalAggregateExecBase.doExecute is a near-verbatim copy of ArrowAggregatePythonExec.doExecute, and the copy hard-codes lockFree = false (~L162) where the original picks it from PYTHON_UDF_PIPELINED_EXECUTION — so the FINAL stage silently misses the pipelined-queue optimization. Extracting the shared grouped-input/queue/join scaffolding (the argMetas dedup loop alone now exists twice within this PR) would prevent further drift.

9. Redundant Arrow→Python conversions — worker.py (~L2299): the same column is to_pylist()-ed once per referencing UDF plus again for grouping keys (N aggregators sharing a column ⇒ N+1 conversions per batch). One {offset: pylist} memo per batch would do.

10. Duplicated Py4J constructor call — udf.py_create_judf (~L535): the 8-line constructor invocation is duplicated in an if/else differing only in the trailing bufferType, with a JVM auxiliary constructor existing solely for the arity. Py4J maps None → null and the Scala param already defaults to null, so a single call passing None would remove both the else branch and the auxiliary constructor.

[cloud-fan]

2 addressed, 0 remaining, 7 new to this AI review. (3 newly introduced, 4 late catches, 0 previously raised.)
4 blocking, 3 non-blocking, 0 nits.
The earlier whole-group-memory and empty-global-input blockers are fixed, but several newly identified runtime and contract gaps should be addressed before merge.

Correctness (6)

• Blocking: General: Exclude incremental aggregators from ordinary Python UDF profiling; the profiler rewrites or introspects the Aggregator object and breaks execution.
• Blocking: General: Normalize PythonAggregate.resultId alongside PythonUDF and PythonUDAF to prevent false plan mismatches.
• Non-blocking: General: Make INVALID_PYTHON_UDF_PLACEMENT prohibit mixing aggregate kinds, not multiple incremental aggregators, which are supported and tested.
• Non-blocking: General: Validate missing and non-struct buffer types at the planner boundary and return a classed input error instead of require/asInstanceOf failures.
• Blocking: General: Reject or honor PythonUDF.is_distinct in the Connect planner instead of silently constructing PythonAggregate with isDistinct=false.
• Blocking: General: Cap and flush the PARTIAL group map in bounded batches; FINAL can safely re-merge duplicate keys emitted by early flushes.

Suggestions (1)

• Non-blocking: General: Cache to_pylist() once per distinct input offset so aggregators sharing a column do not repeatedly materialize it.

Verification

Traced the public API through classic and Connect UDF construction, Catalyst resolution and normalization, the two-stage physical plan, and both Python worker handlers. Verified the two earlier blockers are fixed in the current tree, then checked each new finding against its reachable dispatch or configuration path. Tests were not run as part of this review.

PR metadata suggestions

• Document: map-side buffers are currently retained for every distinct key in a partition and do not spill or flush early.

[HyukjinKwon]

Thanks again @dongjoon-hyun, and @cloud-fan for the AI pass -- these overlapped, so this covers both. All addressed in 56fd4be (pushed). Classic (24) + Connect parity (22) tests pass locally.

Correctness / functional

1. Profiler breaks incremental aggregators -- UserDefinedFunction.__call__ now short-circuits when the eval type is the incremental aggregator: it warns ("Profiling incremental Python aggregators is not supported.") and takes the non-profiled path, so neither the CPU-profiler wrapper (which would drop the Aggregator's zero/reduce/bufferSchema) nor the memory profiler's inspect.getsourcelines(f.__code__) runs. Added CPU + memory profiler fallback tests.
2. NormalizePlan.normalizeExprIds -- now resets PythonAggregate.resultId alongside PythonUDF/PythonUDAF.
3. INVALID_PYTHON_UDF_PLACEMENT message -- reworded; multiple incremental aggregators together are supported (and tested), only mixing with other kinds of aggregate is invalid.
4. Missing/non-struct buffer schema -- builder now returns a classed INVALID_PYTHON_AGGREGATOR_BUFFER_SCHEMA (via QueryCompilationErrors) instead of a bare require/ClassCastException, covering the malformed-Connect-proto and direct-UserDefinedFunction paths. Added a test.
5. Mixed diagnostic drops pandas names -- the fallthrough error now names both the incremental aggregators and any co-offending grouped-agg pandas/arrow UDAFs. Added a mixed pandas+incremental test asserting both names appear.
6. Connect is_distinct -- the Connect planner now rejects is_distinct for PythonAggregate with the same DISTINCT-unsupported error the SQL path raises (a resolved Connect aggregate bypasses FunctionResolution, so it was silently dropped before).

Efficiency / duplication

7. Unbounded map-side combine -- the PARTIAL worker now caps the per-key buffer at maxRecordsPerBatch, flushes early when the cap is reached, and emits end-of-partition buffers in bounded chunks. This is safe because FINAL re-merges any duplicate keys the early flushes produce. Added a test with maxRecordsPerBatch=2 to exercise flushing.
8. doExecute duplication + lockFree -- the FINAL stage now reads PYTHON_UDF_PIPELINED_EXECUTION for the queue's lockFree flag (no longer hard-coded false), and the argMetas dedup is extracted into a shared buildArgMetas used by both stages.
9. Redundant Arrow->Python conversions -- the PARTIAL worker now converts each needed input column at most once per batch (memoized by offset), even when several aggregators or a grouping key share it.
10. Duplicated Py4J constructor -- collapsed into a single call passing None (Py4J maps it to the null the Scala bufferType already defaults to), and removed the now-unused auxiliary JVM constructor.

Two notes:

• SPARK-1121 Only add avro if the build is for Hadoop 0.23.X and SPARK_YARN is set #6: the programmatic is_distinct=true proto path isn't reachable from the PySpark client API, so there's no dedicated client test for it; the user-facing SQL DISTINCT rejection is covered for both classic and Connect by test_distinct_and_filter_rejected.
• Removed reference to incubation in Spark user docs. #2: symmetric one-liner with PythonUDAF; exercised via the plan-comparison path rather than a new standalone test.

PTAL, thanks!

[cloud-fan]

7 addressed, 0 remaining, 1 new to this AI review. (0 newly introduced, 1 late catch, 0 previously raised.)
0 blocking, 1 non-blocking, 0 nits.
The execution and integration paths are coherent after the latest fixes, but the new public implementor contract should state its required zero identity law before release.

Design / architecture (1)

• Non-blocking: python/pyspark/sql/aggregator.py:102: Specify that zero() is the two-sided identity for merge; associativity and commutativity alone do not make partial aggregation partition-independent. -- see inline

Verification

Traced udaf through classic and Connect UDF construction, Catalyst analysis and strategy routing, both physical stages, and the PARTIAL/FINAL worker handlers. Rechecked all prior AI findings against the current tree and verified their fixes. Compared the new Python contract with Scala Aggregator, whose zero documentation explicitly requires b + zero = b.

PR metadata suggestions

• Document: zero() must be the identity for merge (on both sides), because Spark may initialize multiple partial or flushed buffers for one logical group.

[HyukjinKwon]

Merge Summary:

• merged into master 9c5094a
• merged into branch-4.x b18b7d3

Posted by merge_spark_pr.py