dongjoon-hyun commented on PR #57952: URL: https://github.com/apache/spark/pull/57952#issuecomment-5321061853
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](https://github.com/apache/spark/blob/3a1892dd9538d77ded71403a633f77f174624da8/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/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](https://github.com/apache/spark/blob/3a1892dd9538d77ded71403a633f77f174624da8/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/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](https://github.com/apache/spark/blob/3a1892dd9538d77ded71403a633f77f174624da8/sql/core/src/main/scala/org/apache/spark/sql/execution/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](https://github.com/apache/spark/blob/3a1892dd9538d77ded71403a633f77f174624da8/sql/core/src/main/scala/org/apache/spark/sql/execution/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](https://github.com/apache/spark/blob/3a1892dd9538d77ded71403a633f77f174624da8/sql/core/src/main/scala/org/apache/spark/sql/execution/python/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. -- 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]
