[SPARK-58210][SQL] Enable ReplaceHashWithSortAgg and CombineAdjacentAggregation by default

[ulysses-you]

What changes were proposed in this pull request?

This PR enables two physical aggregate rules by default, decouples their configurations, and makes them safe to enable together:

• spark.sql.execution.replaceHashWithSortAgg now defaults to true. It replaces a hash-based aggregate with a sort aggregate when the aggregate's child already satisfies the grouping-key sort order.
• spark.sql.execution.combineAdjacentAggregation now defaults to true and is no longer a fallbackConf of replaceHashWithSortAgg; it is an independent config. It merges an adjacent partial/final aggregate pair (with no shuffle between them) into a single complete-mode aggregate.
• ReplaceHashWithSortAgg switches its plan traversal from transformDown to transformUp. HashAggregateExec does not expose a grouping-key output ordering (outputOrdering = Nil), while SortAggregateExec does (it is order-preserving on the grouping keys). With transformDown, an outer (final) aggregate inspected before its inner (partial) aggregate sees the partial's Nil ordering and is not replaced; with transformUp, the partial is replaced first, so its sort-aggregate output ordering satisfies the final aggregate, which is then also replaced. This lets nested partial/final pairs both be converted when the leaf child is already sorted.
• HiveUDAFFunction is made Complete-safe (see below) so a Complete-mode ObjectHashAggregateExec (produced by combining) does not crash mode-aware Hive UDAFs.
• CentralMomentAgg.merge / Covariance.merge are fixed so merging a non-empty buffer into an empty one no longer overflows to NaN (see below).

Affected tests and golden files are updated, and migration-guide entries are added under "Upgrading from Spark SQL 4.2 to 4.3".

Why are the changes needed?

Both rules improve aggregate execution and have been available but off by default. CombineAdjacentAggregation was introduced (SPARK-43317) as a fallbackConf of replaceHashWithSortAgg, so enabling one enabled both. They are logically independent (one reuses existing ordering, the other merges an adjacent pair), so this PR turns each on by default and gives each its own config, letting users enable/disable them separately.

Enabling combining by default surfaced two issues that the old default hid, both fixed here:

• Mode-aware Hive UDAFs: a Complete-mode ObjectHashAggregateExec calls HiveUDAFFunction.update (PARTIAL1 buffer) and then eval (FINAL evaluator). UDAFs that use different buffer classes per mode threw ClassCastException. HiveUDAFFunction already converted a PARTIAL1 buffer to a FINAL buffer on demand in merge; that conversion is extracted into toFinalBuffer and also applied in eval, so the Complete path terminates against a FINAL buffer.
• Empty-buffer merge overflow: CentralMomentAgg.merge / Covariance.merge compute delta * deltaN * n1 * n2 (or dx * dyN * n1 * n2). When one side is an empty buffer (n1 == 0 or n2 == 0) and the other has a large average, delta * deltaN overflows to Infinity before multiplying by the zero count, and Infinity * 0 = NaN corrupts the moments. The old default (no combining) hit this on single-partition groups of large equal values, returning NaN for var_pop/covar_pop/regr_sxy; the Complete-mode default returned the correct 0.0. The fix forces delta/dx/dy to 0.0 when either side is empty and sets the merged average to the non-empty side, so both paths now return 0.0. This is a standalone correctness fix that also helps the separate partial/final path on multi-partition inputs.

Does this PR introduce any user-facing change?

Yes. On default config, Spark may now plan a sort aggregate instead of a hash aggregate (when the child is already sorted on the grouping keys) and may merge an adjacent partial/final aggregate pair into a single complete-mode aggregate. The physical plan of some queries changes as a result. var_pop/covar_pop/regr_sxy on single-partition groups of large equal values now return 0.0 instead of NaN. To restore the previous behavior, set spark.sql.execution.replaceHashWithSortAgg and/or spark.sql.execution.combineAdjacentAggregation to false (both must be false to fully restore the previous partial/final staging, since the settings are now independent).

How was this patch tested?

Updated existing tests:

• ReplaceHashWithSortAggSuite / CombineAdjacentAggregationSuite: the checkAggs helper now toggles both configs together; the former "falls back to replaceHashWithSortAgg" test is replaced by an independence test. The transformUp traversal is exercised by the "replace partial and final hash aggregate together with sort aggregate" case (SMJ -> partial -> final), which expects both aggregates replaced.
• SQLMetricsSuite (SPARK-25497): pins combineAdjacentAggregation=false so the single-partition query keeps its partial/final structure for the limit/codegen metric assertions.
• PlannerSuite (SPARK-40086): accepts a sort aggregate for the single-partition sorted-input queries while still asserting no extra shuffle.
• HiveUDAFSuite (SPARK-24935): runs on the default (combined) path now that the Hive wrapper is Complete-safe; asserts the single Complete-mode ObjectHashAggregateExec and checks both sort-fallback thresholds.
• DataFrameAggregateSuite: added SPARK-58210: empty-buffer merge must not overflow to NaN for statistical aggregates, covering var_pop, covar_pop, regr_sxy across AQE on/off x combine on/off.

Regenerated golden files via SPARK_GENERATE_GOLDEN_FILES=1:

• sql-tests/results/explain-cbo.sql.out
• tpcds-plan-stability approved plans for the affected TPC-DS queries.

A total of 34 golden plans changed, all caused by the two rules enabled in this PR, with no regressions.

Type	Queries
Pure combine	q14a, q14a.sf100, q22a, q22a.sf100, q33, q33.sf100, q46, q46.sf100 (modified & v1_4), q49, q49.sf100, q49.sf100 (v2_7), q51a, q51a.sf100, q56, q56.sf100, q60, q60.sf100, q66, q66.sf100, q68.sf100, q77a.sf100
Pure replace	q16, q16.sf100, q64.sf100, q94, q94.sf100, q95, q95.sf100, q64.sf100 (v2_7)
Both	q23a.sf100, q23b.sf100, q64, q64 (v2_7)

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

Generated-by: Claude Code

[ulysses-you]

cc @viirya @cloud-fan @sunchao @wangyum thank you

[sunchao]

Summary

Prior state and problem

replaceHashWithSortAgg previously defaulted to false, while combineAdjacentAggregation inherited that setting. This PR separates the configurations and enables both physical optimizations by default, making adjacent Partial/Final aggregation eligible for Complete-mode execution without an explicit opt-in.

Design approach

The configuration independence is sensible, and the bottom-up traversal for ReplaceHashWithSortAgg appears consistent with the ordering guarantees exposed by its children. The risky part is enabling CombineAdjacentAggregation for every adjacent pair without first establishing that each aggregate's update/evaluate lifecycle is equivalent to its partial/merge/evaluate lifecycle.

Correctness / compatibility analysis

I found one release-blocking compatibility failure for mode-aware Hive UDAFs and a separate numerical behavior change for built-in statistical aggregates. Both are reachable on single-partition or already-partitioned inputs, under both normal and adaptive preparation paths. I did not find correctness problems in aggregate modes, result attributes, filters, offsets, required distributions, streaming separation, or the four combinations of the configuration flags.

Key design decisions

Keeping the settings independent is the right direction. Default-on combining should fail closed for aggregate implementations that are not known to have stable Complete semantics, or those implementations should be adapted so Complete uses their proper lifecycle.

Implementation sketch

The patch changes both defaults to true, removes the combine setting's fallback, applies adjacent-aggregate combining independently of hash-to-sort replacement, updates affected tests, and regenerates plan goldens and benchmark outputs.

Behavioral changes worth calling out

The change can do more than alter plan shape: it can crash a supported mode-aware Hive UDAF and can change statistical outputs between NaN and 0.0. Because the settings are independent, disabling only hash-to-sort replacement does not restore the previous aggregate staging.

Suggested improvements

Please make Hive UDAFs Complete-safe or exclude them from the rewrite, deliberately fix or guard the statistical merge behavior with AQE/non-AQE parity tests, and document the independent combine rollback setting in the 4.2-to-4.3 migration guide.

[sunchao]

[P1] Please do not disable this regression test to accommodate the new default. A Complete object aggregate calls HiveUDAFFunction.update, which creates a PARTIAL1-mode buffer, and then eval, which passes that buffer to the separately initialized FINAL evaluator. Mode-aware Hive UDAFs may legally use different buffer classes in those modes, so the default rewrite can turn a previously supported query into a ClassCastException. I reproduced this and also verified that a variant which correctly handles Hive's native COMPLETE lifecycle still fails through Spark's mixed-evaluator path. Please make the Hive wrapper Complete-safe or exclude these UDAFs before enabling the rule by default, and keep this test exercising the default path.

[ulysses-you]

Thanks for the detailed repro. Fixed in the latest push by making the Hive wrapper Complete-safe rather than excluding UDAFs or pinning the test.

HiveUDAFFunction keeps two evaluators: partial1HiveEvaluator (PARTIAL1, buffer created in update) and finalHiveEvaluator (FINAL, buffer created in merge). eval always terminated via the FINAL evaluator, so a buffer that only went through update (the Complete-mode path) was a PARTIAL1-mode buffer handed to terminate, which mode-aware UDAFs reject.

merge already converted a PARTIAL1 buffer to a FINAL buffer on demand. I extracted that into toFinalBuffer and call it from eval as well, so Complete-mode (update without merge) terminates against a FINAL buffer. SPARK-24935 now runs on the default (combined) path and asserts the single Complete-mode ObjectHashAggregateExec, with both sort-fallback thresholds still checked. I verified a mode-aware UDAF variant no longer throws ClassCastException.

The change is in sql/hive/.../hiveUDFs.scala; HiveUDAFSuite, ObjectHashAggregateSuite, and UDAQuerySuite pass.

[sunchao]

[P2] Enabling this rewrite changes built-in aggregate results, not only the physical plan. With AQE off and a single-partition group containing two finite 1e155 values, disabling combining returns NaN for var_pop, covar_pop, and regr_sxy, while Complete returns 0.0. The old NaN comes from overflowing the first merge into the empty Final buffer, so 0.0 is mathematically preferable, but a default physical optimization should not silently introduce config-dependent results. Please fix and test the empty-buffer merge as a deliberate correctness change, or keep these aggregates out of combining until the modes are equivalent.

[ulysses-you]

Great catch, and thanks for pinpointing the empty-buffer merge. Fixed as a deliberate correctness change in the merge expressions rather than excluding the aggregates.

Root cause: CentralMomentAgg.merge computes delta * deltaN * n1 * n2 and Covariance.merge computes dx * dyN * n1 * n2. When one side is an empty buffer (n1 == 0 or n2 == 0) and the other side has a large average, delta * deltaN (or dx * dyN) overflows to Infinity before being multiplied by the zero count, and Infinity * 0 = NaN corrupts the merged moments. The existing newN === 0 guard only covers both-empty. The old default (no combining) hit this on a single-partition group of two large equal values, returning NaN; the Complete-mode default sidestepped the merge and returned the correct 0.0, so the two configs disagreed.

Fix: force delta/dx/dy to 0.0 when either side is empty, and set newAvg/newXAvg/newYAvg directly to the non-empty side's average. The delta terms then vanish cleanly and the merged buffer equals the non-empty side. Both the combined and the separate paths now return 0.0 for var_pop/covar_pop/regr_sxy.

Added SPARK-58210: empty-buffer merge must not overflow to NaN for statistical aggregates in DataFrameAggregateSuite, covering var_pop, covar_pop, regr_sxy across AQE on/off x combine on/off. The group-by, linear-regression, and aggregates_part1 SQL golden files are unchanged (regenerated to confirm).

[sunchao]

[P2] Please document spark.sql.execution.combineAdjacentAggregation here as well. Although the config itself is introduced during 4.3 development, its default-on execution behavior is new for users upgrading from 4.2. The two settings are now independent, so following this rollback instruction and disabling only replaceHashWithSortAgg still leaves adjacent Partial/Final aggregation combined. Please restore the removed entry and state that both settings may need to be false to recover the previous staging.

[ulysses-you]

Good point — restored. The migration guide now has a separate entry for spark.sql.execution.combineAdjacentAggregation under "Upgrading from Spark SQL 4.2 to 4.3", and it states explicitly that the setting is independent of replaceHashWithSortAgg, so disabling only replaceHashWithSortAgg still leaves adjacent aggregation combined; both must be set to false to fully restore the previous partial/final staging.

[viirya]

I traced the two fixes from the last round (Hive UDAF Complete-safety and the empty-buffer overflow) against the PR-head source, and both look correct and complete.

The toFinalBuffer extraction is a clean refactor rather than new logic — it reuses the exact PARTIAL1->FINAL conversion merge already did. For the previously-reachable modes (Final/PartialMerge), the buffer is already canDoMerge=true, so toFinalBuffer returns it untouched and behavior is unchanged; only the newly-reachable Complete path (update-without-merge) gets converted, which is exactly the fix. Non-mode-aware UDAFs stay correct too, and there's no extra work on the Final path.

The statistical-merge fix is also complete. Forcing delta/dx/dy to 0.0 when one side is empty makes every higher-order term (m2/m3/m4, ck) collapse cleanly to the non-empty side, and setting newAvg/newXAvg/newYAvg directly finishes it — the merged buffer equals the non-empty side exactly. The both-sides-non-empty path is byte-identical, so existing results don't regress. The AQE x combine parity test reproduces the reported case directly. The cross-reference comment in Covariance pointing back to CentralMomentAgg is a nice touch for maintainability.

One thing worth documenting: ReplaceHashWithSortAgg switches transformDown to transformUp, but this isn't mentioned in the description, the commit messages, or anywhere in the discussion. It looks intentional and correct to me — SortAggregateExec is order-preserving while HashAggregateExec is not, so replacing a child hash agg with a sort agg changes that child's outputOrdering, and bottom-up traversal lets a parent then see the satisfied ordering and be replaced in turn (each replacement is still gated by the ordering check, so it can only enable more valid replacements, never wrong ones). That's a genuine behavior change that likely accounts for some of the 34 golden-plan diffs — could you add a sentence to the PR description (or an inline comment) explaining why bottom-up is needed here? It'll save the next reader from having to reconstruct the reasoning.

Otherwise LGTM — the config decoupling, the test updates that toggle both settings, and the migration-guide entries all read cleanly.

[ulysses-you]

Thanks @viirya for the thorough trace — agreed on both fixes. Re: transformDown -> transformUp: good call, I've added a bullet to the PR description explaining it. The reasoning matches what you laid out: HashAggregateExec does not expose a grouping-key outputOrdering (Nil), while SortAggregateExec is order-preserving on the grouping keys. With transformDown an outer aggregate inspected before its inner partial aggregate sees the partial's Nil ordering and is not replaced; with transformUp the partial is replaced first, so its sort-aggregate output ordering satisfies the outer aggregate, which is then also replaced. Each replacement is still gated by the orderingSatisfies check, so bottom-up can only enable more valid replacements, never wrong ones. The "replace partial and final hash aggregate together with sort aggregate" case in ReplaceHashWithSortAggSuite (SMJ -> partial -> final, expecting both replaced) covers it.

[sunchao]

Follow-up review of the current head: one remaining statistical-aggregate correctness issue and two independently verified regression-test coverage gaps.

[sunchao]

[P2] Please extend this empty-buffer fix and AQE-by-combining regression to corr and regr_r2 as well. PearsonCorrelation.merge still evaluates dx * dxN * n1 * n2 and dy * dyN * n1 * n2 before the zero-count multiplication, so the same Infinity * 0 corruption survives this change. I reproduced this at this exact head with AQE disabled and Seq((1.0e155, 1.0e-150), (1.000000000000001e155, 2.0e-150)).toDF("x", "y").repartition(1).selectExpr("corr(x, y)", "regr_r2(y, x)"): with combineAdjacentAggregation=true, the result is (1.0, 1.0000000000000004); with it disabled, both results are NaN. Since this PR enables combining by default, a physical-optimization toggle still changes the observable answers. Please guard the Pearson empty-buffer merge as well and cover both expressions across AQE on/off and combining on/off.

[ulysses-you]

Good catch -- fixed by guarding the Pearson merge the same way Covariance/CentralMomentAgg already are, rather than excluding the expressions.

PearsonCorrelation.merge computed dx * dxN * n1 * n2 (and the dy variant) for xMk/yMk before the zero-count multiply, so when one side is an empty buffer and the other has a large average, dx * dxN overflowed to Infinity and Infinity * 0 = NaN corrupted the moments. I now force dx/dy to 0.0 when either side is empty and set newXAvg/newYAvg directly to the non-empty side's average, mirroring the existing fix (with a back-reference comment to CentralMomentAgg).

Added SPARK-58210: empty-buffer merge must not overflow to NaN for Pearson correlation in DataFrameAggregateSuite, using your exact repro ((1.0e155, 1.0e-150), (1.000000000000001e155, 2.0e-150), repartitioned to 1) and asserting corr(x, y) and regr_r2(y, x) are finite and config-independent across AQE on/off x combine on/off.

[sunchao]

[P3] Please construct a fresh query inside each fallback-threshold withSQLConf block. This one DataFrame is executed first with threshold 1, and SparkPlan.executeRDD memoizes the resulting RDD; the subsequent checkAnswer(df, ...) under threshold 100 reuses that already-created execution instead of constructing an ObjectHashAggregateExec execution under the new threshold. I verified that both checks use the same execution-RDD ID, so this does not exercise both fallback and non-fallback despite the test's claim. ObjectHashAggregateSuite also explicitly cautions against reusing a DataFrame while changing this threshold. Rebuilding sql(...) inside each configuration, and ideally asserting the fallback metric, would make the lifecycle fix actually cover both paths.

[ulysses-you]

You're right -- the single reused DataFrame meant executeRDD replayed the threshold-1 execution under threshold 100, so only one path ran. Fixed: I rebuild sql(query) inside each withSQLConf block, and now assert the numTasksFallBacked metric -- > 0 at threshold 1 (fallback) and == 0 at threshold 100 (non-fallback) -- so both paths are actually exercised. Added a numFallbackTasks helper summing the metric across the ObjectHashAggregateExec nodes.

[sunchao]

[P3] Please add a traversal regression with replaceHashWithSortAgg=true and combineAdjacentAggregation=false, and assert that both the partial and final nodes become SortAggregateExec. This helper currently enables combining together with replacement, so the partial/final pair is collapsed before the test can establish that bottom-up traversal replaces both nodes; the together case therefore observes only one sort aggregate. I mutation-tested the exact head by reverting transformUp to transformDown: both the normal and AQE together tests still passed. Consequently the test cited as coverage for the traversal change cannot detect reverting that change. A sorted-input two-node case with combining disabled, run with AQE both on and off, would fail under transformDown and cover the behavior.

[ulysses-you]

Agreed, the together case enabled combining and collapsed the pair before traversal, so it couldn't detect reverting transformUp. Added SPARK-58210: bottom-up traversal replaces both partial and final hash aggregate with replaceHashWithSortAgg=true and combineAdjacentAggregation=false, over the SMJ -> partial -> final shape, asserting 0 hash / 2 sort.

I ran the same mutation you did: reverting transformUp -> transformDown now fails this test in both AQE modes (1 hash agg remains instead of 0), and it passes once reverted back. So the traversal change is now genuinely covered.

[sunchao]

Thanks for addressing the review comments. The Pearson merge fix, independently exercised Hive fallback paths, and bottom-up aggregation regression all look good. LGTM.

[wangyum]

Is it possible splitting this PR into two before merging, because it currently bundles two logically independent changes:

1. The merge-expression correctness fix (standalone bug).
2. The default flip + config decoupling

Bundling them makes the merge fix impossible to backport to 4.2.x / 3.5.x. If we later find a regression from the default flip and want to revert just that, the correctness fix would go with it.

[ulysses-you]

Good call @wangyum, agreed. I've split out two standalone correctness fixes, each independent of the default flip and backportable on its own:

#57460 (SPARK-58291) -- statistical-aggregate empty-buffer merge overflow:

• PearsonCorrelation / Covariance / CentralMomentAgg merge overflow guard
• the two DataFrameAggregateSuite regression tests

#57464 (SPARK-58294) -- Hive UDAF Complete-mode buffer conversion:

• hiveUDFs.scala eval PARTIAL1 -> FINAL buffer conversion
• the HiveUDAFSuite regression test

On the hiveUDFs fix: I initially thought it couldn't be backported independently, but it can -- the Complete-mode path is reachable today by setting spark.sql.execution.combineAdjacentAggregation=true (SPARK-43317), without this PR's default flip. So it stands alone. Its new test explicitly enables that config rather than relying on the flipped default, and I mutation-tested it (reverting the eval fix reproduces the ClassCastException).

Once both are reviewed, I'll rebase this PR to drop those files, leaving only the default flip + config decoupling (plus the golden-file and benchmark updates) here.

Ordering note: #57460 and #57464 should merge first, then this PR rebases on top. If this PR drops the fixes before they land, its own DataFrameAggregateSuite / HiveUDAFSuite cases would reintroduce the behavior they check for.

cc @sunchao @viirya