[SPARK-58210][SQL][FOLLOWUP] Extend CombineAdjacentAggregation to partial merge

[cloud-fan]

What changes were proposed in this pull request?

This is a follow-up to #57363. It extends CombineAdjacentAggregation to combine adjacent
PartialMerge and Final hash aggregates into one Final aggregate. This pattern can be
produced by physical-plan extensions that add an aggregation stage; Spark's built-in query
planning does not currently produce it.

The combined aggregate preserves the upper Final aggregate's child distribution requirement.
Because the lower aggregate is removed, the combined aggregate takes its streaming and
shuffle-partition metadata from that lower aggregate and reads its original child. Filtered
Partial aggregation remains supported, while a PartialMerge pair with filters is not combined.

Why are the changes needed?

CombineAdjacentAggregation currently handles only Partial followed by Final. A compatible
adjacent PartialMerge and Final hash aggregate pair can also be collapsed safely, avoiding an
unnecessary hash aggregation stage.

The upper aggregate's distribution requirement must remain after the rule runs so that later AQE
optimization does not split partitions that the combined aggregate requires to remain clustered.
The lower aggregate's streaming and shuffle-partition metadata remain child-facing properties and
therefore move to the combined aggregate when the lower node is removed.

Does this PR introduce any user-facing change?

Yes. Spark may produce a single final hash aggregate instead of adjacent partial-merge and final
hash aggregates when their grouping and logical lineage are compatible. Query results are
unchanged.

How was this patch tested?

Added coverage to CombineAdjacentAggregationSuite for:

• combining an executable PartialMerge and Final hash aggregate pipeline and comparing the
  uncombined and combined results;
• preserving the upper distribution requirement and lower streaming, shuffle-partition,
  grouping, child, and buffer-offset metadata;
• refusing to combine a filtered PartialMerge pair;
• retaining the distribution requirement under AQE and verifying that shuffle partition specs
  are present but are not partial reducer specs.

Ran:

build/sbt 'sql/testOnly org.apache.spark.sql.execution.CombineAdjacentAggregationSuite' sql/scalastyle

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

Generated-by: Codex (GPT-5)

[dongjoon-hyun]

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.

[ulysses-you]

Combine PartialMerge - Final to one FInal is sound good to me. But is there a valid end to end test ? I can not find a real world query can reach this pattern that all agg functions are PartialMerge and the adjacent agg functions are Final.

[cloud-fan]

Thanks for raising this. OSS does not currently produce a pure adjacent PartialMerge -> Final HashAggregate pair from a built-in query plan. The rule is intended to handle plans produced by physical-plan extensions that add an aggregation stage, and I updated the class documentation to say that explicitly.

The test now constructs a valid executable Partial -> PartialMerge -> Final pipeline, applies the rule to the upper pair, and compares the uncombined and combined execution results. This replaces the previous plan-shape-only coverage. I also added a separate end-to-end query test for the AQE rebalance/skew distribution regression.

[cloud-fan]

Thanks for the detailed review. Addressed in 359d3dc:

1. The combined hash aggregate now retains finalAgg.requiredChildDistributionExpressions. This preserves ClusteredDistribution/AllTuples and lets AQE validation reject skewed rebalance reads that split a grouping key. I added an AQE rebalance/skew regression that checks the answer, confirms the aggregates combine, and verifies no PartialReducerPartitionSpec survives below the combined aggregate.

2. Removed the partialAgg.outputSet != finalAgg.usedInputs check, since it is true by construction and did not validate the post-rewrite input layout.

3. The PartialMerge -> Final result now takes initialInputBufferOffset from the lower aggregate, whose child becomes the combined aggregate child.

I also updated the class documentation. The sort/object-hash paths continue retaining the final node distribution through copy, consistently with the hash path.

For execution coverage, the PartialMerge test now builds a valid executable Partial -> PartialMerge -> Final pipeline and compares results before and after combining, in addition to its metadata and filter checks. OSS does not currently generate this pure adjacency from a built-in query; the documented use case is a physical-plan extension that introduces an extra aggregation stage.

I attempted build/sbt 'sql/testOnly org.apache.spark.sql.execution.CombineAdjacentAggregationSuite' twice, but SBT project loading was blocked before compilation by DNS failures resolving three BOMs from Maven Central (jackson-bom, jjwt-bom, and kubernetes-client-bom). Static diff/style checks pass.

[dongjoon-hyun]

Thank you, @cloud-fan. This extension looks correct to me — I verified the key correctness points:

• Buffer binding: Final-mode aggregates bind buffers positionally from initialInputBufferOffset, and the combined node takes the lower aggregate's offset while its child keeps the same grouping + buffers layout. Correct.
• AQE safety: since this rule runs after EnsureRequirements, keeping the distribution requirement on the combined node is what lets ValidateRequirements block OptimizeSkewInRebalancePartitions from splitting rebalance partitions (which would produce wrong results). The second test covers exactly this scenario.
• Streaming: real streaming aggregation plans have StateStoreSaveExec between Final and PartialMerge, so this pattern cannot match them.
• Mixed-mode (distinct) aggregates are excluded by the forall checks, and filtered PartialMerge pairs are conservatively rejected with test coverage.

A few minor comments (none blocking):

1. The PR description says "preserve the lower aggregate's child distribution", but the code keeps the upper (final) aggregate's requiredChildDistributionExpressions (and the tests assert that). The code is right; the description could be updated. The described "requires the lower aggregate output to exactly match the final aggregate inputs" check also doesn't exist explicitly — it holds implicitly from the parent-child pattern.
2. requiredChildDistributionExpressions = finalAgg.requiredChildDistributionExpressions inside finalAgg.copy(...) is a no-op. If the intent is to document that the upper node's requirement is kept, a comment would be clearer — as written it may read as a typo for partialAgg.….
3. The Complete path now takes isStreaming / numShufflePartitions from the lower aggregate (previously the final's were kept). Practically equivalent for plans Spark produces, but requiredChildDistribution is now computed from fields of two different nodes (exprs from final, the rest from partial) — a short comment on why would help future readers.
4. The IllegalArgumentException branch in combineHashAggregates is unreachable since combinedMode only returns Complete/Final; returning the new aggregateExpressions directly from combinedMode would remove the second match entirely.
5. The scaladoc reads as if PartialMerge + Final combining applies generally, but only HashAggregateExec handles it; SortAggregateExec / ObjectHashAggregateExec remain Partial-only. Worth a note in the doc.

On tests: the new path is covered by a hand-assembled plan (understandable, since Spark core never produces this shape), and the AQE test is effectively a regression test for the pre-existing Complete path — asserting the shuffle actually had skewed partitions would make it non-vacuous, but that can be a follow-up.

Thank you again for the follow-up, @cloud-fan!

[cloud-fan]

Thanks, @dongjoon-hyun. Addressed these comments in 296a850:

• updated the PR description to distinguish the upper aggregate's distribution requirement from
  the lower aggregate's child-facing streaming and shuffle metadata, and removed the claim about
  an explicit output/input equality check;
• removed the redundant requiredChildDistributionExpressions copy argument and added a comment
  explaining why the remaining metadata comes from different nodes;
• replaced the AggregateMode/exception match with an explicit combined aggregate result;
• clarified that PartialMerge + Final is supported only for HashAggregateExec;
• made the AQE partition-spec assertion non-vacuous.

The focused suite (10 tests) and sql/scalastyle both pass.

[cloud-fan]

the failed streaming test is unrelated, thanks for review, merging to master/4.x/4.3

[cloud-fan]

Merge Summary:

• merged into master f88a554
• merged into branch-4.x 99c88b3
• merged into branch-4.3 0d8a1e6

Posted by merge_spark_pr.py