ulysses-you commented on code in PR #57742:
URL: https://github.com/apache/spark/pull/57742#discussion_r3746851652
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -845,29 +1158,56 @@ case class HashAggregateExec(
}
}
- val declareRowBuffer: String = if (isFastHashMapEnabled) {
- val fastRowType = if (isVectorizedHashMapEnabled) {
- classOf[MutableColumnarRow].getName
+ val declareRowBuffer: String = {
+ val declareBuffers = if (isFastHashMapEnabled) {
+ val fastRowType = if (isVectorizedHashMapEnabled) {
+ classOf[MutableColumnarRow].getName
+ } else {
+ "UnsafeRow"
+ }
+ s"""
+ |UnsafeRow $unsafeRowBuffer = null;
+ |$fastRowType $fastRowBuffer = null;
+ """.stripMargin
+ } else {
+ s"UnsafeRow $unsafeRowBuffer = null;"
+ }
+ val declareBypassed = if (adaptivePartialAggEnabled) {
+ s"boolean $adaptiveRowBypassedTerm = false;"
} else {
- "UnsafeRow"
+ ""
}
s"""
- |UnsafeRow $unsafeRowBuffer = null;
- |$fastRowType $fastRowBuffer = null;
+ |$declareBuffers
+ |$declareBypassed
""".stripMargin
- } else {
- s"UnsafeRow $unsafeRowBuffer = null;"
}
// We try to do hash map based in-memory aggregation first. If there is
not enough memory (the
// hash map will return null for new key), we spill the hash map to disk
to free memory, then
// continue to do in-memory aggregation and spilling until all the rows
had been processed.
// Finally, sort the spilled aggregate buffers by key, and merge them
together for same key.
+ //
+ // With adaptive partial aggregation, once pass-through is active
`updateRowInHashMap` fills the
+ // single-row buffer built above; we then emit `key ++ buffer` straight to
the parent so the row
+ // skips both the fast map and the regular map.
+ val emitPassThroughRow = if (adaptivePartialAggEnabled) {
+ val numBypassingRows = metricTerm(ctx, "numBypassingRows")
+ s"""
+ |if ($adaptiveRowBypassedTerm) {
+ | $numBypassingRows.add(1);
+ | $outputFunc(${unsafeRowKeyCode.value}, $unsafeRowBuffer);
Review Comment:
Addressed at this head — the order test now covers both plan shapes
(`splits` ∈ {1, 2}) and both collision positions (`dupAt` ∈ {8, 9}), and the
generated and interpreted paths agree in all 8 cells per combination
(`519cb60a5f7`, `388a2d76bc5`).
##########
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:
Addressed in 87937dbd6d0 — `checkAdaptiveMatchesReference` now asserts
`numBypassingRows > 0` on the `Partial` aggregate per cell, so a cell can no
longer silently stop exercising pass-through.
##########
sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt:
##########
@@ -0,0 +1,56 @@
+================================================================================================
+high-cardinality input, no-spill pass-through (Tier 1)
Review Comment:
Addressed — the benchmark results were regenerated (`b69bfc5ec99`) with the
current section titles and the `minCompaction = 1.05` default.
##########
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:
Addressed in 87937dbd6d0 — the comment no longer claims the distinct
`Partial` aggregate cannot exist; the `count distinct` test exercises it
deliberately.
##########
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:
Addressed in 87937dbd6d0 — the eligibility check now also requires
`!isStreaming`.
##########
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:
Addressed in 519cb60a5f7 — the matrix now folds `inputPartitions` ∈ {1, 2}
into every reference run, and the order-sensitive test varies `splits` and
`dupAt` directly.
##########
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:
Addressed in `519cb60a5f7` and `388a2d76bc5` — the order test covers both
plan shapes and pins path-agreement (codegen == interpreted), deliberately not
a non-bypassed reference, since the merge order is not a promise.
##########
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:
Addressed in 184fc813519 — the comment now matches `next()`: first
pass-through row, then the map (or sort-based) output, then the remaining rows.
--
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]
