ulysses-you commented on code in PR #57742: URL: https://github.com/apache/spark/pull/57742#discussion_r3780377366
########## sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala: ########## @@ -0,0 +1,1256 @@ +/* + * 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. Once pass-through is active the map is frozen, and its + * output always precedes the passed-through rows: a row that collides with a frozen key is held + * behind the map and flushed only after it drains, so every group merges its buffers in the same + * order as a run that never bypasses, including order-sensitive aggregates such as `first`/`last`. + * + * The suite has two halves: + * 1. Correctness: aggregate results are 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). Order-sensitive aggregates are tested against the reference too, including + * under a fan-out child that queues its whole batch behind the frozen map. + * 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 `build` with adaptive partial aggregation disabled (the reference) and then across the + * full configuration matrix with it enabled, asserting every enabled run matches the reference. + * + * `build` takes the number of input partitions, which the matrix varies along with everything + * else, because the plan shape decides which parts of the feature run at all. When the two + * aggregates end up in one whole-stage -- no `Exchange` between them -- the partial aggregate's + * output feeds the Final's `doConsume` directly and never reaches + * `BufferedRowIterator.currentRows`, so `shouldStop()` stays false for the whole build and + * neither `needStopCheck` nor the resumed-build path is exercised. Splitting them puts the + * streamed rows through the output buffer and runs both. + * + * More than one input partition is necessary but not sufficient for that split: a `Range` keyed + * directly on `id` already reports an output partitioning that satisfies the Final aggregate's + * `ClusteredDistribution`, so `EnsureRequirements` inserts no `Exchange` however many partitions + * it has. Tests that want the split shape group on a derived key (a cast, say) so the input + * partitioning no longer satisfies the requirement. + * + * `expectBypass` ties the correctness guarantee to the triggering guarantee: beyond matching the + * reference, every cell must either actually stream rows through (when true) or keep + * aggregating (when false). Without it a test could silently stop exercising pass-through if the + * input stopped being bypassable, and only this assertion makes that fail loudly. + */ + private def checkAdaptiveMatchesReference( + build: Int => DataFrame, + expectBypass: Boolean = true): Unit = { + for { + inputPartitions <- Seq(1, 2) + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + forceSpill <- Seq(true, false) + } { + // The reference is built with the same partitioning, so only the feature differs. + val reference = withSQLConf( + (SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") +: fixedPlanConfs: _*) { + build(inputPartitions).collect().toSeq + } + 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"inputPartitions=$inputPartitions wholeStage=$wholeStage " + + s"twoLevelMap=$twoLevelMap forceSpill=$forceSpill" + withClue(msg) { + // Collect once so the metrics are populated, then check whether the bypass fired for + // this cell. The metric lives on the `Partial`-mode `HashAggregateExec`, so that is the + // operator the assertion reads. + val df = build(inputPartitions) + df.collect() + val skipped = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => + agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L) + }.sum + if (expectBypass) { + assert(skipped > 0, + s"expected rows to bypass partial aggregation, got $skipped bypassed rows") + } else { + assert(skipped == 0, + s"expected no rows to bypass partial aggregation, got $skipped bypassed rows") + } + checkAnswer(df, 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 { parts => + spark.range(0, 200, 1, parts) + .select($"id".cast("string") 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, so + // the bypass metric must stay zero in every cell. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 600, 1, parts) + .select(($"id" % 5).cast("string") 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. The + // overall compaction ratio (~2.0) is above the threshold, but the *first* periodic check still + // sees the leading distinct keys and fires, so the bypass must be observable too. + checkAdaptiveMatchesReference { parts => + spark.range(0, 1000, 1, parts) + .select(($"id" % 500).cast("string") 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 { parts => + spark.range(0, 500, 1, parts) + .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 { parts => + spark.range(0, 400, 1, parts) + .select( + when($"id" % 40 === 0, lit(null)).otherwise($"id").cast("string") 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 { parts => + spark.range(0, 300, 1, parts) + .select($"id".cast("string") 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 (declarative buffer). + // The imperative-buffer case is covered separately (see the `approx_count_distinct` test). + checkAdaptiveMatchesReference { parts => + spark.range(0, 400, 1, parts) + .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 an imperative-buffer aggregate") { + // `approx_count_distinct` uses `HyperLogLogPlusPlus`, an `ImperativeAggregate` whose buffer + // state is written by `initialize(buffer)` rather than by a projection, so a pass-through + // single-row buffer has to be re-initialized with `copyFrom(initialAggregationBuffer)` for + // every row. No declarative aggregate exercises that reset. It also reports + // `supportCodegen = false`, so the operator only ever runs on `TungstenAggregationIterator`; + // keep it in its own test rather than folding it into a codegen cell that would quietly + // become interpreted. + checkAdaptiveMatchesReference { parts => + spark.range(0, 300, 1, parts) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(approx_count_distinct($"v") as "c") + } + } + + 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".cast("string") as "k", ($"id" % 100) as "v") + .createOrReplaceTempView("t") + checkAdaptiveMatchesReference { parts => + 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 { parts => + spark.range(0, 300, 1, parts) + .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 { parts => + spark.range(0, 400, 1, parts) + .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 { parts => + spark.range(0, 1000, 1, parts) + .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".cast("string") 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 df = query() + df.collect() + val skipped = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => + agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L) + }.sum + assert(skipped > 0, + s"expected the large frozen map to eventually bypass, got $skipped bypassed rows") + val reference = withSQLConf( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + query().collect().toSeq + } + checkAnswer(df, reference) + } + } + + test("distinct aggregation stays correct") { + checkAdaptiveMatchesReference { parts => + spark.range(0, 300, 1, parts) + .select($"id".cast("string") 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).cast("string") 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).cast("string") 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".cast("string") 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".cast("string") 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") { + // No rows means the check points never fire; the metric must stay zero. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 0, 1, 1) + .select($"id".cast("string") 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)") { + // The grand-total group repeats on every expanded row, so the compaction ratio stays above the + // threshold and the partial aggregate declines to bypass in every cell. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 400, 1, parts) + .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)") { + // Same as ROLLUP: the grand-total group keeps the ratio above the threshold, so nothing + // bypasses. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 400, 1, parts) + .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 { parts => + 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 { parts => + spark.range(0, 400, 1, parts) + .select(($"id" % 100).cast("string") 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 { parts => + spark.range(0, 100, 1, parts + 1) + .union(spark.range(100, 200, 1, parts + 1)) + .groupBy("id").count() + } + } + + test("results unchanged when an Exchange separates the two aggregates") { + // Grouping on a derived key stops the `Range`'s output partitioning from satisfying the Final + // aggregate's `ClusteredDistribution`, so the plan keeps an `Exchange` and the two aggregates + // land in separate whole-stages. That is the shape where streamed rows actually pass through + // `BufferedRowIterator.currentRows` -- fused, they go straight into the Final's hash map and + // neither `needStopCheck` nor the resumed build is reached. + checkAdaptiveMatchesReference { parts => + spark.range(0, 200, 1, parts) + .select($"id".cast("string") as "k", ($"id" * 2) as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", max($"v") as "m") + } + } + + test("both execution paths agree on an order-sensitive aggregate") { Review Comment: addressed ########## sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala: ########## @@ -626,13 +858,37 @@ case class HashAggregateExec( | long $beforeAgg = System.nanoTime(); | $doAggFuncName(partitionIndex); | $aggTime.add((System.nanoTime() - $beforeAgg) / $NANOS_PER_MILLIS); + | $adaptiveStopCheck |} - |// output the result - |$outputFromFastHashMap - |$outputFromRegularHashMap + |$adaptiveResumeBuild + |$adaptiveFinalOutput """.stripMargin } + // Blocking operators normally suppress the child's `shouldStop()` check because they buffer all + // output. With adaptive partial aggregation, pass-through rows are appended to the output buffer + // while consuming child input, so the stop check is re-enabled to let the child yield between + // rows. + // + // This bounds the buffer only as far as the child honours it. Each passed-through row queues a + // copy and advances the frozen-map output by one row, and that one appended map row makes + // `shouldStop()` true, so a child that checks between rows never queues more than one row. A + // one-to-many child that does not check `shouldStop()` inside its fan-out (`GenerateExec` emits + // `for (index ...) { consume }` and `while (iterator.hasNext()) { consume }` with no check) + // appends every row produced from one input row before it can yield: the fan-out batch lands in + // `BufferedRowIterator.currentRows` (which is not spillable), and the map advances one row per + // queued row, so the buffer grows to roughly the batch width rather than `minRows`. + override def needStopCheck: Boolean = adaptivePartialAggEnabled + + // Blocking operators normally do not copy their result because every output row is drained (via + // `shouldStop()`) before the next one is produced. Adaptive pass-through breaks that assumption: + // when an `Expand` sits below, one input row fans out into several pass-through rows that are all + // appended in the same child loop iteration before any drain, and they all alias the single + // result `UnsafeRow`. Such children report `needCopyResult` themselves, so propagate their + // requirement rather than copying for every adaptive aggregate. + override def needCopyResult: Boolean = adaptivePartialAggEnabled && Review Comment: addressed -- 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]
