diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index 3ba3a6c749ca7..a1c778c4c10e9 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -181,6 +181,67 @@ Missing or inaccurate statistics will hinder Spark's ability to select an optima - **Query plan estimates**: You can inspect Spark's cost estimates in the optimized query plan via [`EXPLAIN COST`](sql-ref-syntax-qry-explain.html) or `DataFrame.explain(mode="cost")`. - **Runtime statistics**: You can inspect these statistics in the [SQL UI](web-ui.html#sql-tab) under the "Details" section as a query is running. Look for `Statistics(..., isRuntime=true)` in the plan. +## Optimizing the Aggregate + +### Adaptive Partial Aggregation + +A grouping aggregation normally runs in two phases: a partial aggregation before the shuffle and a +final aggregation after it. The partial aggregation is only worthwhile when it actually reduces the +number of rows; when the grouping keys are close to unique it maintains -- and possibly spills -- an +aggregation map roughly as large as its input while emitting almost as many rows as it consumed. + +When adaptive partial aggregation is enabled, hash aggregation measures the reduction ratio (the +number of distinct grouping keys divided by the number of processed rows) at runtime and, if the +partial aggregation is not reducing rows enough to be worthwhile, stops populating the aggregation +map and passes the remaining rows through as single-row partial aggregation buffers for the final +aggregation to merge. Query results are unchanged. The ratio is evaluated periodically, and again +right before the aggregation map would spill, so a query that only becomes ineffective later in its +input is still caught. + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Property NameDefaultMeaningSince Version
spark.sql.execution.aggregate.adaptivePartialAggregation.enabledtrue + When true, hash aggregation adaptively bypasses the pre-shuffle partial aggregation at runtime + when it observes that the partial aggregation is not reducing the number of rows enough to be + worthwhile. This applies only to hash aggregation with grouping keys. + 4.3.0
spark.sql.execution.aggregate.adaptivePartialAggregation.sampleRows100000 + The number of input rows to sample before evaluating the reduction ratio. When the ratio is + below the threshold, the next evaluation happens after twice as many rows, so low-cardinality + input is re-checked only rarely. + 4.3.0
spark.sql.execution.aggregate.adaptivePartialAggregation.noSpillReductionRatioThreshold0.9 + The reduction ratio threshold applied while the aggregation map is still fully in memory. If + the ratio is at least this value the partial aggregation is bypassed. A larger value is more + conservative (keeps partial aggregation in more cases). + 4.3.0
spark.sql.execution.aggregate.adaptivePartialAggregation.spillReductionRatioThreshold(value of noSpillReductionRatioThreshold) + The reduction ratio threshold applied when the aggregation map is about to spill. Setting it + lower makes the bypass more likely once spilling is imminent, and 0 always bypasses instead of + spilling. + 4.3.0
+ ## Optimizing the Join Strategy ### Automatically Broadcasting Joins diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index beb8d5ee14581..6afafa87195e9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -4156,6 +4156,64 @@ object SQLConf { .booleanConf .createWithDefault(false) + val ADAPTIVE_PARTIAL_AGGREGATION_ENABLED = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.enabled") + .doc("When true, hash aggregation adaptively bypasses the pre-shuffle partial aggregation " + + "at runtime when it observes that the partial aggregation is not reducing the number of " + + "rows enough to be worthwhile. Once bypassed, the remaining input rows are passed " + + "through as single-row partial aggregation buffers for the final aggregation to merge, " + + "which avoids the cost of maintaining and spilling a large aggregation map with little " + + "reduction benefit. This applies only to hash aggregation with grouping keys.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(true) + + val ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_ROWS = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.sampleRows") + .doc("The number of input rows to sample before evaluating the reduction ratio for the " + + s"no-spill tier of adaptive partial aggregation (see " + + s"'${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}'). From this many rows on, if the ratio " + + "of distinct grouping keys to processed rows is at least " + + s"'spark.sql.execution.aggregate.adaptivePartialAggregation." + + "noSpillReductionRatioThreshold', partial aggregation is bypassed for the rest of the " + + "input. When the ratio is below the threshold, the next evaluation happens after twice " + + "as many rows, so low-cardinality input is re-checked only rarely.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .intConf + .checkValue(_ > 0, "The sample row count must be positive.") + .createWithDefault(100000) + + val ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation." + + "noSpillReductionRatioThreshold") + .doc("The reduction ratio threshold used by the no-spill tier of adaptive partial " + + s"aggregation (see '${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}'). The reduction ratio " + + "is the number of distinct grouping keys divided by the number of processed rows. After " + + s"sampling '${ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_ROWS.key}' rows without spilling, if " + + "the ratio is at least this value the partial aggregation is bypassed. A larger value " + + "is more conservative (keeps partial aggregation in more cases).") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .doubleConf + .checkValue(v => v >= 0.0 && v <= 1.0, "The reduction ratio threshold must be in [0.0, 1.0].") + .createWithDefault(0.9) + + val ADAPTIVE_PARTIAL_AGGREGATION_SPILL_REDUCTION_RATIO_THRESHOLD = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation." + + "spillReductionRatioThreshold") + .doc("The reduction ratio threshold used by the on-spill tier of adaptive partial " + + s"aggregation (see '${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}'). When the aggregation " + + "map is about to spill, if the ratio of distinct grouping keys to processed rows is at " + + "least this value the partial aggregation is bypassed for the rest of the input. It " + + s"falls back to '${ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD.key}' " + + "so that both tiers apply one policy by default. Setting it lower makes the bypass more " + + "likely once spilling is imminent, and 0 always bypasses instead of spilling.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .fallbackConf(ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD) + val JSON_GENERATOR_IGNORE_NULL_FIELDS = buildConf("spark.sql.jsonGenerator.ignoreNullFields") .doc("Whether to ignore null fields when generating JSON objects in JSON data source and " + @@ -8903,6 +8961,18 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def bypassPartialAggregation: Boolean = getConf(BYPASS_PARTIAL_AGGREGATION) + def adaptivePartialAggregationEnabled: Boolean = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_ENABLED) + + def adaptivePartialAggregationSampleRows: Int = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_ROWS) + + def adaptivePartialAggregationNoSpillReductionRatioThreshold: Double = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD) + + def adaptivePartialAggregationSpillReductionRatioThreshold: Double = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_SPILL_REDUCTION_RATIO_THRESHOLD) + def objectAggSortBasedFallbackThreshold: Int = getConf(OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD) def variableSubstituteEnabled: Boolean = getConf(VARIABLE_SUBSTITUTE_ENABLED) diff --git a/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk21-results.txt b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk21-results.txt new file mode 100644 index 0000000000000..e5dfdc5633a47 --- /dev/null +++ b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk21-results.txt @@ -0,0 +1,56 @@ +================================================================================================ +high-cardinality input, no-spill pass-through (Tier 1) +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +codegen = true, adaptive = F 4094 4126 45 2.0 488.0 1.0X +codegen = true, adaptive = T 2402 2424 32 3.5 286.3 1.7X +codegen = false, adaptive = F 4989 4994 6 1.7 594.8 0.8X +codegen = false, adaptive = T 3183 3193 13 2.6 379.5 1.3X + + +================================================================================================ +low-cardinality input, no-spill pass-through (Tier 1) +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 287 299 11 58.4 17.1 1.0X +codegen = true, adaptive = T 282 290 5 59.4 16.8 1.0X +codegen = false, adaptive = F 1328 1329 2 12.6 79.2 0.2X +codegen = false, adaptive = T 1342 1350 12 12.5 80.0 0.2X + + +================================================================================================ +high-cardinality input, on-spill pass-through (Tier 2) +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 8850 8911 87 0.9 1055.0 1.0X +codegen = true, adaptive = T 4448 4570 173 1.9 530.3 2.0X +codegen = false, adaptive = F 9261 9357 136 0.9 1104.0 1.0X +codegen = false, adaptive = T 5276 5355 112 1.6 629.0 1.7X + + +================================================================================================ +low-cardinality input, on-spill pass-through (Tier 2) +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 789 806 16 21.3 47.0 1.0X +codegen = true, adaptive = T 813 835 28 20.6 48.5 1.0X +codegen = false, adaptive = F 1351 1425 104 12.4 80.5 0.6X +codegen = false, adaptive = T 1346 1350 6 12.5 80.2 0.6X + + diff --git a/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk25-results.txt b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk25-results.txt new file mode 100644 index 0000000000000..57d9b2dff3d63 --- /dev/null +++ b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk25-results.txt @@ -0,0 +1,56 @@ +================================================================================================ +high-cardinality input, no-spill pass-through (Tier 1) +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +codegen = true, adaptive = F 4083 4087 6 2.1 486.7 1.0X +codegen = true, adaptive = T 2431 2443 17 3.5 289.8 1.7X +codegen = false, adaptive = F 4913 4924 14 1.7 585.7 0.8X +codegen = false, adaptive = T 3214 3220 9 2.6 383.1 1.3X + + +================================================================================================ +low-cardinality input, no-spill pass-through (Tier 1) +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 257 267 8 65.3 15.3 1.0X +codegen = true, adaptive = T 282 292 7 59.5 16.8 0.9X +codegen = false, adaptive = F 1290 1290 0 13.0 76.9 0.2X +codegen = false, adaptive = T 1298 1301 4 12.9 77.4 0.2X + + +================================================================================================ +high-cardinality input, on-spill pass-through (Tier 2) +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 7932 7986 77 1.1 945.5 1.0X +codegen = true, adaptive = T 4246 4354 152 2.0 506.2 1.9X +codegen = false, adaptive = F 9290 9386 136 0.9 1107.4 0.9X +codegen = false, adaptive = T 5244 5298 76 1.6 625.2 1.5X + + +================================================================================================ +low-cardinality input, on-spill pass-through (Tier 2) +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 764 774 17 22.0 45.5 1.0X +codegen = true, adaptive = T 786 794 8 21.3 46.9 1.0X +codegen = false, adaptive = F 1362 1362 0 12.3 81.2 0.6X +codegen = false, adaptive = T 1363 1368 7 12.3 81.2 0.6X + + diff --git a/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt new file mode 100644 index 0000000000000..d6c49962b5ea5 --- /dev/null +++ b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt @@ -0,0 +1,56 @@ +================================================================================================ +high-cardinality input, no-spill pass-through (Tier 1) +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +codegen = true, adaptive = F 3975 4060 120 2.1 473.8 1.0X +codegen = true, adaptive = T 2381 2404 32 3.5 283.8 1.7X +codegen = false, adaptive = F 4850 4854 5 1.7 578.2 0.8X +codegen = false, adaptive = T 3081 3086 7 2.7 367.3 1.3X + + +================================================================================================ +low-cardinality input, no-spill pass-through (Tier 1) +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 287 318 24 58.4 17.1 1.0X +codegen = true, adaptive = T 312 320 7 53.8 18.6 0.9X +codegen = false, adaptive = F 1261 1263 3 13.3 75.1 0.2X +codegen = false, adaptive = T 1301 1303 3 12.9 77.6 0.2X + + +================================================================================================ +high-cardinality input, on-spill pass-through (Tier 2) +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 8319 8410 128 1.0 991.7 1.0X +codegen = true, adaptive = T 4421 4464 61 1.9 527.0 1.9X +codegen = false, adaptive = F 9254 9314 85 0.9 1103.2 0.9X +codegen = false, adaptive = T 5251 5257 8 1.6 626.0 1.6X + + +================================================================================================ +low-cardinality input, on-spill pass-through (Tier 2) +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 808 816 11 20.8 48.1 1.0X +codegen = true, adaptive = T 826 842 19 20.3 49.3 1.0X +codegen = false, adaptive = F 1332 1343 15 12.6 79.4 0.6X +codegen = false, adaptive = T 1303 1304 1 12.9 77.7 0.6X + + diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java index af8d5a4610f64..d850d0d18befe 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java @@ -227,6 +227,14 @@ public double getAvgHashProbesPerKey() { return map.getAvgHashProbesPerKey(); } + /** + * Returns the number of distinct keys currently stored in the underlying `BytesToBytesMap`. + * Used by adaptive partial aggregation to estimate the pre-shuffle reduction ratio. + */ + public int getNumKeys() { + return map.numKeys(); + } + /** * Sorts the map's records in place, spill them to disk, and returns an [[UnsafeKVExternalSorter]] * diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala index 62c4f896f2ee4..6bc21ff2f9c76 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala @@ -72,7 +72,9 @@ case class HashAggregateExec( "aggTime" -> SQLMetrics.createTimingMetric(sparkContext, "time in aggregation build"), "avgHashProbe" -> SQLMetrics.createAverageMetric(sparkContext, "avg hash probes per key"), - "numTasksFallBacked" -> SQLMetrics.createMetric(sparkContext, "number of sort fallback tasks")) + "numTasksFallBacked" -> SQLMetrics.createMetric(sparkContext, "number of sort fallback tasks"), + "numBypassingRows" -> + SQLMetrics.createMetric(sparkContext, "number of bypassing rows")) // This is for testing. We force TungstenAggregationIterator to fall back to the unsafe row hash // map and/or the sort-based aggregation once it has processed a given number of input rows. @@ -94,6 +96,7 @@ case class HashAggregateExec( val avgHashProbe = longMetric("avgHashProbe") val aggTime = longMetric("aggTime") val numTasksFallBacked = longMetric("numTasksFallBacked") + val numBypassingRows = longMetric("numBypassingRows") child.execute().mapPartitionsWithIndex { (partIndex, iter) => @@ -121,7 +124,9 @@ case class HashAggregateExec( peakMemory, spillSize, avgHashProbe, - numTasksFallBacked) + numTasksFallBacked, + numBypassingRows, + adaptivePartialAggConfig) if (!hasInput && groupingExpressions.isEmpty) { numOutputRows += 1 Iterator.single[UnsafeRow](aggregationIterator.outputForEmptyGroupingKeyWithoutInput()) @@ -141,6 +146,47 @@ case class HashAggregateExec( .map(_.asInstanceOf[DeclarativeAggregate]) private val bufferSchema = DataTypeUtils.fromAttributes(aggregateBufferAttributes) + /** + * Runtime configuration for adaptive partial aggregation, or `None` when it does not apply to + * this operator. When defined, 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]]). + * + * Adaptive partial aggregation 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`/`PartialMerge` + * have no such downstream to fall back on. + * - 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 adaptivePartialAggConfig: Option[AdaptivePartialAggregationConfig] = { + val applicable = 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) + if (applicable) { + Some(AdaptivePartialAggregationConfig( + sampleRows = conf.adaptivePartialAggregationSampleRows, + noSpillReductionRatioThreshold = + conf.adaptivePartialAggregationNoSpillReductionRatioThreshold, + spillReductionRatioThreshold = + conf.adaptivePartialAggregationSpillReductionRatioThreshold)) + } else { + None + } + } + // The name for Fast HashMap private var fastHashMapTerm: String = _ private var isFastHashMapEnabled: Boolean = false @@ -154,6 +200,31 @@ case class HashAggregateExec( private var hashMapTerm: String = _ private var sorterTerm: String = _ + // Codegen state for adaptive partial aggregation. When the pre-shuffle reduction ratio of the + // regular (second-level) hash map is too low, the operator stops populating the map and instead + // streams each remaining row through as a single-row partial buffer for the Final aggregate to + // merge. Only the regular map is governed: the append-only fast hash map keeps absorbing hot + // keys, so pass-through only ever applies to the fast-miss stream and the fast path is never + // regressed for high-reduction inputs. + private var adaptivePassThroughTerm: String = _ + private var regularMapRowCountTerm: String = _ + private var adaptiveChildrenConsumedTerm: String = _ + // Whether the map output has already been emitted (and the maps freed). Once pass-through is + // active the maps are frozen, so they are output as soon as pass-through fires to release their + // memory before the remaining input is streamed; this flag lets the final output skip them. + private var adaptiveMapOutputDoneTerm: String = _ + // Whether the map iterators have been set up (`finishHashMap`). The map-output function may be + // re-entered when its loops return via `shouldStop()` to drain the buffer, and `finishAggregate` + // destructs the map, so the setup must run only once. + private var adaptiveMapSetupDoneTerm: String = _ + // The next regular-map row count at which the no-spill tier re-evaluates the reduction ratio. + // It starts at `sampleRows` and doubles after each sub-threshold check, so the ratio is checked + // only rarely once the input proves low-cardinality. + private var adaptiveNextSampleRowTerm: String = _ + // The name of the generated output function, promoted to a field so `doConsumeWithKeys` can emit + // pass-through rows directly from within the build loop. + private var outputFunc: String = _ + /** * This is called by generated Java class, should be public. */ @@ -436,6 +507,20 @@ case class HashAggregateExec( protected override def doProduceWithKeys(ctx: CodegenContext): String = { val initAgg = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initAgg") + if (adaptivePartialAggConfig.isDefined) { + adaptivePassThroughTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptivePassThrough") + regularMapRowCountTerm = ctx.addMutableState(CodeGenerator.JAVA_LONG, "regularMapRowCount") + adaptiveNextSampleRowTerm = + ctx.addMutableState(CodeGenerator.JAVA_LONG, "adaptiveNextSampleRow", + v => s"$v = ${adaptivePartialAggConfig.get.sampleRows}L;") + adaptiveChildrenConsumedTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveChildrenConsumed") + adaptiveMapOutputDoneTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveMapOutputDone") + adaptiveMapSetupDoneTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveMapSetupDone") + } if (conf.enableTwoLevelAggMap) { enableTwoLevelHashMap() } else if (conf.enableVectorizedHashMap) { @@ -535,19 +620,30 @@ case class HashAggregateExec( // `addNewFunction` spills this helper into a nested class (as can happen // once the outer class passes the code-size threshold), the bare field // reference fails with `IllegalAccessError`. + + // Generate code for output. This must happen before the `doAgg` helper below, because with + // adaptive partial aggregation enabled, `doConsumeWithKeys` (invoked from the child's produce + // inside `doAgg`) emits pass-through rows by calling this output function directly. + val keyTerm = ctx.freshName("aggKey") + val bufferTerm = ctx.freshName("aggBuffer") + outputFunc = generateResultFunction(ctx) + + // After the child input is consumed, finish the build: with adaptive partial aggregation mark + // that the child is fully consumed (to support re-entry; the map iterators are set up inside + // the map-output function), otherwise set up the map iterators for the output below. + val postChildProduce = if (adaptivePartialAggConfig.isDefined) { + s"$adaptiveChildrenConsumedTerm = true;" + } else { + finishHashMap + } val doAggFuncName = ctx.addNewFunction(doAgg, s""" |private void $doAgg(int partitionIndex) throws java.io.IOException { | ${child.asInstanceOf[CodegenSupport].produce(ctx, this)} - | $finishHashMap + | $postChildProduce |} """.stripMargin) - // generate code for output - val keyTerm = ctx.freshName("aggKey") - val bufferTerm = ctx.freshName("aggBuffer") - val outputFunc = generateResultFunction(ctx) - val limitNotReachedCondition = limitNotReachedCond def outputFromFastHashMap: String = { @@ -615,8 +711,78 @@ case class HashAggregateExec( """.stripMargin } + // With adaptive partial aggregation the maps are frozen once pass-through is active, so their + // output (which also frees them) can happen as soon as pass-through fires, releasing the memory + // before the remaining input is streamed. The output loops are wrapped in a function so the + // same code runs either early (once pass-through freezes the maps) or at the end of the build. + // The done flag is set inside, after the loops, so a mid-output drain (the loops return via + // `shouldStop()`) leaves it unset and the caller resumes the map iterator on re-entry; once it + // is set the maps have been fully output and freed and will not be touched again. The iterator + // setup (`finishHashMap`, which destructs the map) is guarded to run only once. + val outputMapFuncName = if (adaptivePartialAggConfig.isDefined) { + val name = ctx.freshName("outputMap") + ctx.addNewFunction(name, + s""" + |private void $name() throws java.io.IOException { + | if (!$adaptiveMapSetupDoneTerm) { + | $finishHashMap + | $adaptiveMapSetupDoneTerm = true; + | } + | $outputFromFastHashMap + | $outputFromRegularHashMap + | $adaptiveMapOutputDoneTerm = true; + |} + """.stripMargin) + } else { + "" + } + val aggTime = metricTerm(ctx, "aggTime") val beforeAgg = ctx.freshName("beforeAgg") + // With adaptive partial aggregation, `doAgg` may start appending pass-through rows to the + // output buffer mid-build. In that case `shouldStop()` becomes true and we must return so the + // buffered rows are drained; the build is resumed on re-entry (guarded by `childrenConsumed`) + // until the child input is exhausted, only then falling through to the map output below. + val adaptiveStopCheck = if (adaptivePartialAggConfig.isDefined) { + "if (shouldStop()) return;" + } else { + "" + } + // Once pass-through is active the maps are frozen, so output them (releasing their memory) + // exactly once: early in `adaptiveResumeBuild` (resuming the build means it returned because + // pass-through filled the buffer, so pass-through is already active) or at the end in + // `adaptiveFinalOutput` when they were never output early. The output loops return via + // `shouldStop()` when the buffer fills, so the done flag is set inside the output function and + // re-entry resumes the map iterator. + val adaptiveOutputMap = if (adaptivePartialAggConfig.isDefined) { + s""" + |if (!$adaptiveMapOutputDoneTerm) { + | $outputMapFuncName(); + | if (shouldStop()) return; + |} + """.stripMargin + } else { + "" + } + val adaptiveResumeBuild = if (adaptivePartialAggConfig.isDefined) { + s""" + |if (!$adaptiveChildrenConsumedTerm) { + | $adaptiveOutputMap + | $doAggFuncName(partitionIndex); + | if (shouldStop()) return; + |} + """.stripMargin + } else { + "" + } + val adaptiveFinalOutput = if (adaptivePartialAggConfig.isDefined) { + adaptiveOutputMap + } else { + s""" + |$outputFromFastHashMap + |$outputFromRegularHashMap + """.stripMargin + } s""" |if (!$initAgg) { | $initAgg = true; @@ -626,13 +792,27 @@ 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 must be re-enabled to keep the buffer bounded. + override def needStopCheck: Boolean = + adaptivePartialAggConfig.isDefined || super.needStopCheck + + // 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`. Copy the result so the buffered rows do not collapse into the last one. + override def needCopyResult: Boolean = + adaptivePartialAggConfig.isDefined || super.needCopyResult + protected override def doConsumeWithKeys(ctx: CodegenContext, input: Seq[ExprCode]): String = { // create grouping key val unsafeRowKeyCode = GenerateUnsafeProjection.createCode( @@ -644,6 +824,18 @@ case class HashAggregateExec( val unsafeRowBuffer = ctx.freshName("unsafeRowAggBuffer") val fastRowBuffer = ctx.freshName("fastAggBuffer") + // For adaptive partial aggregation pass-through, each bypassed row is emitted as a single-row + // partial buffer: start from the initial aggregation buffer, apply the update expressions once, + // and output `key ++ buffer` for the Final aggregate to merge. This projects the initial + // buffer. + val emptyAggBufferCode = if (adaptivePartialAggConfig.isDefined) { + GenerateUnsafeProjection.createCode(ctx, declFunctions.flatMap(f => f.initialValues)) + } else { + null + } + // Per-row local flag marking that the current row is being streamed through (held by no map). + val adaptiveRowBypassedTerm = ctx.freshName("adaptiveRowBypassed") + // To individually generate code for each aggregate function, an element in `updateExprs` holds // all the expressions for the buffer of an aggregation function. val updateExprs = aggregateExpressions.map { e => @@ -663,46 +855,124 @@ case class HashAggregateExec( case _ => ("true", "", "") } - val findOrInsertRegularHashMap: String = - s""" - |// generate grouping key - |${unsafeRowKeyCode.code} - |int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode(); - |if ($checkFallbackForBytesToBytesMap) { - | // try to get the buffer from hash map - | $unsafeRowBuffer = - | $hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys, $unsafeRowKeyHash); - |} - |// Can't allocate buffer from the hash map. Spill the map and fallback to sort-based - |// aggregation after processing all input rows. - |if ($unsafeRowBuffer == null) { - | if ($sorterTerm == null) { - | $sorterTerm = $hashMapTerm.destructAndCreateExternalSorter(); - | } else { - | $sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter()); - | } - | $resetCounter - | // the hash map had be spilled, it should have enough memory now, - | // try to allocate buffer again. - | $unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow( - | $unsafeRowKeys, $unsafeRowKeyHash); - | if ($unsafeRowBuffer == null) { - | // failed to allocate the first page - | throw QueryExecutionErrors.aggregateOutOfMemoryError(); - | } - |} - """.stripMargin + val findOrInsertRegularHashMap: String = { + // Assumes the grouping key projection (`unsafeRowKeyCode.code`) has already run for this row, + // so `unsafeRowKeyCode.value` holds the current key. The projection is emitted exactly once + // per regular-map row (see below); emitting it in more than one runtime branch is unsafe + // because the projection's subexpression/writer state assigned in one branch would be read + // stale from another (e.g. the adaptive pass-through path would reuse the last probed key). + val probeRegularMap = + s""" + |int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode(); + |if ($checkFallbackForBytesToBytesMap) { + | // try to get the buffer from hash map + | $unsafeRowBuffer = + | $hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys, $unsafeRowKeyHash); + |} + """.stripMargin + + val spillMap = + s""" + |if ($sorterTerm == null) { + | $sorterTerm = $hashMapTerm.destructAndCreateExternalSorter(); + |} else { + | $sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter()); + |} + |$resetCounter + |// the hash map had been spilled, so it should have enough memory now, + |// try to allocate buffer again. + |$unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow( + | $unsafeRowKeys, $unsafeRowKeyHash); + |if ($unsafeRowBuffer == null) { + | // failed to allocate the first page + | throw QueryExecutionErrors.aggregateOutOfMemoryError(); + |} + """.stripMargin + + if (adaptivePartialAggConfig.isDefined) { + val cfg = adaptivePartialAggConfig.get + // Adaptive partial aggregation governs only this regular (second-level) map. Count the + // rows that enter it (a fast-map miss, or every row when the fast map is off) and use + // `regularMap.getNumKeys() / regularRows` as the pre-shuffle reduction ratio. + // - Tier 2 (on-spill): when the map cannot allocate for a new key (it would otherwise + // spill), bypass instead if the ratio is at least `spillReductionRatioThreshold`. + // - Tier 1 (no-spill): from `sampleRows` regular rows on, bypass if the ratio is at + // least `noSpillReductionRatioThreshold`. The sampling window doubles after each + // sub-threshold check, so low-cardinality input is re-evaluated only rarely while a + // late high-cardinality tail can still trigger the bypass. + // Both tiers fire only before any spill (`sorter == null`): once the map has spilled, the + // reduction-ratio estimate no longer covers the spilled rows, and pass-through must never + // coexist with sort-based aggregation. When the map is full after a spill, the map spills + // again as usual. + // The key projection runs once here so `unsafeRowKeyCode.value` is valid for both the + // probe below and the pass-through buffer built by the caller. + s""" + |// generate grouping key + |${unsafeRowKeyCode.code} + |if (!$adaptivePassThroughTerm) { + | $probeRegularMap + | if ($unsafeRowBuffer == null) { + | if ($sorterTerm == null && $regularMapRowCountTerm > 0 && + | (double) $hashMapTerm.getNumKeys() >= + | $regularMapRowCountTerm * ${cfg.spillReductionRatioThreshold}D) { + | $adaptivePassThroughTerm = true; + | } else { + | $spillMap + | } + | } + | if ($unsafeRowBuffer != null) { + | $regularMapRowCountTerm += 1; + | if ($sorterTerm == null && + | $regularMapRowCountTerm == $adaptiveNextSampleRowTerm) { + | if ((double) $hashMapTerm.getNumKeys() >= + | $regularMapRowCountTerm * ${cfg.noSpillReductionRatioThreshold}D) { + | $adaptivePassThroughTerm = true; + | } else { + | $adaptiveNextSampleRowTerm = $adaptiveNextSampleRowTerm * 2; + | } + | } + | } + |} + """.stripMargin + } else { + s""" + |// generate grouping key + |${unsafeRowKeyCode.code} + |$probeRegularMap + |// Can't allocate buffer from the hash map. Spill the map and fallback to sort-based + |// aggregation after processing all input rows. + |if ($unsafeRowBuffer == null) { + | $spillMap + |} + """.stripMargin + } + } val findOrInsertHashMap: String = { - if (isFastHashMapEnabled) { + val findCode = if (isFastHashMapEnabled) { // If fast hash map is on, we first generate code to probe and update the fast hash map. // If the probe is successful the corresponding fast row buffer will hold the mutable row. + // Once adaptive pass-through is active, skip the fast map entirely so the row is streamed + // through instead of being inserted anywhere. + val fastMapProbe = + s""" + |${fastRowKeys.map(_.code).mkString("\n")} + |if (${fastRowKeys.map("!" + _.isNull).mkString(" && ")}) { + | $fastRowBuffer = $fastHashMapTerm.findOrInsert( + | ${fastRowKeys.map(_.value).mkString(", ")}); + |} + """.stripMargin + val guardedFastMapProbe = if (adaptivePartialAggConfig.isDefined) { + s""" + |if (!$adaptivePassThroughTerm) { + | $fastMapProbe + |} + """.stripMargin + } else { + fastMapProbe + } s""" - |${fastRowKeys.map(_.code).mkString("\n")} - |if (${fastRowKeys.map("!" + _.isNull).mkString(" && ")}) { - | $fastRowBuffer = $fastHashMapTerm.findOrInsert( - | ${fastRowKeys.map(_.value).mkString(", ")}); - |} + |$guardedFastMapProbe |// Cannot find the key in fast hash map, try regular hash map. |if ($fastRowBuffer == null) { | $findOrInsertRegularHashMap @@ -711,6 +981,31 @@ case class HashAggregateExec( } else { findOrInsertRegularHashMap } + + // When pass-through is active, a row that no map holds must be streamed through. + // `rowBypassed` marks exactly those rows: the fast map and regular map probes are skipped + // (guarded above), so both buffers stay null. The Tier-1 transition row is excluded on + // purpose -- its probe already inserted the key into the regular map, so it is aggregated + // there and must not be re-emitted. + val createPassThroughBuffer = if (adaptivePartialAggConfig.isDefined) { + // The grouping key was already projected in `findOrInsertRegularHashMap` + // (`unsafeRowKeyCode.code`), so `unsafeRowKeyCode.value` holds this row's key. Only build + // the single-row partial buffer here. + s""" + |if ($adaptivePassThroughTerm && $unsafeRowBuffer == null) { + | $adaptiveRowBypassedTerm = true; + | ${emptyAggBufferCode.code} + | $unsafeRowBuffer = ${emptyAggBufferCode.value}; + |} + """.stripMargin + } else { + "" + } + + s""" + |$findCode + |$createPassThroughBuffer + """.stripMargin } val inputAttrs = aggregateBufferAttributes ++ inputAttributes @@ -845,29 +1140,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 { - "UnsafeRow" + s"UnsafeRow $unsafeRowBuffer = null;" + } + val declareBypassed = if (adaptivePartialAggConfig.isDefined) { + s"boolean $adaptiveRowBypassedTerm = false;" + } else { + "" } 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 (adaptivePartialAggConfig.isDefined) { + val numBypassingRows = metricTerm(ctx, "numBypassingRows") + s""" + |if ($adaptiveRowBypassedTerm) { + | $numBypassingRows.add(1); + | $outputFunc(${unsafeRowKeyCode.value}, $unsafeRowBuffer); + |} + """.stripMargin + } else { + "" + } s""" |$declareRowBuffer |$findOrInsertHashMap |$incCounter |$updateRowInHashMap + |$emitPassThroughRow """.stripMargin } @@ -897,3 +1219,24 @@ case class HashAggregateExec( override protected def withNewChildInternal(newChild: SparkPlan): HashAggregateExec = copy(child = newChild) } + +/** + * Runtime parameters that control adaptive partial aggregation for a single [[HashAggregateExec]]. + * + * The aggregation samples the pre-shuffle reduction ratio (distinct grouping keys / processed + * rows) and bypasses partial aggregation when the ratio is too high to be worthwhile, using two + * tiers: + * - no-spill tier: evaluated from `sampleRows` rows on while the aggregation map is still fully + * in memory; the sampling window doubles after each sub-threshold check, so a low-cardinality + * input is re-evaluated only rarely while a late high-cardinality tail can still be caught. + * In-memory partial aggregation is cheap, so this tier uses the more conservative + * `noSpillReductionRatioThreshold`. + * - on-spill tier: evaluated when the aggregation map is about to spill. Partial aggregation now + * pays disk I/O costs, so this tier uses the more aggressive `spillReductionRatioThreshold`. + * + * Once either tier triggers, partial aggregation is bypassed for the rest of the input. + */ +case class AdaptivePartialAggregationConfig( + sampleRows: Int, + noSpillReductionRatioThreshold: Double, + spillReductionRatioThreshold: Double) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala index 00d18a2f79a81..450578984d6b1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala @@ -95,7 +95,9 @@ class TungstenAggregationIterator( peakMemory: SQLMetric, spillSize: SQLMetric, avgHashProbe: SQLMetric, - numTasksFallBacked: SQLMetric) + numTasksFallBacked: SQLMetric, + numBypassingRows: SQLMetric = null, + adaptivePartialAggConfig: Option[AdaptivePartialAggregationConfig] = None) extends AggregationIterator( partIndex, groupingExpressions, @@ -179,6 +181,21 @@ class TungstenAggregationIterator( // hashMap. If there is not enough memory, it will multiple hash-maps, spilling // after each becomes full then using sort to merge these spills, finally do sort // based aggregation. + // + // When adaptive partial aggregation is enabled (see [[AdaptivePartialAggregationConfig]]), the + // processing may stop early and switch to pass-through mode: the remaining input rows are not + // added to the map but are instead emitted as single-row partial buffers by the output stage + // (see `passThroughOutput`). Two tiers decide the switch, both evaluated only before any spill + // has happened (so `externalSorter == null`), which keeps the reduction-ratio estimate based on + // the full set of processed rows and guarantees `passThrough` never coexists with sort-based + // aggregation: + // - no-spill tier: from `sampleRows` rows on while the map is in memory, if + // distinctKeys / processedRows >= noSpillReductionRatioThreshold; the sampling window doubles + // after each sub-threshold check, so a low-cardinality input is re-evaluated only rarely. + // - on-spill tier: when the map is about to spill for the first time, if + // distinctKeys / processedRows >= spillReductionRatioThreshold. In this case we do NOT spill; + // the full in-memory map is kept for normal output and the row that could not be inserted + // becomes the first pass-through row. private def processInputs(fallbackStartsAt: (Int, Int)): Unit = { if (groupingExpressions.isEmpty) { // If there is no grouping expressions, we can just reuse the same buffer over and over again. @@ -191,7 +208,19 @@ class TungstenAggregationIterator( } } else { var i = 0 - while (inputIter.hasNext) { + var processedRows = 0L + // Eligibility and the thresholds are fixed for the lifetime of this iterator, so unwrap them + // once instead of re-checking the `Option` for every row. + val adaptiveEnabled = adaptivePartialAggConfig.isDefined + val noSpillRatioThreshold = + adaptivePartialAggConfig.map(_.noSpillReductionRatioThreshold).getOrElse(0.0) + val spillRatioThreshold = + adaptivePartialAggConfig.map(_.spillReductionRatioThreshold).getOrElse(0.0) + // The next row count at which the no-spill tier re-evaluates the reduction ratio. It starts + // at `sampleRows` and doubles after each sub-threshold check, so the ratio is checked only + // rarely once the input proves low-cardinality. + var nextSampleRow = adaptivePartialAggConfig.map(_.sampleRows.toLong).getOrElse(0L) + while (inputIter.hasNext && !passThrough) { val newInput = inputIter.next() val groupingKey = groupingProjection.apply(newInput) var buffer: UnsafeRow = null @@ -199,21 +228,46 @@ class TungstenAggregationIterator( buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) } if (buffer == null) { - val sorter = hashMap.destructAndCreateExternalSorter() - if (externalSorter == null) { - externalSorter = sorter + // The map is full and would normally spill. On the first spill, adaptive partial + // aggregation may instead bypass: keep the in-memory map as-is, pass this row and all + // remaining rows through, and skip the spill entirely. + if (adaptiveEnabled && externalSorter == null && processedRows > 0 && + hashMap.getNumKeys().toDouble >= processedRows * spillRatioThreshold) { + passThrough = true + // `newInput` could not be inserted; stash a copy as the first pass-through row so it + // is not lost when we drain the rest of `inputIter`. + pendingPassThroughRow = newInput.copy() } else { - externalSorter.merge(sorter) + val sorter = hashMap.destructAndCreateExternalSorter() + if (externalSorter == null) { + externalSorter = sorter + } else { + externalSorter.merge(sorter) + } + i = 0 + buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) + if (buffer == null) { + // failed to allocate the first page + throw QueryExecutionErrors.aggregateOutOfMemoryError() + } } - i = 0 - buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) - if (buffer == null) { - // failed to allocate the first page - throw QueryExecutionErrors.aggregateOutOfMemoryError() + } + if (!passThrough) { + processRow(buffer, newInput) + i += 1 + processedRows += 1 + // No-spill tier: from the sampling window on, if the map is still fully in memory and + // the reduction ratio is too high to be worthwhile, bypass partial aggregation for the + // rest. The window doubles after each sub-threshold check so low-cardinality input is + // re-evaluated only rarely while a late high-cardinality tail can still be caught. + if (adaptiveEnabled && externalSorter == null && processedRows == nextSampleRow) { + if (hashMap.getNumKeys().toDouble >= processedRows * noSpillRatioThreshold) { + passThrough = true + } else { + nextSampleRow = nextSampleRow * 2 + } } } - processRow(buffer, newInput) - i += 1 } if (externalSorter != null) { @@ -354,6 +408,49 @@ 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 adaptive tier. + // Because both tiers only trigger before any spill, pass-through never coexists with sort-based + // aggregation, so the output order is: map entries first, then the pass-through rows. + private[this] var passThrough: Boolean = false + + // The row that could not be inserted at the on-spill tier trigger point. It is stashed here (as + // a copy) so it becomes the first pass-through row rather than being lost. + private[this] var pendingPassThroughRow: InternalRow = null + + // A reused aggregation buffer for building single-row partial buffers during pass-through. It is + // re-initialized from `initialAggregationBuffer` for every passed-through row. + private[this] lazy val passThroughAggregationBuffer: UnsafeRow = createNewAggregationBuffer() + + // Whether there are remaining pass-through rows to emit. + private def passThroughHasNext: Boolean = + passThrough && (pendingPassThroughRow != null || inputIter.hasNext) + + // Emits the next input row as a single-row partial aggregation buffer, i.e. a group of size one. + // The output (grouping key ++ buffer) is a valid partial buffer that the downstream Final + // aggregation merges, so the result is identical to running partial aggregation on this row. + private def nextPassThroughOutput(): UnsafeRow = { + val row = if (pendingPassThroughRow != null) { + val stashed = pendingPassThroughRow + pendingPassThroughRow = null + stashed + } else { + inputIter.next() + } + val groupingKey = groupingProjection.apply(row) + // Reset the buffer to initial values, then update it with this single row. + passThroughAggregationBuffer.copyFrom(initialAggregationBuffer) + processRow(passThroughAggregationBuffer, row) + if (numBypassingRows != null) { + numBypassingRows += 1 + } + generateOutput(groupingKey, passThroughAggregationBuffer) + } + /////////////////////////////////////////////////////////////////////////// // Part 6: Loads input rows and setup aggregationBufferMapIterator if we // have not switched to sort-based aggregation. @@ -398,7 +495,8 @@ class TungstenAggregationIterator( /////////////////////////////////////////////////////////////////////////// override final def hasNext: Boolean = { - (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) + (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) || + passThroughHasNext } override final def next(): UnsafeRow = { @@ -412,7 +510,7 @@ class TungstenAggregationIterator( sortBasedAggregationBuffer.copyFrom(initialAggregationBuffer) outputRow - } else { + } else if (mapIteratorHasNext) { // We did not fall back to sort-based aggregation. val result = generateOutput( @@ -426,13 +524,17 @@ class TungstenAggregationIterator( if (!mapIteratorHasNext) { // If there is no input from aggregationBufferMapIterator, we copy current result. val resultCopy = result.copy() - // Then, we free the map. + // Then, we free the map. Pass-through (if any) does not use the map. hashMap.free() resultCopy } else { result } + } else { + // Adaptive partial aggregation bypassed partial aggregation: emit the remaining input + // rows as single-row partial buffers. + nextPassThroughOutput() } numOutputRows += 1 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala new file mode 100644 index 0000000000000..05907997a419d --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala @@ -0,0 +1,833 @@ +/* + * 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 decision tiers 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 on-spill (Tier 2) 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 sample so the no-spill (Tier 1) path triggers on modest inputs. + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_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 Tier 1 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 Tier 2 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( + skipped = partialAggs.map(_.metrics("numBypassingRows").value).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("numBypassingRows").value + }.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 + * sample, threading a descriptive clue for failure messages. + * + * The fast (first-level) map is append-only and never spills; adaptive partial aggregation + * governs only the regular (second-level) map. 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 on-spill tier); when 0 the regular map does not spill. + */ + private def forEachCodegenAndMap( + sampleRows: Int = 8, + regularFallback: Int = 0, + noSpillThreshold: 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 (noSpillThreshold >= 0.0) { + Seq(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD.key -> + noSpillThreshold.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_SAMPLE_ROWS.key -> sampleRows.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 both tiers. + 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") { + checkAdaptiveMatchesReference { () => + spark.range(0, 400, 1, 1) + .select( + when($"id" % 4 === 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 + // sampling tier 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 sample lets the map accumulate many keys before the no-spill tier 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_SAMPLE_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 reduction ratio, so we deliberately do not port that + // exclusion. These tests assert results stay correct with the exclusion absent. + + 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)") { + withTempView("t") { + spark.range(0, 400, 1, 1) + .select(($"id" % 180) as "k1", ($"id" % 6) 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), ())""".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") + } + } + + ///////////////////////////////////////////////////////////////////////////// + // Part 2: Triggering -- the bypass fires when, and only when, it should. + ///////////////////////////////////////////////////////////////////////////// + + test("pass-through fires for high-cardinality input, not for low-cardinality input") { + forEachCodegenAndMap() { clue => + // Fully distinct keys: partial aggregation reduces nothing, so rows must bypass. + val highCard = () => spark.range(0, 200, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(highCard) > 0, + "expected some rows to bypass partial aggregation for high-cardinality input") + } + // Few distinct keys, high reduction: partial aggregation is effective, nothing bypasses. + val lowCard = () => spark.range(0, 600, 1, 1) + .select(($"id" % 5) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(lowCard) == 0, + "expected no rows to bypass partial aggregation for low-cardinality input") + } + } + } + + test("group-by-only pass-through fires for high-cardinality input") { + forEachCodegenAndMap() { clue => + val distinctKeys = () => spark.range(0, 200, 1, 1).select($"id" as "k").distinct() + withClue(clue) { + assert(numBypassingRows(distinctKeys) > 0, + "expected group-by-only rows to bypass partial aggregation for high-cardinality input") + } + } + } + + test("filtered aggregate functions are eligible for pass-through") { + // The filter clause does not change eligibility: a partial aggregate over `FILTER (WHERE ...)` + // functions still bypasses on high-cardinality input, and the per-row filter guard runs inside + // the pass-through single-row buffer update. + forEachCodegenAndMap() { clue => + withTempView("t") { + spark.range(0, 200, 1, 1) + .select($"id" as "k", ($"id" % 100) as "v") + .createOrReplaceTempView("t") + val df = () => spark.sql( + """SELECT k, sum(v) FILTER (WHERE v % 2 = 0) AS s + |FROM t GROUP BY k""".stripMargin) + withClue(clue) { + assert(numBypassingRows(df) > 0, + "expected filtered-aggregate rows to bypass partial aggregation for high-cardinality " + + "input") + } + } + } + } + + test("pass-through fires for high-cardinality input below an Expand") { + // The static PR#28804 heuristic would have refused to skip whenever an Expand was present; our + // runtime decision skips because the expanded rows genuinely do not reduce. GROUPING SETS over + // two single-column, fully-distinct sets is used (rather than ROLLUP/CUBE) so there is no + // grand-total group dragging the reduction ratio below the threshold: every expanded row is a + // fresh key, so all configurations bypass. + forEachCodegenAndMap() { clue => + withTempView("t") { + spark.range(0, 200, 1, 1) + .select($"id".as("k1"), ($"id" + 1000).as("k2"), $"id".as("v")) + .createOrReplaceTempView("t") + val gs = () => spark.sql( + """SELECT k1, k2, sum(v) AS s + |FROM t + |GROUP BY k1, k2 GROUPING SETS ((k1), (k2))""".stripMargin) + withClue(clue) { + assert(numBypassingRows(gs) > 0, + "expected rows below an Expand to bypass partial aggregation for high-cardinality " + + "input") + } + } + } + } + + test("no pass-through when the feature is disabled") { + // The metric must stay zero across the whole matrix when the switch is off, even for input that + // would otherwise bypass. + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString) ++ fixedPlanConfs): _*) { + val df = () => spark.range(0, 200, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") { + assert(numBypassingRows(df) == 0, + "no rows should bypass partial aggregation when the feature is disabled") + } + } + } + } + + test("no pass-through for a global aggregation with no grouping keys") { + // Global aggregation is ineligible (`groupingExpressions` is empty): there is a single group, + // so there is nothing to stream through. The partial aggregate must never bypass regardless of + // codegen or map settings, even under a forced fallback. + forEachCodegenAndMap(regularFallback = 16) { clue => + val df = () => spark.range(0, 200, 1, 1) + .agg(sum($"id") as "s", count(lit(1)) as "c") + withClue(clue) { + assert(numBypassingRows(df) == 0, + "a global aggregation is not eligible and must never bypass") + } + } + } + + test("Tier 1 (no-spill) fires when the sample shows no reduction, without spilling") { + // No forced regular-map spill: only the no-spill sampling tier can trigger the bypass. Fully + // distinct keys over a small sample cross `noSpillReductionRatioThreshold`, so rows bypass and + // the regular map never spills. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(df) + assert(c.skipped > 0, "Tier 1 should bypass fully distinct input under the sample") + assert(c.spillBytes == 0, "Tier 1 must decide before any spill happens") + assert(c.tasksFallBacked == 0, "Tier 1 must not fall back to sort-based aggregation") + } + } + } + + test("Tier 2 (on-spill) fires when the map would spill on high-cardinality input") { + // Force the regular map to fall back quickly. High-cardinality input that reaches the fallback + // point should bypass via the on-spill tier rather than spilling. Use a sample larger than the + // input so Tier 1 cannot fire first and the on-spill tier is the one exercised. + forEachCodegenAndMap(sampleRows = 100000, regularFallback = 16) { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(df) + assert(c.skipped > 0, + "Tier 2 should bypass high-cardinality input at the spill boundary") + // The whole point of Tier 2 is to bypass *instead of* falling back to sort-based + // aggregation, so the sorter is never created. `numTasksFallBacked` is the reliable + // per-operator signal for that (the `spillSize` metric on the interpreted path is derived + // from the task-cumulative memory-spill counter and can be inflated by unrelated spilling + // such as the downstream shuffle write, so it is not asserted here). + assert(c.tasksFallBacked == 0, "Tier 2 must replace the sort fallback, not trigger it") + } + } + } + + test("without the feature the same input really does fall back to sort") { + // Sanity check for the Tier 2 assertion above: with adaptive disabled, the identical + // high-cardinality input under the same forced fallback genuinely falls back to sort-based + // aggregation. This proves Tier 2's `tasksFallBacked == 0` reflects the bypass and not merely + // an input that never reached the spill boundary. + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + val fastCap = if (twoLevelMap) 4 else 1 + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> s"$fastCap, 16") ++ + fixedPlanConfs): _*) { + val df = () => spark.range(0, 200, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") { + val c = runAndReadCounters(df) + assert(c.skipped == 0, "feature disabled: nothing should bypass") + assert(c.tasksFallBacked > 0, + "feature disabled: the forced fallback should trigger sort-based aggregation") + } + } + } + } + + test("spill tier decides identically at the exact ratio boundary with codegen on and off") { + // The on-spill tier evaluates the ratio over the rows already aggregated, excluding the failed + // insertion that becomes the first pass-through row, so both execution paths must judge the + // same row set and reach the same decision. `id % 40` over 400 rows gives 40 distinct keys + // when the map fills at 50 aggregated rows, i.e. a ratio of exactly 0.8: at the threshold the + // bypass fires, just above it (0.85) it does not. + Seq(0.8 -> true, 0.85 -> false).foreach { case (threshold, shouldBypass) => + val skippedPerCodegen = Seq(true, false).map { wholeStage => + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> "false", + // A sample larger than the input keeps the no-spill tier out of the picture. + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_ROWS.key -> "100000", + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_SPILL_REDUCTION_RATIO_THRESHOLD.key -> + threshold.toString, + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 50") ++ + fixedPlanConfs): _*) { + val df = () => spark.range(0, 400, 1, 1) + .select(($"id" % 40) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(s"threshold=$threshold wholeStage=$wholeStage") { + numBypassingRows(df) + } + } + } + withClue(s"threshold=$threshold skipped=$skippedPerCodegen") { + assert(skippedPerCodegen.forall(_ > 0) == shouldBypass, + s"expected bypass=$shouldBypass at the ratio boundary") + assert(skippedPerCodegen.map(_ > 0).distinct.length == 1, + "codegen and interpreted paths must reach the same decision at the boundary") + } + } + } + + test("a zero threshold always bypasses once the sample has been processed") { + // 0 is the most aggressive setting: the ratio check `distinctKeys >= rows * 0` always holds, + // so even a low-cardinality input that the default threshold keeps aggregating is bypassed + // right after the sample. The results must still match the feature-off reference. + forEachCodegenAndMap(noSpillThreshold = 0.0) { clue => + val df = () => spark.range(0, 600, 1, 1) + .select(($"id" % 5) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(df) > 0, + "a zero threshold must bypass even a low-cardinality input") + } + } + } + + test("larger sample defers the decision so a small high-cardinality input is not bypassed") { + // With a sample larger than the whole input and no regular-map spill forced, the Tier 1 check + // point is never reached, so nothing bypasses even though the keys are fully distinct. + forEachCodegenAndMap(sampleRows = 100000) { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(df) == 0, + "no bypass expected before the sample size is reached") + } + } + } + + test("partial aggregate output row count reflects the bypass (independent of the skip metric)") { + // `numOutputRows` on the partial aggregate is the ground truth: it is driven by the normal + // aggregation output path, not by our self-reported `numBypassingRows` metric. This test cross + // checks the two and pins the observable data-side effect of bypassing. + val numRows = 200 + forEachCodegenAndMap() { clue => + // Fully distinct keys: once the bypass fires the operator stops collapsing rows, so the + // partial aggregate emits far more than the handful of keys a real aggregation would. Read + // all counters from a single execution so the metrics are not double-counted. + val highCard = () => spark.range(0, numRows, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(highCard) + assert(c.skipped > 0, "high-cardinality input should bypass") + // Every partial output row is either a real (aggregated) group or a bypassed row, so the + // partial output count must be at least the number of bypassed rows, and it climbs toward + // the input row count -- well above the heavy reduction a kept aggregation would give. + assert(c.partialOutputRows >= c.skipped, + s"partial output ${c.partialOutputRows} should be >= bypassed rows ${c.skipped}") + assert(c.partialOutputRows > numRows / 2, + s"partial output (${c.partialOutputRows}) should climb toward the input row count") + } + + // Low-cardinality reference: partial aggregation stays effective, so its output equals the + // small number of distinct keys and nothing is bypassed. + val lowCard = () => spark.range(0, 600, 1, 1) + .select(($"id" % 5) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(lowCard) + assert(c.skipped == 0, "low-cardinality input should not bypass") + assert(c.partialOutputRows == 5, + "an effective partial aggregate should emit exactly the distinct key count") + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/AdaptivePartialAggregationBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/AdaptivePartialAggregationBenchmark.scala new file mode 100644 index 0000000000000..2b33c37ce781f --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/AdaptivePartialAggregationBenchmark.scala @@ -0,0 +1,137 @@ +/* + * 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.benchmark + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.internal.SQLConf + +/** + * Benchmark comparing runtime adaptive partial aggregation (see + * [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]) against the static pre-shuffle partial + * aggregation. When the partial aggregation is not reducing rows, the operator streams the + * remaining rows through as single-row partial buffers instead of maintaining (and possibly + * spilling) a large aggregation map. + * + * Each scenario runs the query across the full matrix of whole-stage codegen on/off and the + * feature disabled (`adaptive = F`, the pre-change baseline) vs enabled (`adaptive = T`), over a + * {high, low}-cardinality x {no-spill, on-spill} grid: + * - high-cardinality, no spill: the no-spill tier bypasses, which should win. + * - low-cardinality, no spill: nothing bypasses, which must not regress. + * - high-cardinality, forced regular-map spill: the on-spill tier bypasses instead of spilling, + * which should win. + * - low-cardinality, forced regular-map spill: the ratio is too low for the on-spill tier to + * bypass, so both runs spill identically (no regression). + * + * To run this benchmark: + * {{{ + * 1. build/sbt "sql/Test/runMain + * org.apache.spark.sql.execution.benchmark.AdaptivePartialAggregationBenchmark" + * 2. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/Test/runMain + * org.apache.spark.sql.execution.benchmark.AdaptivePartialAggregationBenchmark" + * Results will be written to "benchmarks/AdaptivePartialAggregationBenchmark-results.txt". + * }}} + */ +object AdaptivePartialAggregationBenchmark extends SqlBasedBenchmark { + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + // The upstream `CombineAdjacentAggregation` and `ReplaceHashWithSortAgg` rules would collapse + // or convert these single-partition hash aggregates, so both are disabled to keep the + // Partial+Final `HashAggregateExec` structure the adaptive feature governs. + val fixedPlanConfs = Seq( + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false", + SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "false") + + // Adds the (whole-stage codegen, adaptive switch) matrix for `query`. `extraConf` is applied + // to all four cases so the only differences are the two axes. + def addCodegenAdaptiveCases( + benchmark: Benchmark, + query: () => DataFrame, + extraConf: Seq[(String, String)] = Nil): Unit = { + for { + wholeStage <- Seq(true, false) + adaptive <- Seq(false, true) + } { + val adaptiveLabel = if (adaptive) "T" else "F" + val label = s"codegen = $wholeStage, adaptive = $adaptiveLabel" + benchmark.addCase(label) { _ => + withSQLConf( + (Seq( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> adaptive.toString) ++ + fixedPlanConfs ++ extraConf): _*) { + query().noop() + } + } + } + } + + // Fully distinct keys make partial aggregation useless, so the no-spill (Tier 1) sampling tier + // bypasses: the feature should be faster than the baseline that maintains a map entry per row. + runBenchmark("high-cardinality input, no-spill pass-through (Tier 1)") { + val N = 8L << 20 + val benchmark = new Benchmark("adaptive partial agg, high card, no spill", N, + output = output) + addCodegenAdaptiveCases(benchmark, () => distinctKeyedDf(N)) + benchmark.run() + } + + // 1000 distinct keys over a large input: partial aggregation reduces a lot, the no-spill tier + // never fires, and the two runs must match (no regression). + runBenchmark("low-cardinality input, no-spill pass-through (Tier 1)") { + val N = 16L << 20 + val benchmark = new Benchmark("adaptive partial agg, low card, no spill", N, + output = output) + addCodegenAdaptiveCases(benchmark, () => + spark.range(N).selectExpr("id % 1000 as k", "id as v").groupBy("k").agg("v" -> "sum")) + benchmark.run() + } + + // Force the regular map to spill quickly and disable the no-spill tier (huge sample). With + // fully distinct keys the reduction ratio is 1.0, so at the spill boundary the on-spill + // tier (Tier 2) bypasses instead of spilling; the baseline spills repeatedly and falls back + // to sort-based aggregation. + runBenchmark("high-cardinality input, on-spill pass-through (Tier 2)") { + val N = 8L << 20 + val benchmark = new Benchmark("adaptive partial agg, high card, spill", N, output = output) + val tier2Conf = Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_ROWS.key -> Int.MaxValue.toString, + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 1048576") + addCodegenAdaptiveCases(benchmark, () => distinctKeyedDf(N), extraConf = tier2Conf) + benchmark.run() + } + + // Force the regular map to spill quickly on low-cardinality input. The reduction ratio is tiny + // (1000 distinct keys), so even at the spill boundary the on-spill tier correctly does not + // bypass: both runs spill and fall back to sort-based aggregation identically (no regression). + runBenchmark("low-cardinality input, on-spill pass-through (Tier 2)") { + val N = 16L << 20 + val benchmark = new Benchmark("adaptive partial agg, low card, spill", N, output = output) + val tier2Conf = Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_ROWS.key -> Int.MaxValue.toString, + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 1048576") + addCodegenAdaptiveCases(benchmark, () => + spark.range(N).selectExpr("id % 1000 as k", "id as v").groupBy("k").agg("v" -> "sum"), + extraConf = tier2Conf) + benchmark.run() + } + } + + private def distinctKeyedDf(N: Long): DataFrame = + spark.range(N).selectExpr("id as k", "id as v").groupBy("k").agg("v" -> "sum") +}