dongjoon-hyun commented on PR #57859: URL: https://github.com/apache/spark/pull/57859#issuecomment-5229628748
Thank you for extending this rule, @cloud-fan. The `PartialMerge` + `Final` -> `Final` combination itself looks semantically sound to me (merge is associative), but I have a few concerns. **1. Weakening the required distribution of the existing `Partial` + `Final` path looks risky under AQE.** `combineHashAggregates` now applies `requiredChildDistributionExpressions = partialAgg.requiredChildDistributionExpressions` in both modes. For a `Partial` aggregate this is `None` (`UnspecifiedDistribution`), while the previously merged code kept `finalAgg`'s `Some(groupingKeys)`. The combined `Complete`/`Final` aggregate still produces final results in a single pass, so clustering by the grouping keys remains its correctness requirement. That requirement is the safety net in AQE: `optimizeQueryStage` validates `AQEShuffleReadRule` results via `ValidateRequirements`, which checks each node's `requiredChildDistribution`. Concretely, with `SELECT k, count(*) FROM (SELECT /*+ rebalance(k) */ * FROM t) GROUP BY k`, the rebalance shuffle satisfies the final aggregate's requirement, the two aggregates become adjacent, and the rule fires. Before this PR the combined node required `ClusteredDistribution(k)`, so `OptimizeSkewInRebalancePartitions` splitting a skewed partition would fail validation and be reverted. With this PR the combined node claims `Unspecified`, the split passes validation, the same key can land in multiple partitions, and the one-shot aggregate emits duplicate groups. The global-aggregation case (`Some(Nil)` = `AllTuples` -> `None`) has the same issue. I think the combined node should keep `finalAgg`'s `requiredChildDistributionExpressions` in both modes — removing the lower aggregate doesn't change the combined operator's own correctness requirement. **2. The `partialAgg.outputSet != finalAgg.usedInputs` guard is always false.** `HashAggregateExec.usedInputs` is `inputSet` (from `AggregateCodegenSupport`), i.e. `AttributeSet(child.output)`, and in this pattern `finalAgg.child` *is* `partialAgg`, so the two sides are identical by construction. The check never rejects anything. If the intent is "the lower aggregate's output exactly matches the final aggregate's inputs", that holds trivially; the layout that actually matters after combining is `partialAgg.child.output`, which this check doesn't look at. **3. `initialInputBufferOffset` should arguably come from `partialAgg`.** The combined `Final` aggregate reads `partialAgg.child`'s rows, so the buffer offset should be relative to that layout, i.e. `partialAgg.initialInputBufferOffset`. Today both offsets equal the grouping length by `AggUtils` construction, so they coincide, but the rule doesn't verify that. Using `partialAgg`'s offset (or guarding on equality) would be more robust. Minor comments: - The class doc still only describes `Partial` + `Final` -> `Complete`; the new pattern is worth documenting there, and the sort/object-hash branches don't get the new metadata handling — whatever we decide for (1) should be consistent across the three branches. - The new test builds the `PartialMerge` plan synthetically and only checks plan shape; it never executes the combined plan. An end-to-end test (and one covering the AQE rebalance/skew scenario in (1)) would be valuable. Could you also share a concrete query where the planner naturally produces adjacent `PartialMerge` + `Final` hash aggregates? The `PartialMerge` nodes from `AggUtils` seem to be either mode-mixed (distinct) or separated by streaming state operators. -- 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]
