[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 for adjacent hash
aggregates in two ways:

• combine PartialMerge followed by Final into one Final aggregate;
• preserve the lower aggregate's child distribution, streaming, and shuffle-partition metadata
  when removing that lower aggregate.

The rule also requires the lower aggregate output to exactly match the final aggregate inputs.
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 pair can also be collapsed safely, avoiding an unnecessary hash
aggregation stage. When the lower aggregate is removed, its child-facing distribution and streaming
metadata must move to the combined node because the combined node now reads that lower aggregate's
original child.

This broadens the rule to optimize another safe adjacent-aggregation pattern while keeping the
combined operator's child requirements consistent with the plan subtree it now consumes.

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, lineage, and input/output attributes are compatible. Query
results are unchanged.

How was this patch tested?

Added coverage to CombineAdjacentAggregationSuite for:

• combining PartialMerge and Final hash aggregates;
• preserving distribution, streaming, shuffle, result, and buffer-offset metadata;
• refusing to combine a filtered PartialMerge pair.

Ran:

build/sbt 'sql/testOnly org.apache.spark.sql.execution.CombineAdjacentAggregationSuite'
build/sbt 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!