peter-toth commented on code in PR #57742: URL: https://github.com/apache/spark/pull/57742#discussion_r3735865415
########## sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala: ########## @@ -0,0 +1,1015 @@ +/* + * 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 { Review Comment: **Finding 12.** This matrix varies codegen, the fast map and a forced spill, but never the one thing that decides where the partial aggregate's output goes: the number of input partitions. Every `spark.range` in the suite passes `1` except the two in `results unchanged under a fused Union` (`:555`), and with a single input partition `SinglePartition` already satisfies the Final's `ClusteredDistribution`, so `EnsureRequirements` inserts no `Exchange` and the two aggregates fuse into one `WholeStageCodegen`. Same query, `87937dbd6d0`, only the split count differing: ``` splits=1 splits=2 *(1) HashAggregate(functions=[first(v)]) *(2) HashAggregate(functions=[first(v)]) +- *(1) HashAggregate(partial_first(v)) +- AQEShuffleRead coalesced +- *(1) Project +- ShuffleQueryStage 0 +- *(1) Range (0, 40, splits=1) +- Exchange hashpartitioning(k, 5) +- *(1) HashAggregate(partial_first(v)) +- *(1) Project +- *(1) Range (..., splits=2) ``` That one fact removes two of the three mechanisms this PR adds, in 43 of the 44 tests: - **The output buffer never fills.** Fused, the partial aggregate's `outputFunc` feeds the Final's `doConsume` — a hash-map insert. Nothing reaches `BufferedRowIterator.currentRows`, so `shouldStop()` stays false for the whole build: `needStopCheck` bounds nothing, and bypassed rows land in the Final's spillable map rather than on the heap. - **`adaptiveResumeBuild` never runs.** With `shouldStop()` false, `doAgg` consumes the partition in one call and `postChildProduce` sets `childrenConsumed`, so `outputMap()` is only ever reached via `adaptiveFinalOutput` — after every pass-through row. The "output them early ... releasing the memory before the remaining input is streamed" path (`HashAggregateExec.scala:736-739`) is unexercised, and in this shape the map is in fact held to the end. You already hit this empirically from the other side: the fused-`Union` test needed multi-partition inputs before a streamed row could fill the output buffer at all. It generalises — the 8-cell verification behind finding 1 and the `currentRows` probe on finding 10's thread both only ever ran the fused shape. Minimum fix is a couple of dedicated multi-partition tests, one of them order-sensitive. The thorough fix folds the partition count into this matrix — `build: Int => DataFrame` and ```scala for { inputPartitions <- Seq(1, 2) wholeStage <- Seq(true, false) twoLevelMap <- Seq(true, false) forceSpill <- Seq(true, false) } { ``` with each caller threading it into its `spark.range(..., 1, inputPartitions)`, and the reference run built the same way so it stays comparable. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala: ########## @@ -0,0 +1,1015 @@ +/* + * 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) + } { + val spillConf = if (forceSpill) { + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> forceSpillFallback) + } else { + Nil + } + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + // Small `minRows` so the periodic check runs on modest inputs. + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "8") ++ + spillConf ++ fixedPlanConfs): _*) { + val msg = s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap forceSpill=$forceSpill" + withClue(msg) { + checkAnswer(build(), reference) + } + } + } + } + + /** + * The observable per-run counters we assert on, all read from the partial `HashAggregateExec` in + * a single execution so the metrics are not double-counted: + * - `skipped`: our self-reported `numBypassingRows` metric. + * - `partialOutputRows`: the partial aggregate's own `numOutputRows`. An independent, + * pre-existing counter driven by the normal output path, so it is the ground truth for + * whether rows were streamed through -- it equals the distinct key count when aggregation is + * effective and climbs toward the input row count once the operator bypasses. + * - `spillBytes`: the partial aggregate's `spillSize`. Reliable only when no fallback is + * forced: on the interpreted path this is derived from the task-cumulative memory-spill + * counter, so a forced fallback (or downstream shuffle-write spill) can inflate it. + * Asserted only by the periodic check test, which forces no fallback; use + * `tasksFallBacked` otherwise. + * - `tasksFallBacked`: the partial aggregate's `numTasksFallBacked`, incremented only when the + * regular map actually falls back into sort-based aggregation. When the spill check bypasses + * at the spill boundary the sorter is never created, so this stays 0 -- direct, per-operator + * evidence the bypass replaced the sort fallback. + */ + private case class AggCounters( + skipped: Long, + partialOutputRows: Long, + spillBytes: Long, + tasksFallBacked: Long) + + // Verifies `df` (an already-collected bypassing run) produces the same results as the feature-off + // reference. `build` is re-run for the reference so it gets a genuinely non-adaptive plan rather + // than reusing the bypassing run's cached one. + private def checkAgainstReference(df: DataFrame, build: () => DataFrame): Unit = { + val reference = withSQLConf( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + build().collect().toSeq + } + checkAnswer(df, reference) + } + + private def runAndReadCounters(build: () => DataFrame): AggCounters = { + // The triggering tests assert on metrics, so also verify the bypassing run produces the same + // results as the feature-off reference. + val df = build() + df.collect() + val partialAggs = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => agg + } + // A partial aggregate is always present for the grouped queries these tests use. + assert(partialAggs.nonEmpty, "expected a partial HashAggregateExec in the plan") + val counters = AggCounters( + // The metric is only registered on aggregates the feature applies to; an aggregate without + // it bypassed nothing. + skipped = partialAggs.map(_.metrics.get("numBypassingRows").map(_.value).getOrElse(0L)).sum, + partialOutputRows = partialAggs.map(_.metrics("numOutputRows").value).sum, + spillBytes = partialAggs.map(_.metrics("spillSize").value).sum, + tasksFallBacked = partialAggs.map(_.metrics("numTasksFallBacked").value).sum) + checkAgainstReference(df, build) + counters + } + + private def numBypassingRows(build: () => DataFrame): Long = runAndReadCounters(build).skipped + + // Returns the bypassed-row count per Partial-mode `HashAggregateExec`, keyed by the number of + // grouping keys, and verifies the run matches the feature-off reference. A `count(DISTINCT ...)` + // group-by has two such Partial phases -- the de-duplication partial (grouping on key + distinct + // columns) and the distinct partial (grouping on the keys only) -- so their bypasses can be told + // apart by the grouping key count. + private def bypassRowsByGroupingKeyCount(build: () => DataFrame): Map[Int, Long] = { + val df = build() + df.collect() + val byKeyCount = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => + agg.groupingExpressions.length -> + agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L) + }.groupBy(_._1).map { case (n, pairs) => n -> pairs.map(_._2).sum } + checkAgainstReference(df, build) + byKeyCount + } + + /** + * Runs `body` once per (wholeStage, twoLevelMap) combination with the feature enabled and a small + * `minRows`, threading a descriptive clue for failure messages. + * + * The fast (first-level) map is append-only and never spills, so only the regular (second-level) + * map can reach a spill boundary. With the default fast-map capacity (2^16) a small + * high-cardinality input would be fully absorbed by the fast map and never reach the regular map, + * so nothing could ever bypass. To make the triggering tests meaningful when the two-level map is + * on, we shrink the fast map via the first field of `testFallbackStartsAt` so rows fall through + * to the regular map. `regularFallback` optionally sets the second field to also force the + * regular map to spill (for the spill check); when 0 the regular map does not spill. + */ + private def forEachCodegenAndMap( + minRows: Long = 8, + regularFallback: Int = 0, + minCompaction: Double = -1.0)( + body: String => Unit): Unit = { + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + // Shrink the fast map to 4 keys when it is on so rows reach the regular map. The second field + // controls regular-map spilling; 0 means "never" (a large sentinel). + val fallbackConf = if (twoLevelMap || regularFallback > 0) { + val fastCap = if (twoLevelMap) 4 else 1 + val regular = if (regularFallback > 0) regularFallback else Int.MaxValue + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> s"$fastCap, $regular") + } else { + Nil + } + // A negative value means "leave the threshold at its default". + val thresholdConf = if (minCompaction >= 0.0) { + Seq(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key -> minCompaction.toString) + } else { + Nil + } + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> minRows.toString) ++ + fallbackConf ++ thresholdConf ++ fixedPlanConfs): _*) { + body(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") + } + } + } + + ///////////////////////////////////////////////////////////////////////////// + // Part 1: Correctness -- results identical to the feature-off reference. + ///////////////////////////////////////////////////////////////////////////// + + test("results unchanged for high-cardinality input that bypasses partial aggregation") { + // Every grouping key is distinct, so partial aggregation reduces nothing and should be + // bypassed by the periodic check. + checkAdaptiveMatchesReference { () => + spark.range(0, 200, 1, 1) + .select($"id" as "k", ($"id" * 2) as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", max($"v") as "m") + } + } + + test("results unchanged for low-cardinality input that keeps partial aggregation") { + // Few distinct keys, high reduction: partial aggregation is effective and should be kept. + checkAdaptiveMatchesReference { () => + spark.range(0, 600, 1, 1) + .select(($"id" % 5) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", min($"v") as "mn", max($"v") as "mx") + } + } + + test("results unchanged for medium-cardinality input near the reduction threshold") { + // Roughly half the rows are distinct keys; exercises the boundary of the ratio checks. + checkAdaptiveMatchesReference { () => + spark.range(0, 1000, 1, 1) + .select(($"id" % 500) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with multiple grouping keys and string keys") { + checkAdaptiveMatchesReference { () => + spark.range(0, 500, 1, 1) + .select( + concat(lit("g"), ($"id" % 300).cast("string")) as "k1", + ($"id" % 7) as "k2", + $"id" as "v") + .groupBy($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with nullable grouping keys") { + // Nulls are sparse enough (1 in 40) that the keys stay close to unique and the input really + // does bypass; a denser null key would lift the compaction ratio above the threshold and the + // test would never engage the feature. + checkAdaptiveMatchesReference { () => + spark.range(0, 400, 1, 1) + .select( + when($"id" % 40 === 0, lit(null)).otherwise($"id") as "k", + $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with average (multi-slot buffer) aggregate") { + // avg has a two-slot partial buffer (sum, count); pass-through buffers must carry all slots. + checkAdaptiveMatchesReference { () => + spark.range(0, 300, 1, 1) + .select($"id" as "k", ($"id" + 1) as "v") + .groupBy($"k") + .agg(avg($"v") as "a", sum($"v") as "s") + } + } + + test("results unchanged with a mix of many aggregate functions and buffer types") { + // Exercises a wide pass-through buffer spanning several aggregate buffer layouts at once: + // sum (decimal), avg (double), count, min/max, first/last, and stddev (imperative buffer). + checkAdaptiveMatchesReference { () => + spark.range(0, 400, 1, 1) + .select( + $"id" as "k", + ($"id" % 97).cast("decimal(10,2)") as "d", + ($"id" % 13).cast("double") as "dbl") + .groupBy($"k") + .agg( + sum($"d") as "sd", + avg($"dbl") as "ad", + count(lit(1)) as "c", + min($"dbl") as "mn", + max($"dbl") as "mx", + first($"dbl") as "f", + last($"dbl") as "l", + stddev($"dbl") as "sd2") + } + } + + test("results unchanged with filtered aggregate functions") { + // A `FILTER (WHERE ...)` aggregate is compiled into a per-row guard around the buffer update + // rather than a separate filtering operator: `If(filter, update, buffer)` in the interpreted + // path and an `if (!cond) continue` guard in the generated code. Pass-through reuses those + // exact update expressions, so a bypassed row whose filter is false contributes nothing to its + // single-row buffer. The all-true and all-false filters pin the two extremes, and the fully + // distinct grouping keys ensure rows bypass (in the regular-map-only configurations) so the + // filter guard actually runs in the pass-through path. + withTempView("t") { + spark.range(0, 400, 1, 1) + .select($"id" as "k", ($"id" % 100) as "v") + .createOrReplaceTempView("t") + checkAdaptiveMatchesReference { () => + spark.sql( + """SELECT k, + | sum(v) FILTER (WHERE v % 2 = 0) AS s_even, + | count(1) FILTER (WHERE v > 50) AS c_gt50, + | avg(v) FILTER (WHERE v > 25) AS a_gt25, + | sum(v) FILTER (WHERE true) AS s_all, + | sum(v) FILTER (WHERE false) AS s_none + |FROM t GROUP BY k""".stripMargin) + } + } + } + + test("results unchanged with decimal and date grouping keys") { + checkAdaptiveMatchesReference { () => + spark.range(0, 300, 1, 1) + .select( + ($"id" % 280).cast("decimal(12,3)") as "k1", + date_add(lit(java.sql.Date.valueOf("2020-01-01")), ($"id" % 250).cast("int")) as "k2", + $"id" as "v") + .groupBy($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged for group-by-only (distinct) with no aggregate functions") { + // No aggregate functions: the pass-through buffer is a zero-column UnsafeRow, so the output is + // just the grouping key. High-cardinality keys should bypass, and the de-duplicated result must + // still match the reference. + checkAdaptiveMatchesReference { () => + spark.range(0, 400, 1, 1) + .select(($"id" % 350) as "k1", ($"id" % 11) as "k2") + .distinct() + } + } + + test("results unchanged for group-by-only with duplicate keys (Final phase must not bypass)") { + // A group-by-only aggregate has an empty `aggregateExpressions`, so checking the aggregate + // modes alone is vacuously true and could wrongly admit the `Final` phase of the two-phase + // plan. With duplicate keys, a bypassing `Final` would skip its de-duplication and return + // duplicate rows. The two-level map off variants route the rows to the regular map so the + // periodic check fires and the regression would show up. + checkAdaptiveMatchesReference { () => + spark.range(0, 1000, 1, 1) + .select(($"id" % 10) as "c") + .distinct() + } + } + + test("results unchanged when a large frozen map is output before pass-through streaming") { + // A larger `minRows` lets the map accumulate many keys before the periodic check bypasses, + // so the early map output (which also frees the map) spans several drain cycles and re-enters + // the map-output function; the results must still match the feature-off reference. + val query = () => spark.range(0, 400000, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "200000", + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> "false") ++ fixedPlanConfs): _*) { + val reference = withSQLConf( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + query().collect().toSeq + } + checkAnswer(query(), reference) + } + } + + test("distinct aggregation stays correct") { + checkAdaptiveMatchesReference { () => + spark.range(0, 300, 1, 1) + .select($"id" as "k", ($"id" % 50) as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd", sum($"v") as "s") + } + } + + test("distinct aggregation bypasses on high-cardinality input") { + // The `PartialMerge` phase of the multi-phase distinct plan always aggregates (it is not + // `Partial` mode and requires a distribution), so the rows reaching the distinct `Partial` + // phase are de-duplicated and pass-through carries exactly one distinct value each. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 1000, 1, 1) + .select(($"id" % 100) as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + assert(numBypassingRows(df) > 0, + "expected a distinct partial aggregation to bypass for high-cardinality input") + } + } + } + + test("count distinct: the de-duplication partial aggregate bypasses") { + // `count(DISTINCT v) GROUP BY k` plans two `Partial` phases: the de-duplication partial groups + // on (k, v) and the distinct partial groups on (k). Fully distinct (k, v) pairs make the + // de-duplication partial (2 grouping keys) reduce nothing, so it must bypass. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select(($"id" % 4) as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the (k, v) de-duplication partial to bypass, got $byKeyCount") + } + } + } + + test("count distinct: the distinct partial aggregate bypasses") { + // Mirror of the test above for the other phase: with many distinct keys but few distinct + // values per key, the (k, v) de-duplication partial reduces well while the distinct partial + // (1 grouping key) sees a fresh key per row and must bypass. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select($"id" as "k", ($"id" % 2) as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (grouping on k) to bypass, got $byKeyCount") + } + } + } + + test("count distinct: both partial aggregates bypass and results stay correct") { + // Fully distinct keys and fully distinct values: neither partial phase reduces anything, so + // both bypass in the same execution. The de-duplication partial keeps the (k, v) pairs unique + // and the distinct partial counts them, so the result must still match the reference. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the (k, v) de-duplication partial to bypass, got $byKeyCount") + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (grouping on k) to bypass, got $byKeyCount") + } + } + } + + test("global aggregation (no grouping keys) is never bypassed and stays correct") { + withSQLConf(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true") { + checkAnswer( + spark.range(0, 100, 1, 1).agg(sum($"id") as "s", count(lit(1)) as "c"), + Row(4950L, 100L)) + } + } + + test("results unchanged with an empty input") { + checkAdaptiveMatchesReference { () => + spark.range(0, 0, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + // The following four tests cover plans where an `ExpandExec` sits below the partial aggregate + // (ROLLUP / CUBE / GROUPING SETS / multi-distinct). PR apache/spark#28804 statically disabled its + // skip-partial-aggregate optimization whenever an Expand was present, but that was a performance + // heuristic guarding its *static* row sampling, not a correctness requirement. Our decision is + // made at runtime from the observed compaction ratio, so we deliberately do not port that + // exclusion. These tests assert results stay correct with the exclusion absent. + // + // The ROLLUP and CUBE cases below use two grouping columns, where the grand-total set keeps the + // compaction ratio high enough that they decline to bypass -- they cover the eligible-but- + // declining side. (Widening the rollup lowers the ratio: with five distinct columns the same + // shape does bypass.) The GROUPING SETS and multi-distinct tests, and `pass-through fires for + // high-cardinality input below an Expand`, cover an Expand that bypasses. + + test("results unchanged for ROLLUP (Expand below partial aggregate)") { + checkAdaptiveMatchesReference { () => + spark.range(0, 400, 1, 1) + .select(($"id" % 200) as "k1", ($"id" % 7) as "k2", $"id" as "v") + .rollup($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged for CUBE (Expand below partial aggregate)") { + checkAdaptiveMatchesReference { () => + spark.range(0, 400, 1, 1) + .select(($"id" % 150) as "k1", ($"id" % 5) as "k2", $"id" as "v") + .cube($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged for GROUPING SETS (Expand below partial aggregate)") { + // No `()` grouping set, and both keys distinct, so every expanded row is a fresh key and the + // input genuinely bypasses. A grand-total set would collapse all rows into one group and lift + // the compaction ratio above the threshold (see the ROLLUP and CUBE tests below). + withTempView("t") { + spark.range(0, 400, 1, 1) + .select($"id" as "k1", ($"id" + 1000) as "k2", $"id" as "v") + .createOrReplaceTempView("t") + checkAdaptiveMatchesReference { () => + spark.sql( + """SELECT k1, k2, sum(v) AS s, count(1) AS c + |FROM t + |GROUP BY k1, k2 GROUPING SETS ((k1, k2), (k1), (k2))""".stripMargin) + } + } + } + + test("results unchanged for multi-distinct (Expand below partial aggregate)") { + checkAdaptiveMatchesReference { () => + spark.range(0, 400, 1, 1) + .select(($"id" % 100) as "k", ($"id" % 30) as "a", ($"id" % 40) as "b") + .groupBy($"k") + .agg(countDistinct($"a") as "da", countDistinct($"b") as "db", sum($"a") as "s") + } + } + + test("results unchanged under a fused Union (child yields from a nested helper)") { + // `UnionExec` wraps each child's produce in its own helper, so a streamed row that fills the + // output buffer returns only as far as the aggregate's build loop. Reaching the end of the + // child's produce therefore does not mean the input is exhausted, and treating it as such + // drops the rest of the partition. + checkAdaptiveMatchesReference { () => + spark.range(0, 100, 1, 2).union(spark.range(100, 200, 1, 2)).groupBy("id").count() + } + } + + test("both execution paths agree on an order-sensitive aggregate") { + // Bypassing emits a group's streamed rows ahead of the ones its map still holds, so + // `first`/`last` see a different merge order than an unbypassed run. Whether that order is + // guaranteed is a separate question -- what must hold is that the generated and interpreted + // paths agree, or `spark.sql.codegen.wholeStage` would be observable in the result. + // + // Key -1 takes the first and last input rows; the rest are distinct, so the bypass fires early + // and leaves row 0 in the map while row 8 streams past it. + val query = () => spark.range(0, 9, 1, 1) + .select( + when($"id" === 0 || $"id" === 8, lit(-1L)).otherwise($"id") as "k", + $"id" as "v") + .groupBy($"k") + .agg(first($"v") as "f", last($"v") as "l") + + val results = for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + forceSpill <- Seq(true, false) + } yield { + val spillConf = if (forceSpill) { + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> forceSpillFallback) + } else { + Nil + } + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "2") ++ + spillConf ++ fixedPlanConfs): _*) { + query().collect().toSeq.sortBy(_.getLong(0)) + } + } + assert(results.distinct.length == 1, Review Comment: **Finding 13.** This asserts the eight cells agree with each other, but not that they agree with a non-bypassed run — and they do not. Measured on `87937dbd6d0` for exactly this query, reading the `k = -1` row: | run | `f` | `l` | | --- | --- | --- | | `ADAPTIVE_PARTIAL_AGGREGATION_ENABLED=false` | **0** | **8** | | all 8 cells with it on | **8** | **0** | So the value the test pins is the inverted one: `first` returns the last input row of the group and `last` returns the first. The suite's own reference-comparison helper would catch it — swap the assert for ```scala checkAdaptiveMatchesReference(query) ``` and it fails at this head (`minRows` would need to come from the helper, or stay a local `withSQLConf` around it). That is the assertion this test wants: agreeing with each other is necessary but the reference is what says the order is right. Second gap: because the input is `spark.range(0, 9, 1, 1)` there is no `Exchange` here (finding 12), so the eight cells are eight variations of the fused shape. The divergence that is still live needs an `Exchange` plus a colliding row that is not the *first* bypassed row — `dupAt = 9` rather than `8` — which is why this test passes while the paths still disagree. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala: ########## @@ -191,29 +209,67 @@ class TungstenAggregationIterator( } } else { var i = 0 - while (inputIter.hasNext) { + var processedRows = 0L + val minRows = adaptiveMinRows + // The processed-row count at which the compaction ratio is evaluated next. It advances by + // `minRows` after every check, and restarts after a spill so the new in-memory map epoch is + // judged on its own rows. `minRows = 0` disables the periodic check: the count is only ever + // compared after being incremented past 0, so it never matches and only the spill check + // below remains. + var nextCheckRow = minRows + // The partial aggregation is ineffective when it does not collapse `minCompaction` rows into + // one key. There is no fast map on this path, so the map's keys are all the operator holds. + def ineffective(): Boolean = + processedRows < hashMap.getNumKeys().toDouble * adaptiveMinCompaction + while (inputIter.hasNext && !passThrough) { Review Comment: **Finding 15.** An alternative to freezing the map, offered because it settles finding 1's merge-order question outright rather than picking a side — you asked for a preferred shape on finding 2's thread and this is mine, with the honest cost below. Today the flip means *stop aggregating*: the loop exits here, the map is frozen and freed, and every remaining row streams. The variant is that the flip means *stop growing the map* — keep looping to exhaustion and, per row, look the key up **without inserting**: ```scala while (inputIter.hasNext) { val newInput = inputIter.next() val groupingKey = groupingProjection.apply(newInput) val buffer = if (!passThrough) { hashMap.getAggregationBufferFromUnsafeRow(groupingKey) // may insert } else { hashMap.getAggregationBufferIfPresent(groupingKey) // never inserts } if (buffer != null) processRow(buffer, newInput) else emitSingleRowBuffer(newInput) ... } ``` `BytesToBytesMap.lookup` is already lookup-only, so `UnsafeFixedWidthAggregationMap` needs a thin `getAggregationBufferIfPresent`; `RowBasedHashMapGenerator` and `VectorizedHashMapGenerator` both already generate a find-or-insert probe, so a find-only variant is a small addition to each. What it buys: - **Merge order is preserved by construction — finding 1 stops being a question.** A group is either wholly in the map (one buffer) or wholly streamed (single-row buffers, in input order), never both. There is then no order to choose between the two execution paths, in either plan shape, and no `first`/`last` change against a non-bypassed run. Findings 13 and 14 go with it. - **Finding 6's sentence becomes true as written.** An input that turns high-cardinality late is handled because its new keys stream; the periodic check no longer has to notice the turn for the operator to behave. - **Finding 2 is bounded rather than fixed.** A wrong decision costs only the keys that are genuinely *new* after it. On the input in that thread the 100 hot tail keys are already resident when the flip fires, so they keep aggregating and the partial output goes back to ~100k rows instead of 500,000. But invert it — 100k distinct keys, then 400k rows over 100 *brand-new* keys — and all 400k still stream. So this narrows the exposure, it does not remove it, and finding 2 still needs its own answer. The honest cost, against your own benchmark: | per row after the flip | feature off | current bypass | this variant | | --- | --- | --- | --- | | hash + probe | yes, into a table that keeps growing | **no** | yes, into a *frozen* table | | insert / rehash / growth | yes | no | no | | spill + sort-merge fallback | yes, once the map fills | no | no | | map held | to the end of the task | freed early | to the end of the task | | rows crossing the shuffle | #distinct keys | #rows after the flip | #rows whose key is absent | - **It gives up the probe-skip, which is not a rounding error.** `AdaptivePartialAggregationBenchmark-results.txt` shows high-cardinality codegen at 473.8 -> 283.8 ns/row **with no spill at all** (1.7x), against 991.7 -> 527.0 ns/row when spilling. So ~275 ns/row of the win is spill avoidance and ~190 ns/row is per-row map work that today is skipped entirely. This variant recovers the spill part in full and only part of the map-work part: it drops the insert, the rehash and the growth, and probes stay cheap because the table stops growing, but the probe itself stays. What fraction survives I cannot get from these numbers — re-running the existing benchmark with it is what decides whether it is worth the trade. - **The early map free goes away.** Peak memory becomes the map's frozen size for the rest of the task. (In the fused shape that is already what happens today — finding 12 — so it is only a regression where an `Exchange` exists.) - **More surface:** three map implementations need the find-only accessor, and `numBypassingRows` changes meaning (rows whose key was absent, not rows after the flip). Not asking for it in this PR, and if the benchmark says the probe-skip is carrying most of the win then it is the wrong trade and finding 1 has to be settled directly instead. I am raising it because it is the only shape I found where the order question does not have to be answered at all. ########## 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. Review Comment: **Finding 14.** This still describes the pre-`87937dbd6d0` order and now contradicts `next()` at `:504-507`, which says the opposite. Worth fixing whichever way finding 1 lands, since this is the comment a later reader will trust. ```suggestion // It may coexist with earlier spills, so the output order is: the pass-through rows first, then // the sort-based (or map) output. ``` -- 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]
