[SPARK-58511][SQL] Bypass ineffective pre-shuffle partial aggregation at runtime - #57742
[SPARK-58511][SQL] Bypass ineffective pre-shuffle partial aggregation at runtime#57742ulysses-you wants to merge 2 commits into
Conversation
… at runtime When a pre-shuffle partial aggregation is not reducing rows (the distinct-key ratio is too high), maintaining an aggregation map is not worthwhile. This change makes hash aggregation detect that at runtime and bypass the partial aggregation: the remaining input rows are passed through as single-row partial buffers that the downstream Final aggregation merges, avoiding the cost of maintaining and spilling a large map. Two decision tiers, both evaluated only while the regular (second-level) map is still in memory: - no-spill tier: from a sample of rows on, bypass if the reduction ratio is at least `noSpillReductionRatioThreshold`. The sampling window doubles after each sub-threshold check, so low-cardinality input is re-checked only rarely while a late high-cardinality tail can still be caught. - on-spill tier: when the map is about to spill, bypass instead if the ratio is at least `spillReductionRatioThreshold`. Only pre-shuffle `Partial` hash aggregation with grouping keys is eligible. Both the codegen path (HashAggregateExec) and the interpreted path (TungstenAggregationIterator) are covered. Once pass-through is active the maps are frozen, so they are output (and their memory released) before the remaining input is streamed. New configs under spark.sql.execution.aggregate.adaptivePartialAggregation.* (enabled by default), a `numBypassingRows` metric, tests, and a benchmark are added. Co-Authored-By: Claude <noreply@anthropic.com>
|
cc @viirya @dongjoon-hyun @sunchao @LuciferYang @cloud-fan thank you |
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
1 blocking, 2 non-blocking, 2 nits.
The optimization is promising and broadly tested, but the policy should be simplified before adding several independently tunable configuration contracts.
Design / architecture (1)
- sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:4188: Start with one minimum-row setting and one reduction threshold shared by periodic and pre-spill decisions. -- see inline
Correctness (1)
- sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:913: The spill-tier ratio uses different row sets in code-generated and interpreted execution. -- see inline
Suggestions (1)
- sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala:257: The interpreted hot loop repeatedly checks and unwraps immutable configuration. -- see inline
Nits: 2 minor items (see inline comments).
Verification
I traced eligibility and pass-through behavior through both code-generated and interpreted execution, including the pre-spill decision, map output, single-row partial-buffer construction, and downstream final aggregation. The two paths currently use different spill-tier denominators. No tests were run as part of this review.
| .checkValue(_ > 0, "The sample row count must be positive.") | ||
| .createWithDefault(100000) | ||
|
|
||
| val ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD = |
There was a problem hiding this comment.
Could we start with one minimum-row setting and one reduction threshold, applying the same policy both periodically and immediately before spilling? The separate no-spill and spill thresholds plus exponential resampling add configuration and behavioral complexity without showing that these dimensions must be independently tunable. At the spill boundary, switch to pass-through when the common policy says aggregation is ineffective; otherwise spill normally. This also gives the code-generated and interpreted paths one invariant to implement and test.
There was a problem hiding this comment.
Adding two configs is a trade-off, due to that it is difficult to accurately determine how much data reduction rate can offset the performance loss caused by sort and spill. There is a way to use the spill threshold, setting spill thresholds to 0, so it can always passing through directly when spill happens.
I set both of these configs to the same default value to make their behavior the same by default.
There was a problem hiding this comment.
Thanks, the shared default plus the concrete spill-cost rationale addresses my concern about introducing two policies by default.
| |// generate grouping key | ||
| |${unsafeRowKeyCode.code} | ||
| |if (!$adaptivePassThroughTerm) { | ||
| | $regularMapRowCountTerm += 1; |
There was a problem hiding this comment.
Evaluate the spill ratio over the same row set in both execution paths. This increment counts the failed insertion that becomes the first pass-through row, while the interpreted path evaluates before counting it. Move the increment after the spill decision and add an exact threshold-boundary test with codegen enabled and disabled.
There was a problem hiding this comment.
good catch, addressed
There was a problem hiding this comment.
Confirmed: both paths now evaluate the spill ratio over the pre-failure rows, and the boundary test covers codegen on and off.
| // the reduction ratio is too high to be worthwhile, bypass partial aggregation for the | ||
| // rest. The window doubles after each sub-threshold check so low-cardinality input is | ||
| // re-evaluated only rarely while a late high-cardinality tail can still be caught. | ||
| if (adaptivePartialAggConfig.isDefined && externalSorter == null && |
There was a problem hiding this comment.
Extract stable primitive settings before entering the input loop. The current Option.isDefined/get checks repeat for every aggregated row even though eligibility and thresholds cannot change during the iterator's lifetime.
There was a problem hiding this comment.
good point, addressed
There was a problem hiding this comment.
Confirmed, the immutable adaptive settings are now extracted before the hot loop.
| | $sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter()); | ||
| |} | ||
| |$resetCounter | ||
| |// the hash map had be spilled, it should have enough memory now, |
There was a problem hiding this comment.
| |// the hash map had be spilled, it should have enough memory now, | |
| |// the hash map had been spilled, so it should have enough memory now, |
|
|
||
| // Force the regular map to spill quickly and disable the no-spill tier (huge sample). With | ||
| // fully distinct keys the reduction ratio is 1.0, so at the spill boundary the on-spill | ||
| // (Tier 2) tier bypasses instead of spilling; the baseline spills repeatedly and falls back |
There was a problem hiding this comment.
| // (Tier 2) tier bypasses instead of spilling; the baseline spills repeatedly and falls back | |
| // tier (Tier 2) bypasses instead of spilling; the baseline spills repeatedly and falls back |
- Make `spillReductionRatioThreshold` fall back to `noSpillReductionRatioThreshold` so both tiers apply one policy by default. - Evaluate the spill-tier ratio over the same row set in both execution paths: the codegen path now counts only the rows that made it into the map, matching the interpreted path, and skips the decision until at least one row is counted. Add a test asserting both paths decide identically at the exact ratio boundary. - Unwrap the adaptive config once before the interpreted input loop instead of re-checking the `Option` for every row. - Fix two comment typos. Co-Authored-By: Claude <noreply@anthropic.com>
| s"falls back to '${ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD.key}' " + | ||
| "so that both tiers apply one policy by default. Setting it lower makes the bypass more " + | ||
| "likely once spilling is imminent, and 0 always bypasses instead of spilling.") | ||
| .version("4.3.0") |
There was a problem hiding this comment.
Just leaving a note that branch-4.3 has already been cut, so please check whether these changes are intended to go there as well @ulysses-you
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
4 addressed, 1 remaining, 3 new. (0 newly introduced, 3 late catches, 0 previously raised.)
3 blocking, 1 non-blocking, 0 nits.
The partial-buffer mechanism is sound, but the decision policy and spill lifecycle should converge with DBR before this becomes a second implementation contract.
Remaining from prior review (1)
- Could we start with one minimum-row setting and one reduction threshold? DBR already implements this feature and this OSS implementation will eventually be imported there, so please align normal and spill-boundary checks on the same
rows / keys <= minCompactionpredicate afterminRowsrather than creating a second operational contract. -- existing thread
Design / architecture (2)
- sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:206: Measure the operator's total processed rows and total in-memory keys rather than only regular-map traffic. -- see inline
- sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:925: Keep adaptive pass-through available for new in-memory map epochs after earlier spills. -- see inline
Suggestions (1)
- sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:814: Propagate the child's result-copy requirement instead of forcing copies for every adaptive aggregate. -- see inline
Verification
I compared the OSS implementation with DBR's code-generated and interpreted adaptive partial aggregation paths. DBR uses one minimum-row setting and one compaction threshold, counts fast-map and regular-map keys over all processed rows, evaluates new in-memory map epochs after spills, and propagates the child's result-copy requirement. No tests were run as part of this review.
| // Codegen state for adaptive partial aggregation. When the pre-shuffle reduction ratio of the | ||
| // regular (second-level) hash map is too low, the operator stops populating the map and instead | ||
| // streams each remaining row through as a single-row partial buffer for the Final aggregate to | ||
| // merge. Only the regular map is governed: the append-only fast hash map keeps absorbing hot |
There was a problem hiding this comment.
Please measure all processed rows against the total number of in-memory keys, including both maps. DBR uses the operator-level invariant—equivalent to fastMap.rowCount + regularMap.getNumKeys over all processed rows—so limiting both sides to regular-map traffic makes two-level-map routing change the decision and creates a second policy to maintain.
| | } | ||
| | if ($unsafeRowBuffer != null) { | ||
| | $regularMapRowCountTerm += 1; | ||
| | if ($sorterTerm == null && |
There was a problem hiding this comment.
Please keep adaptive pass-through available after earlier spills. DBR intentionally resets the processed-row count after each spill and evaluates the new in-memory map epoch; if it activates, the existing sorter and current map drain before the remaining rows pass through. Permanently gating on sorter == null materially diverges for inputs whose cardinality becomes unfavorable later.
| // appended in the same child loop iteration before any drain, and they all alias the single | ||
| // result `UnsafeRow`. Copy the result so the buffered rows do not collapse into the last one. | ||
| override def needCopyResult: Boolean = | ||
| adaptivePartialAggConfig.isDefined || super.needCopyResult |
There was a problem hiding this comment.
Please propagate the child's copy requirement instead of enabling copies for every adaptive aggregate. DBR gates this on the child, and ExpandExec already reports needCopyResult = true, so doing the same preserves the aliasing fix without adding row.copy() to ordinary output.
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
I traced the pass-through mechanism end to end on both paths and the core is correct: a bypassed row builds a fresh initialValues buffer, runs the normal update, and is emitted as (key, single-row partial buffer) — exactly what the partial map would hold for a one-row group, so the Final merge produces the same result. The eligibility gate is right (requiredChildDistributionExpressions.isEmpty correctly excludes the group-by-only Final phase, with a good comment on why the mode check alone would admit it vacuously), the Expand-aliasing copy is genuinely needed, and the interpreted path mirrors the codegen contract (copyFrom(initialAggregationBuffer) + processRow). The correctness matrix in the test suite is broad and always compares against the feature-off reference, so results can't silently diverge. No result-correctness concerns from me.
The remaining open items are all about policy/efficiency/metrics rather than result correctness — worth converging but none change query output:
aggTimeunder-reports when pass-through fires (new). The timer wraps only the firstdoAgginsideif (!initAgg)(HashAggregateExec.scala:792-794); once pass-through activates mid-build, the bulk of the remaining input is consumed throughadaptiveResumeBuild'sdoAggcall (line 771), which is outside the timer. For a high-cardinality input that bypasses early, most of the build work isn't counted, so the SQL UI's aggregation time is materially low. Consider timing each resumed build (or the whole build lifecycle).- Config
version("4.3.0")vs. the actual target.dev/next_version_candidates.pynow reportsbranch-4.x -> 4.4.0(branch-4.3 is cut, as @uros-b noted), so unless this is intended for a branch-4.3 backport the configs should say4.4.0— otherwise the SQL config docs claim a released version already has this feature. This is still an open question on the thread. - I agree with @cloud-fan's two remaining design points, and I'd frame them as effectiveness rather than correctness so they're easier to weigh: the permanent
sorter == nullgate means an input whose cardinality turns unfavorable only after an early spill can never bypass (the doc's "a late high-cardinality tail can still be caught" holds only on the no-spill path); andneedCopyResult = adaptivePartialAggConfig.isDefined || super.needCopyResultover-copies — it's safe, but sinceHashAggregateExecis a blocking operatorsuper.needCopyResultis alwaysfalseand won't pick up the child, so gating on the child's requirement (whichExpandExecalready reports astrue) preserves the aliasing fix while dropping the copy for ordinary single-output children.
The mechanism itself is sound and the tests are thorough; these are convergence/accuracy items rather than blockers on result correctness, but the version tag and the aggTime gap are worth resolving before merge.
4 participants
What changes were proposed in this pull request?
This PR makes hash aggregation detect at runtime that a pre-shuffle partial aggregation is not reducing rows, and bypass it: the remaining input rows are passed through as single-row partial buffers that the downstream
Finalaggregation merges, so the output contract is unchanged.Two decision tiers, both evaluated only while the regular (second-level) map is still fully in memory (
sorter == null):sampleRowsregular-map rows on, bypass ifdistinctKeys / regularRows >= noSpillReductionRatioThreshold. The sampling window doubles after each sub-threshold check, so a low-cardinality input is re-evaluated only rarely while a late high-cardinality tail can still be caught.spillReductionRatioThreshold. This threshold is more aggressive because a spilling partial aggregation starts paying disk I/O.Only a pre-shuffle
Partialhash aggregation with grouping keys is eligible (requiredChildDistributionExpressions.isEmptyidentifies that phase and, importantly, keeps a group-by-only aggregate'sFinalphase out - itsaggregateExpressionsis empty, so a mode-only check would admit it vacuously). Both the codegen path (HashAggregateExec) and the interpreted path (TungstenAggregationIterator) are covered.Only the regular (second-level) map is governed, and the reduction ratio uses the rows that entered it as the denominator. The append-only fast hash map never spills, so a fast-map hit counts in neither the numerator nor the denominator; using total rows would keep the ratio below the threshold even for fully distinct input, because the fast map absorbs the first 2^16 keys.
Once pass-through is active the maps are frozen, so they are output - which also frees their memory - before the remaining input is streamed, rather than being held until the end.
DISTINCTaggregates are eligible: in the multi-phase distinct plan the intermediatePartialMergephase is notPartialmode (and requires a distribution), so it always aggregates and de-duplicates, and the rows reaching the distinctPartialphase therefore carry exactly one distinct value each.New configs, all under
spark.sql.execution.aggregate.adaptivePartialAggregation.*:enabledtruesampleRows100000noSpillReductionRatioThreshold0.9spillReductionRatioThresholdfallback noSpillReductionRatioThresholdA
numBypassingRowsSQL metric reports how many rows bypassed.Why are the changes needed?
For high-cardinality grouping keys the pre-shuffle partial aggregation reduces little or nothing, but still pays for maintaining - and often spilling - an aggregation map as large as the input. Bypassing it at runtime removes that cost while keeping the two-phase plan intact, so the decision needs no planner-side statistics and adapts per task.
Benchmark (
AdaptivePartialAggregationBenchmark, run in GitHub Actions, JDK 25):The bypassing cases win in both tiers and on both execution paths; the low-cardinality cases, where the tiers correctly decline to bypass, show no regression beyond the small per-row sampling overhead.
This is related to, but independent of, the static
spark.sql.execution.bypassPartialAggregation(SPARK-57688), which drops the partial aggregation at planning time. The runtime version keeps the partial aggregation when it does reduce rows and only bypasses when the observed ratio says it does not; the two can be used together.Does this PR introduce any user-facing change?
No, query results are unchanged. The feature is enabled by default and only changes how the pre-shuffle partial aggregation is executed, plus the new
numBypassingRowsmetric in the SQL UI.How was this patch tested?
New
AdaptivePartialAggregationSuite(34 tests), in two halves:FILTER (WHERE ...)), key types (string, decimal, date, nullable), group-by-only aggregates with duplicate keys, empty input, andExpand-bearing plans (ROLLUP / CUBE / GROUPING SETS / multi-distinct).numBypassingRowsmetric proves the bypass fires when (and only when) it should - high cardinality bypasses, low cardinality does not, the feature switch and eligibility rules are honored, and both tiers work. These tests also compare against the feature-off reference so a bypassing run can never pass on metrics alone. Forcount(DISTINCT v) GROUP BY kthe bypasses of the twoPartialphases are told apart by grouping-key count and asserted separately.New
AdaptivePartialAggregationBenchmarkcovering the {high, low}-cardinality x {no-spill, on-spill} grid, each across codegen on/off and the adaptive switch; results are included.Generated code for simple query
select c2, count(*) from t3 group by c2:Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.5)