peter-toth commented on code in PR #57742: URL: https://github.com/apache/spark/pull/57742#discussion_r3734471967
########## sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala: ########## @@ -0,0 +1,961 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.aggregate + +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.aggregate.Partial +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests for runtime adaptive partial aggregation + * (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). When a partial aggregate is not reducing + * rows, the operator stops aggregating and streams the remaining rows through as single-row partial + * buffers for the Final aggregate to merge. It must never change results. + * + * The suite has two halves: + * 1. Correctness: output is identical to the reference (feature-off) run across the full matrix + * of codegen on/off, two-level map on/off, and spill/no-spill, over a range of aggregate + * shapes, key types, and `Expand`-bearing plans (ROLLUP / CUBE / GROUPING SETS / + * multi-distinct). + * 2. Triggering: the `numBypassingRows` metric proves the bypass actually fires when (and only + * when) it should -- high-cardinality input bypasses, low-cardinality input keeps aggregating, + * the feature switch and eligibility rules are honored, and both check points work. + */ +class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession + with AdaptiveSparkPlanHelper { + + import testImplicits._ + + // A `testFallbackStartsAt` setting ("fastMapCounter, regularMapCounter") that makes the regular + // map fall back (spill) periodically, exercising the spill-check decision path in both the + // codegen and interpreted aggregation paths. Kept moderate so low-cardinality inputs (which are + // never bypassed and therefore really spill) do not open an unbounded number of spill readers. + private val forceSpillFallback = "4, 16" + + // The upstream `CombineAdjacentAggregation` and `ReplaceHashWithSortAgg` rules would change the + // plan of these small single-partition queries away from a Partial+Final `HashAggregateExec`: + // the former merges the two adjacent phases (no shuffle in between) into a single `Complete` + // aggregate, and the latter converts a hash aggregate to a sort aggregate when the input is + // already sorted by the grouping key (a `Range` over an ascending `id` key). The adaptive + // feature lives in the partial hash aggregation, so both rules are disabled to keep that + // structure in the tests. + private val fixedPlanConfs = Seq( + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false", + SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "false") + + /** + * Runs `df` with adaptive partial aggregation disabled (the reference) and then across the full + * configuration matrix with it enabled, asserting every enabled run matches the reference. + */ + private def checkAdaptiveMatchesReference(build: () => DataFrame): Unit = { + val reference = withSQLConf( + (SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") +: fixedPlanConfs: _*) { + build().collect().toSeq + } + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + forceSpill <- Seq(true, false) Review Comment: **Finding 3.** `checkAdaptiveMatchesReference` never asserts that the bypass actually fired, and for four of its callers it never does. I instrumented the exact helper on `f2616ecb556` and read `numBypassingRows` off the partial `HashAggregateExec` in each of the 8 cells: | test | bypassed rows (all 8 cells) | | --- | --- | | `results unchanged with nullable grouping keys` (:266) | **0** | | `results unchanged for ROLLUP` (:493) | **0** | | `results unchanged for CUBE` (:502) | **0** | | `results unchanged for GROUPING SETS` (:511) | **0** | | the other callers (high-cardinality, many-agg-functions, string keys, decimal/date keys, avg, filtered, group-by-only, distinct, medium-cardinality, fused union, generator) | 168 – 1192 | The three `Expand` ones matter most: the comment above them says they "assert results stay correct with the exclusion absent" — i.e. they are the justification for not porting apache/spark#28804's `Expand` exclusion — but with zero bypassed rows they assert nothing about it. (`multi-distinct` does bypass, and the separate `pass-through fires for high-cardinality input below an Expand` triggering test does cover `Expand` + bypass against the reference, so the risk is covered; these four tests just aren't the thing covering it.) The inputs are all above the threshold by construction: `nullable grouping keys` has 301 keys over 400 rows (ratio 1.33), and `ROLLUP` expands 400 rows into 1200 over 601 keys (ratio 2.0). Two changes: 1. Give `checkAdaptiveMatchesReference` an `expectBypass: Boolean = true` parameter and assert `numBypassingRows > 0` (pass `false` for `low-cardinality`, `group-by-only with duplicate keys` and `empty input`, which are correctly meant not to fire). That is what stops these from silently going vacuous again. 2. Adjust the four inputs so they do fire. For the nullable one, a sparser null makes it distinct-dominated — this is worth getting right, because a null key is the one case where the fast-map probe is skipped (`if (${fastRowKeys.map("!" + _.isNull)...})`), so a bypassed null key is a genuinely distinct path: ```scala checkAdaptiveMatchesReference { () => spark.range(0, 400, 1, 1) .select( when($"id" % 100 === 0, lit(null)).otherwise($"id") as "k", $"id" as "v") .groupBy($"k") .agg(sum($"v") as "s", count(lit(1)) as "c") } ``` For ROLLUP/CUBE/GROUPING SETS, either make the grouping columns fully distinct so the grand-total group can't lift the ratio, or set `minCompaction` high for those three the way `forEachCodegenAndMap` already lets you. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala: ########## @@ -711,6 +995,54 @@ case class HashAggregateExec( } else { findOrInsertRegularHashMap } + + // Every row is either accepted by an aggregation map or streamed through -- the fast map + // serves a row without it ever reaching the regular map, so both buffers are consulted to + // tell the two apart. + // + // An accepted row counts toward the compaction ratio, so the numerator matches the + // operator-level denominator. Counting inside the regular-map branch alone would drop the + // rows the fast map absorbed from the ratio and bypass an aggregation that is in fact + // reducing. A row no map holds is streamed once pass-through is active: both probes are + // skipped (guarded above), so neither buffer is set, and `rowBypassed` marks exactly those + // rows. The row that fails to insert at the spill boundary lands here too, while the row + // that merely flipped pass-through at the check point is already aggregated in the map that + // took it and must not be re-emitted. + val countOrPassThroughRow = if (adaptivePartialAggEnabled) { + val heldByAMap = if (isFastHashMapEnabled) { + s"($fastRowBuffer != null || $unsafeRowBuffer != null)" + } else { + s"($unsafeRowBuffer != null)" + } + // The grouping key was already projected in `findOrInsertRegularHashMap` + // (`unsafeRowKeyCode.code`), so `unsafeRowKeyCode.value` holds this row's key. Only build + // the single-row partial buffer here. + s""" + |if ($heldByAMap) { + | if (!$adaptivePassThroughTerm) { + | $processedRowsTerm += 1; + | if ($processedRowsTerm == $adaptiveNextCheckRowTerm) { Review Comment: **Finding 2.** The first periodic check fires after `minRows` (100,000) rows no matter how many rows the task will process, and the flip is irreversible — nothing ever re-arms `adaptivePassThrough`. So a distinct-heavy prefix commits the rest of the task, and the downside is unbounded while the upside is bounded by the cost of maintaining the map. Measured on `f2616ecb556`, single partition, **all settings at their defaults** (100k distinct keys, then 400k rows over 100 keys — overall compaction ratio 500000/100100 = 5.0, i.e. the partial aggregation is clearly worth keeping): ```scala spark.range(0, 500000, 1, 1) .select(when($"id" < 100000, $"id").otherwise($"id" % 100) as "k", $"id" as "v") .groupBy($"k").agg(sum($"v")) ``` | adaptive | partial `numOutputRows` | `numBypassingRows` | | --- | --- | --- | | false | 100,000 | 0 | | true | **500,000** | 400,000 | Identical with `wholeStage` on and off. That is 5x the shuffle rows and 5x the merge work in the `Final` aggregate, for a query where the feature should not have fired at all. The 1.05 threshold doesn't protect against this: the prefix genuinely has a ratio of 1.0, the check is just looking at an unrepresentative sample. Two things would each help a lot, and they compose: - Make the check point grow, so a short unrepresentative sample can't commit the task — e.g. `adaptiveNextCheckRow *= 2` instead of `+= minRows`, keeping `minRows` as the first check point. A task then re-evaluates at 100k, 200k, 400k, … and the sample backing the decision always scales with the input actually seen. - Allow the bypass to be un-armed. The map is still there and frozen; on a later check point you could re-probe and resume aggregating if the observed distinct-key rate improved. If you'd rather not, please say so in the config doc — "once bypassed, the rest of the task's input is passed through" is a property users need to know when tuning `minRows`. Since the feature is on by default I think one of these should land in this PR rather than as a follow-up. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala: ########## @@ -141,6 +148,40 @@ case class HashAggregateExec( .map(_.asInstanceOf[DeclarativeAggregate]) private val bufferSchema = DataTypeUtils.fromAttributes(aggregateBufferAttributes) + /** + * Whether adaptive partial aggregation applies to this operator. When it does, the aggregation + * may bypass partial aggregation at runtime and pass the remaining input rows through as + * single-row partial buffers (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). It only + * applies to a pre-shuffle partial aggregation with grouping keys: + * - `Partial` mode only: the downstream `Final` aggregation merges the passed-through + * single-row buffers, so the output contract is unchanged. `Final`/`Complete` produce the + * result themselves and have no such downstream. `PartialMerge` does have one and could be + * supported by passing its incoming buffer through unchanged, but that is left for later. + * - grouping keys present: a global aggregation produces a single output row, so partial + * aggregation achieves the maximum reduction and must never be bypassed. + * - DISTINCT aggregate functions are allowed: the intermediate `PartialMerge` phase of the + * multi-phase distinct plan is not `Partial` mode (and requires a distribution), so it always + * aggregates and de-duplicates, and the passed-through rows from the distinct `Partial` phase + * therefore carry exactly one distinct value each. + */ + private val adaptivePartialAggEnabled: Boolean = { + conf.adaptivePartialAggregationEnabled && + groupingExpressions.nonEmpty && + // Only the pre-shuffle partial aggregation has a downstream `Final` to merge passed-through + // single-row buffers. `requiredChildDistributionExpressions` is `None` exactly for that + // pre-shuffle phase and `Some` for the `Final`/`Complete` phase. This check is what keeps a + // group-by-only aggregate (no aggregate functions, so an empty `aggregateExpressions`) from + // being admitted vacuously: `aggregateExpressions.forall(_.mode == Partial)` alone is true + // for the empty list, which would wrongly make the `Final` phase eligible as well. + requiredChildDistributionExpressions.isEmpty && Review Comment: **Finding 8.** The comment above says `requiredChildDistributionExpressions` "is `None` exactly for that pre-shuffle phase and `Some` for the `Final`/`Complete` phase". The second half holds, the first doesn't: `AggUtils.planAggregateWithOneDistinct` builds `partialDistinctAggregate` (step 3, `AggUtils.scala:291`) by calling `createAggregate` without the argument, so it defaults to `None` even though its child is `partialMergeAggregate` and it therefore sits *after* the shuffle. That aggregate is precisely what your `count distinct: the distinct partial aggregate bypasses` test asserts on, so the gate is relying on a case the comment says can't exist. The gate is still correct, just for a simpler reason worth stating instead: all-`Partial` modes imply there is a downstream `PartialMerge`/`Final` that merges buffers, wherever the aggregate sits relative to the shuffle, and `requiredChildDistributionExpressions.isEmpty` is needed only to keep a group-by-only `Final` (empty `aggregateExpressions`, so `forall` is vacuous) out. Rewording it that way also removes the pre-shuffle framing that finding 9 trips over. ########## sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt: ########## @@ -0,0 +1,56 @@ +================================================================================================ +high-cardinality input, no-spill pass-through (Tier 1) Review Comment: **Finding 4.** These section titles come from `runBenchmark(...)`, and none of the four match the benchmark as committed: | results file | `AdaptivePartialAggregationBenchmark.scala` | | --- | --- | | `high-cardinality input, no-spill pass-through (Tier 1)` | `high-cardinality input, pass-through at the periodic check` | | `low-cardinality input, no-spill pass-through (Tier 1)` | `low-cardinality input, pass-through at the periodic check` | | `high-cardinality input, on-spill pass-through (Tier 2)` | `high-cardinality input, pass-through at the spill check` | | `low-cardinality input, on-spill pass-through (Tier 2)` | `low-cardinality input, pass-through at the spill check` | So all three `*-results.txt` files were generated before the two-tier policy was collapsed in `e1bbdd717e0`, which also means they were measured with `minCompaction` defaulting to `1.1` rather than the `1.05` you just switched to. These numbers are the evidence for turning a runtime optimization on by default, so please regenerate them against the head: ``` SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/Test/runMain \ org.apache.spark.sql.execution.benchmark.AdaptivePartialAggregationBenchmark" ``` Unrelated nit while you're in that file: the spill scenarios disable the periodic check with `ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS -> Long.MaxValue.toString`, but `0` is the documented sentinel for that now — using it here would also exercise the sentinel. ########## docs/sql-performance-tuning.md: ########## @@ -181,6 +181,58 @@ Missing or inaccurate statistics will hinder Spark's ability to select an optima - **Query plan estimates**: You can inspect Spark's cost estimates in the optimized query plan via [`EXPLAIN COST`](sql-ref-syntax-qry-explain.html) or `DataFrame.explain(mode="cost")`. - **Runtime statistics**: You can inspect these statistics in the [SQL UI](web-ui.html#sql-tab) under the "Details" section as a query is running. Look for `Statistics(..., isRuntime=true)` in the plan. +## Optimizing the Aggregate + +### Adaptive Partial Aggregation + +A grouping aggregation normally runs in two phases: a partial aggregation before the shuffle and a +final aggregation after it. The partial aggregation is only worthwhile when it actually reduces the +number of rows; when the grouping keys are close to unique it maintains -- and possibly spills -- an +aggregation map roughly as large as its input while emitting almost as many rows as it consumed. + +When adaptive partial aggregation is enabled, hash aggregation measures the compaction ratio (the +number of processed rows divided by the number of keys held in its aggregation maps) at runtime +and, if the partial aggregation is not collapsing enough rows to be worthwhile, stops populating +the aggregation map and passes the remaining rows through as single-row partial aggregation buffers +for the final aggregation to merge. Query results are unchanged. The ratio is evaluated Review Comment: **Finding 6.** "a query that only becomes ineffective later in its input is still caught" is true of the spill check but not of the periodic one, and the sentence attaches it to both. The periodic check compares the epoch's *cumulative* rows against the epoch's *cumulative* keys, so a favourable prefix keeps the ratio above the threshold indefinitely: 1M rows over 100 keys followed by 1M fully distinct rows gives `2000000 < 1000100 * 1.05`, i.e. false, at every check point after the turn — it never bypasses, even though the marginal ratio over the last 100k rows is 1.0. Only a spill (which resets `processedRows` and empties the map) makes a late turn visible, which is what `TungstenAggregationIterator.scala:196-199` correctly says. Either narrow the claim here to the spill check, or measure the ratio marginally — `minRows / (keysNow - keysAtLastCheck)`, with the key count snapshotted at each check the way `adaptiveFastKeysAtSpill` already snapshots at each spill. The marginal form is what makes "caught later" actually true, and it makes the periodic check independent of how long the epoch has been running. Note it pulls in the opposite direction from finding 2 (it decides on the most recent window rather than everything seen), so the growing check interval there is worth pairing with it. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala: ########## @@ -354,6 +410,47 @@ class TungstenAggregationIterator( } } + /////////////////////////////////////////////////////////////////////////// + // Part 5b: Methods and fields used by adaptive partial aggregation pass-through. + /////////////////////////////////////////////////////////////////////////// + + // Indicates that partial aggregation has been bypassed and the remaining input rows should be + // passed through as single-row partial buffers. Set in `processInputs` by either check point. + // It may coexist with earlier spills, so the output order is: sort-based (or map) output first, + // then the pass-through rows. + private[this] var passThrough: Boolean = false + + // The row that could not be inserted at the spill check. It is stashed here (as + // a copy) so it becomes the first pass-through row rather than being lost. + private[this] var pendingPassThroughRow: InternalRow = null + + // A reused aggregation buffer for building single-row partial buffers during pass-through. It is + // re-initialized from `initialAggregationBuffer` for every passed-through row. + private[this] lazy val passThroughAggregationBuffer: UnsafeRow = createNewAggregationBuffer() + + // Whether there are remaining pass-through rows to emit. + private def passThroughHasNext: Boolean = + passThrough && (pendingPassThroughRow != null || inputIter.hasNext) + + // Emits the next input row as a single-row partial aggregation buffer, i.e. a group of size one. + // The output (grouping key ++ buffer) is a valid partial buffer that the downstream Final + // aggregation merges, so the result is identical to running partial aggregation on this row. + private def nextPassThroughOutput(): UnsafeRow = { Review Comment: **Finding 7.** @viirya's `aggTime` point is fixed for codegen (`beforeResumedAgg` in `doProduceWithKeys`) but the interpreted path still has the gap, and pass-through made it worse than it was. `HashAggregateExec.doExecute` measures ```scala val beforeAgg = System.nanoTime() ...new TungstenAggregationIterator(...) aggTime += NANOSECONDS.toMillis(System.nanoTime() - beforeAgg) ``` which used to cover everything, because `processInputs` ran to completion in the constructor. Now, once `passThrough` is set, the loop exits early and the remaining input is pulled row by row from `nextPassThroughOutput()` during `next()` — after `aggTime` has already been recorded. For a high-cardinality task that bypasses at the first check point that's ~all of the build work missing from the SQL UI, and it's exactly the case the feature targets, so the metric will understate the very workload people will be looking at. Accumulating here would cover it: ```scala private def nextPassThroughOutput(): UnsafeRow = { val startTime = System.nanoTime() ... val out = generateOutput(groupingKey, passThroughAggregationBuffer) aggTime += NANOSECONDS.toMillis(System.nanoTime() - startTime) out } ``` (needs `aggTime` threaded in alongside `numBypassingRows`; a per-row `nanoTime` may be too coarse, in which case timing the whole drain from `hasNext` returning false would do.) ########## sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala: ########## @@ -141,6 +148,40 @@ case class HashAggregateExec( .map(_.asInstanceOf[DeclarativeAggregate]) private val bufferSchema = DataTypeUtils.fromAttributes(aggregateBufferAttributes) + /** + * Whether adaptive partial aggregation applies to this operator. When it does, the aggregation + * may bypass partial aggregation at runtime and pass the remaining input rows through as + * single-row partial buffers (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). It only + * applies to a pre-shuffle partial aggregation with grouping keys: + * - `Partial` mode only: the downstream `Final` aggregation merges the passed-through + * single-row buffers, so the output contract is unchanged. `Final`/`Complete` produce the + * result themselves and have no such downstream. `PartialMerge` does have one and could be + * supported by passing its incoming buffer through unchanged, but that is left for later. + * - grouping keys present: a global aggregation produces a single output row, so partial + * aggregation achieves the maximum reduction and must never be bypassed. + * - DISTINCT aggregate functions are allowed: the intermediate `PartialMerge` phase of the + * multi-phase distinct plan is not `Partial` mode (and requires a distribution), so it always + * aggregates and de-duplicates, and the passed-through rows from the distinct `Partial` phase + * therefore carry exactly one distinct value each. + */ + private val adaptivePartialAggEnabled: Boolean = { + conf.adaptivePartialAggregationEnabled && + groupingExpressions.nonEmpty && + // Only the pre-shuffle partial aggregation has a downstream `Final` to merge passed-through + // single-row buffers. `requiredChildDistributionExpressions` is `None` exactly for that + // pre-shuffle phase and `Some` for the `Final`/`Complete` phase. This check is what keeps a + // group-by-only aggregate (no aggregate functions, so an empty `aggregateExpressions`) from + // being admitted vacuously: `aggregateExpressions.forall(_.mode == Partial)` alone is true + // for the empty list, which would wrongly make the `Final` phase eligible as well. + requiredChildDistributionExpressions.isEmpty && + aggregateExpressions.forall(a => a.mode == Partial) Review Comment: **Finding 9.** This gate also admits plans I don't think the PR intends to cover yet, and nothing in the suite reaches them: - `AggUtils.planStreamingAggregation`'s `partialAggregate` (`AggUtils.scala:490-500`) calls `createStreamingAggregate` with no `requiredChildDistributionExpressions` and all-`Partial` modes, so **every structured-streaming grouping aggregation** now gets runtime pass-through by default. - Same for `planStreamingAggregationForSession`'s `partialAggregate`, and for the batch `session_window` path, where `mayAppendMergingSessionExec` puts a `MergingSessionsExec` directly above the `None`/all-`Partial` partial aggregate. I traced both and I *believe* they're semantically fine — the downstream `PartialMerge`/`MergingSessionsExec` merges buffers, and `MergingSessionsExec.requiredChildOrdering` is non-empty so `EnsureRequirements` still inserts the sort it needs. But note the static sibling deliberately opts out of session windows (`AggUtils.scala:144-148`: "The bypass is skipped when a session_window grouping key is present"), and this one silently opts in, with more rows crossing the shuffle into the state store path. Either add a streaming case (a `MemoryStream` aggregation with a small `minRows`, asserting the output matches the feature-off run) plus a `session_window` case, or gate on `!isStreaming` and no `SessionWindow.marker` for now and leave a TODO. Given the feature is on by default I'd rather not have streaming enter it untested. -- 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]
