From a62ca30775cdc2d8c8113365daa5e2fc165545af Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Tue, 4 Aug 2026 13:38:33 +0800 Subject: [PATCH 1/5] [SPARK-58511][SQL] Bypass ineffective pre-shuffle partial aggregation at runtime When a pre-shuffle partial aggregation is not reducing rows (the distinct-key ratio is too high), maintaining an aggregation map is not worthwhile. This change makes hash aggregation detect that at runtime and bypass the partial aggregation: the remaining input rows are passed through as single-row partial buffers that the downstream Final aggregation merges, avoiding the cost of maintaining and spilling a large map. Two decision tiers, both evaluated only while the regular (second-level) map is still in memory: - no-spill tier: from a sample of rows on, bypass if the reduction ratio is at least `noSpillReductionRatioThreshold`. The sampling window doubles after each sub-threshold check, so low-cardinality input is re-checked only rarely while a late high-cardinality tail can still be caught. - on-spill tier: when the map is about to spill, bypass instead if the ratio is at least `spillReductionRatioThreshold`. Only pre-shuffle `Partial` hash aggregation with grouping keys is eligible. Both the codegen path (HashAggregateExec) and the interpreted path (TungstenAggregationIterator) are covered. Once pass-through is active the maps are frozen, so they are output (and their memory released) before the remaining input is streamed. New configs under spark.sql.execution.aggregate.adaptivePartialAggregation.* (enabled by default), a `numBypassingRows` metric, tests, and a benchmark are added. Co-Authored-By: Claude --- .../apache/spark/sql/internal/SQLConf.scala | 72 ++ ...tialAggregationBenchmark-jdk21-results.txt | 56 ++ ...tialAggregationBenchmark-jdk25-results.txt | 56 ++ ...ivePartialAggregationBenchmark-results.txt | 56 ++ .../UnsafeFixedWidthAggregationMap.java | 8 + .../aggregate/HashAggregateExec.scala | 449 ++++++++-- .../TungstenAggregationIterator.scala | 130 ++- .../AdaptivePartialAggregationSuite.scala | 770 ++++++++++++++++++ .../AdaptivePartialAggregationBenchmark.scala | 137 ++++ 9 files changed, 1664 insertions(+), 70 deletions(-) create mode 100644 sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk21-results.txt create mode 100644 sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk25-results.txt create mode 100644 sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/AdaptivePartialAggregationBenchmark.scala 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..83cba3e809dc4 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,66 @@ 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.95) + + 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. This " + + "threshold is more aggressive (lower) than the no-spill tier because once the map " + + "spills, partial aggregation starts paying disk I/O costs, so it is worth bypassing " + + "with less reduction benefit.") + .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.8) + 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 +8963,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..3921da46fd699 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,122 @@ 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 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 + + 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) { + | $regularMapRowCountTerm += 1; + | $probeRegularMap + | if ($unsafeRowBuffer == null) { + | // The map is full and would spill. Pre-spill, decide whether to bypass instead. + | if ($sorterTerm == null && + | (double) $hashMapTerm.getNumKeys() >= + | $regularMapRowCountTerm * ${cfg.spillReductionRatioThreshold}D) { + | $adaptivePassThroughTerm = true; + | } else { + | $spillMap + | } + | } else 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 +979,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 +1138,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 +1217,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..4f1ce1068ed12 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,12 @@ class TungstenAggregationIterator( } } else { var i = 0 - while (inputIter.hasNext) { + var processedRows = 0L + // 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 +221,49 @@ 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 (adaptivePartialAggConfig.isDefined && externalSorter == null && processedRows > 0 && + hashMap.getNumKeys().toDouble >= + processedRows * adaptivePartialAggConfig.get.spillReductionRatioThreshold) { + 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 (adaptivePartialAggConfig.isDefined && externalSorter == null && + processedRows == nextSampleRow) { + if (hashMap.getNumKeys().toDouble >= + processedRows * adaptivePartialAggConfig.get.noSpillReductionRatioThreshold) { + passThrough = true + } else { + nextSampleRow = nextSampleRow * 2 + } } } - processRow(buffer, newInput) - i += 1 } if (externalSorter != null) { @@ -354,6 +404,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 +491,8 @@ class TungstenAggregationIterator( /////////////////////////////////////////////////////////////////////////// override final def hasNext: Boolean = { - (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) + (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) || + passThroughHasNext } override final def next(): UnsafeRow = { @@ -412,7 +506,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 +520,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..047ba28eecc3e --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala @@ -0,0 +1,770 @@ +/* + * 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)( + 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 + } + 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 ++ 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("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..25f139aed2439 --- /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 2) tier 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") +} From ce6e4dcd5dc73079f7d6745021e2d3043ac70322 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Tue, 4 Aug 2026 19:28:00 +0800 Subject: [PATCH 2/5] Address review comments - Make `spillReductionRatioThreshold` fall back to `noSpillReductionRatioThreshold` so both tiers apply one policy by default. - Evaluate the spill-tier ratio over the same row set in both execution paths: the codegen path now counts only the rows that made it into the map, matching the interpreted path, and skips the decision until at least one row is counted. Add a test asserting both paths decide identically at the exact ratio boundary. - Unwrap the adaptive config once before the interpreted input loop instead of re-checking the `Option` for every row. - Fix two comment typos. Co-Authored-By: Claude --- docs/sql-performance-tuning.md | 61 +++++++++++++++++ .../apache/spark/sql/internal/SQLConf.scala | 16 ++--- .../aggregate/HashAggregateExec.scala | 24 ++++--- .../TungstenAggregationIterator.scala | 18 +++-- .../AdaptivePartialAggregationSuite.scala | 67 ++++++++++++++++++- .../AdaptivePartialAggregationBenchmark.scala | 2 +- 6 files changed, 158 insertions(+), 30 deletions(-) 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 83cba3e809dc4..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 @@ -4197,8 +4197,8 @@ object SQLConf { .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.95) + .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." + @@ -4206,15 +4206,13 @@ object SQLConf { .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. This " + - "threshold is more aggressive (lower) than the no-spill tier because once the map " + - "spills, partial aggregation starts paying disk I/O costs, so it is worth bypassing " + - "with less reduction benefit.") + "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) - .doubleConf - .checkValue(v => v > 0.0 && v <= 1.0, "The reduction ratio threshold must be in (0.0, 1.0].") - .createWithDefault(0.8) + .fallbackConf(ADAPTIVE_PARTIAL_AGGREGATION_NO_SPILL_REDUCTION_RATIO_THRESHOLD) val JSON_GENERATOR_IGNORE_NULL_FIELDS = buildConf("spark.sql.jsonGenerator.ignoreNullFields") 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 3921da46fd699..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 @@ -879,7 +879,7 @@ case class HashAggregateExec( | $sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter()); |} |$resetCounter - |// the hash map had be spilled, it should have enough memory now, + |// the hash map had been spilled, so it should have enough memory now, |// try to allocate buffer again. |$unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow( | $unsafeRowKeys, $unsafeRowKeyHash); @@ -910,24 +910,26 @@ case class HashAggregateExec( |// generate grouping key |${unsafeRowKeyCode.code} |if (!$adaptivePassThroughTerm) { - | $regularMapRowCountTerm += 1; | $probeRegularMap | if ($unsafeRowBuffer == null) { - | // The map is full and would spill. Pre-spill, decide whether to bypass instead. - | if ($sorterTerm == null && + | if ($sorterTerm == null && $regularMapRowCountTerm > 0 && | (double) $hashMapTerm.getNumKeys() >= | $regularMapRowCountTerm * ${cfg.spillReductionRatioThreshold}D) { | $adaptivePassThroughTerm = true; | } else { | $spillMap | } - | } else if ($sorterTerm == null && - | $regularMapRowCountTerm == $adaptiveNextSampleRowTerm) { - | if ((double) $hashMapTerm.getNumKeys() >= - | $regularMapRowCountTerm * ${cfg.noSpillReductionRatioThreshold}D) { - | $adaptivePassThroughTerm = true; - | } else { - | $adaptiveNextSampleRowTerm = $adaptiveNextSampleRowTerm * 2; + | } + | if ($unsafeRowBuffer != null) { + | $regularMapRowCountTerm += 1; + | if ($sorterTerm == null && + | $regularMapRowCountTerm == $adaptiveNextSampleRowTerm) { + | if ((double) $hashMapTerm.getNumKeys() >= + | $regularMapRowCountTerm * ${cfg.noSpillReductionRatioThreshold}D) { + | $adaptivePassThroughTerm = true; + | } else { + | $adaptiveNextSampleRowTerm = $adaptiveNextSampleRowTerm * 2; + | } | } | } |} 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 4f1ce1068ed12..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 @@ -209,6 +209,13 @@ class TungstenAggregationIterator( } else { var i = 0 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. @@ -224,9 +231,8 @@ class TungstenAggregationIterator( // 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 (adaptivePartialAggConfig.isDefined && externalSorter == null && processedRows > 0 && - hashMap.getNumKeys().toDouble >= - processedRows * adaptivePartialAggConfig.get.spillReductionRatioThreshold) { + 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`. @@ -254,10 +260,8 @@ class TungstenAggregationIterator( // 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 (adaptivePartialAggConfig.isDefined && externalSorter == null && - processedRows == nextSampleRow) { - if (hashMap.getNumKeys().toDouble >= - processedRows * adaptivePartialAggConfig.get.noSpillReductionRatioThreshold) { + if (adaptiveEnabled && externalSorter == null && processedRows == nextSampleRow) { + if (hashMap.getNumKeys().toDouble >= processedRows * noSpillRatioThreshold) { passThrough = true } else { nextSampleRow = nextSampleRow * 2 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 index 047ba28eecc3e..05907997a419d 100644 --- 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 @@ -179,7 +179,10 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession * 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)( + private def forEachCodegenAndMap( + sampleRows: Int = 8, + regularFallback: Int = 0, + noSpillThreshold: Double = -1.0)( body: String => Unit): Unit = { for { wholeStage <- Seq(true, false) @@ -194,13 +197,20 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } 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 ++ fixedPlanConfs): _*) { + fallbackConf ++ thresholdConf ++ fixedPlanConfs): _*) { body(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") } } @@ -713,6 +723,59 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } } + 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. 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 index 25f139aed2439..2b33c37ce781f 100644 --- 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 @@ -104,7 +104,7 @@ object AdaptivePartialAggregationBenchmark extends SqlBasedBenchmark { // 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 2) tier bypasses instead of spilling; the baseline spills repeatedly and falls back + // 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 From e1bbdd717e072d5d5d8ef8d702f4cc40a55657ce Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Wed, 5 Aug 2026 13:54:40 +0800 Subject: [PATCH 3/5] Simplify the adaptive partial aggregation policy Collapse the four tuning configs into `minRows` and `minCompaction`, following the contract cloud-fan described: measure the compaction ratio at the operator level over all processed rows, and re-evaluate it in a fresh epoch after each spill. - Drop `AdaptivePartialAggregationConfig` now that only two values are threaded through, passing them to `TungstenAggregationIterator` as plain fields with no default values. - Make `minRows` a `longConf` so a large value genuinely disables the periodic check. - Reflow the comments and docs, and fix two bad line wraps. Co-Authored-By: Claude --- docs/sql-performance-tuning.md | 49 ++--- .../apache/spark/sql/internal/SQLConf.scala | 73 +++---- .../aggregate/RowBasedAggregateHashMap.java | 8 + .../aggregate/HashAggregateExec.scala | 182 ++++++++---------- .../TungstenAggregationIterator.scala | 92 +++++---- .../VectorizedHashMapGenerator.scala | 4 + .../AdaptivePartialAggregationSuite.scala | 144 +++++++++----- .../AdaptivePartialAggregationBenchmark.scala | 38 ++-- 8 files changed, 291 insertions(+), 299 deletions(-) diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index a1c778c4c10e9..51470e0179ae0 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -190,13 +190,13 @@ final aggregation after it. The partial aggregation is only worthwhile when it a 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. +When adaptive partial aggregation is enabled, hash aggregation measures the compaction ratio (the +number of processed rows divided by the number of keys held in its aggregation maps) at runtime +and, if the partial aggregation is not collapsing enough rows to be worthwhile, stops populating +the aggregation map and passes the remaining rows through as single-row partial aggregation buffers +for the final aggregation to merge. Query results are unchanged. The ratio is evaluated +periodically, and again right before the aggregation map would spill -- in which case the spill is +skipped entirely -- so a query that only becomes ineffective later in its input is still caught. @@ -208,37 +208,28 @@ input is still caught. 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. - + - + - + - - + + - - - - - - - +
Property NameDefaultMeaningSince Version
4.3.04.4.0
spark.sql.execution.aggregate.adaptivePartialAggregation.sampleRowsspark.sql.execution.aggregate.adaptivePartialAggregation.minRows 100000 - 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. + The number of rows to process before the compaction ratio is evaluated, so a decision is never + made on too few rows. The ratio is re-evaluated every time this many further rows have been + processed. 4.3.04.4.0
spark.sql.execution.aggregate.adaptivePartialAggregation.noSpillReductionRatioThreshold0.9spark.sql.execution.aggregate.adaptivePartialAggregation.minCompaction1.1 - 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). + The minimum compaction ratio required to keep the partial aggregation. A ratio of 10 means the + partial aggregation collapses ten rows into one key; when the observed ratio is below this + value the partial aggregation is bypassed for the rest of the input. A larger value bypasses + more aggressively. 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.04.4.0
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 6afafa87195e9..8b547ee63fb37 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 @@ -4164,55 +4164,37 @@ object SQLConf { "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") + .version("4.4.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") + val ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.minRows") + .doc("The number of rows to process before adaptive partial aggregation (see " + + s"'${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}') evaluates the compaction ratio. The " + + "ratio is evaluated once this many rows have been processed since the previous " + + "evaluation, so a decision is never made on too few rows.") + .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) - .intConf - .checkValue(_ > 0, "The sample row count must be positive.") + .longConf + .checkValue(_ > 0, "The minimum 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") + val ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.minCompaction") + .doc("The minimum compaction ratio required to keep the pre-shuffle partial aggregation " + + s"(see '${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}'). The compaction ratio is the " + + "number of processed rows divided by the number of keys held in the aggregation maps, " + + "so a ratio of 10 means the partial aggregation collapses ten rows into one. When the " + + s"ratio is below this value after '${ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key}' rows, " + + "or when the aggregation map is about to spill, the partial aggregation is bypassed for " + + "the rest of the input. A larger value bypasses more aggressively.") + .version("4.4.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) + .checkValue(_ >= 1.0, "The minimum compaction ratio must be at least 1.0.") + .createWithDefault(1.1) val JSON_GENERATOR_IGNORE_NULL_FIELDS = buildConf("spark.sql.jsonGenerator.ignoreNullFields") @@ -8964,14 +8946,11 @@ class SQLConf extends Serializable with Logging with SqlApiConf { 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 adaptivePartialAggregationMinRows: Long = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS) - def adaptivePartialAggregationSpillReductionRatioThreshold: Double = - getConf(ADAPTIVE_PARTIAL_AGGREGATION_SPILL_REDUCTION_RATIO_THRESHOLD) + def adaptivePartialAggregationMinCompaction: Double = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION) def objectAggSortBasedFallbackThreshold: Int = getConf(OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD) diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java b/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java index 3f67d050a8c8e..859a2fdf60d91 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java @@ -110,6 +110,14 @@ public final KVIterator rowIterator() { return batch.rowIterator(); } + /** + * Returns the number of distinct keys currently held by this map. Used by adaptive partial + * aggregation to measure the reduction ratio across both aggregation maps. + */ + public final int getNumKeys() { + return numRows; + } + @Override public final void close() { batch.close(); 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 6bc21ff2f9c76..41094eed14e03 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 @@ -126,7 +126,9 @@ case class HashAggregateExec( avgHashProbe, numTasksFallBacked, numBypassingRows, - adaptivePartialAggConfig) + adaptivePartialAggEnabled, + adaptiveMinRows, + adaptiveMinCompaction) if (!hasInput && groupingExpressions.isEmpty) { numOutputRows += 1 Iterator.single[UnsafeRow](aggregationIterator.outputForEmptyGroupingKeyWithoutInput()) @@ -147,16 +149,13 @@ case class HashAggregateExec( 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: + * Whether adaptive partial aggregation applies to this operator. When it does, the aggregation + * may bypass partial aggregation at runtime and pass the remaining input rows through as + * single-row partial buffers (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). It only + * applies to a pre-shuffle partial aggregation with grouping keys: * - `Partial` mode only: the downstream `Final` aggregation merges the passed-through - * single-row buffers, so the output contract is unchanged. `Final`/`Complete`/`PartialMerge` - * have no such downstream to fall back on. + * 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 @@ -164,8 +163,8 @@ case class HashAggregateExec( * 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 && + private val adaptivePartialAggEnabled: Boolean = { + conf.adaptivePartialAggregationEnabled && groupingExpressions.nonEmpty && // Only the pre-shuffle partial aggregation has a downstream `Final` to merge passed-through // single-row buffers. `requiredChildDistributionExpressions` is `None` exactly for that @@ -175,18 +174,13 @@ case class HashAggregateExec( // 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 number of rows between two compaction-ratio evaluations, and the ratio below which the + // partial aggregation is considered ineffective. Only read when the feature applies. + private val adaptiveMinRows: Long = conf.adaptivePartialAggregationMinRows + private val adaptiveMinCompaction: Double = conf.adaptivePartialAggregationMinCompaction + // The name for Fast HashMap private var fastHashMapTerm: String = _ private var isFastHashMapEnabled: Boolean = false @@ -200,14 +194,13 @@ 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. + // Codegen state for adaptive partial aggregation. When the aggregation maps stop collapsing + // enough rows, the operator stops populating them and instead streams each remaining row through + // as a single-row partial buffer for the Final aggregate to merge. The compaction ratio is + // measured at the operator level: all processed rows against the keys held by both the fast and + // the regular map, so two-level-map routing does not change the decision. private var adaptivePassThroughTerm: String = _ - private var regularMapRowCountTerm: String = _ + private var processedRowsTerm: 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 @@ -217,10 +210,10 @@ case class HashAggregateExec( // 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 processed-row count at which the compaction ratio is evaluated next. It advances by + // `minRows` after every check, and the count is reset after a spill so the new in-memory map + // epoch is judged on its own rows. + private var adaptiveNextCheckRowTerm: 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 = _ @@ -507,13 +500,13 @@ case class HashAggregateExec( protected override def doProduceWithKeys(ctx: CodegenContext): String = { val initAgg = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initAgg") - if (adaptivePartialAggConfig.isDefined) { + if (adaptivePartialAggEnabled) { 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;") + processedRowsTerm = ctx.addMutableState(CodeGenerator.JAVA_LONG, "processedRows") + adaptiveNextCheckRowTerm = + ctx.addMutableState(CodeGenerator.JAVA_LONG, "adaptiveNextCheckRow", + v => s"$v = ${adaptiveMinRows}L;") adaptiveChildrenConsumedTerm = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveChildrenConsumed") adaptiveMapOutputDoneTerm = @@ -631,7 +624,7 @@ case class HashAggregateExec( // 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) { + val postChildProduce = if (adaptivePartialAggEnabled) { s"$adaptiveChildrenConsumedTerm = true;" } else { finishHashMap @@ -719,7 +712,7 @@ case class HashAggregateExec( // `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 outputMapFuncName = if (adaptivePartialAggEnabled) { val name = ctx.freshName("outputMap") ctx.addNewFunction(name, s""" @@ -743,7 +736,7 @@ case class HashAggregateExec( // 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) { + val adaptiveStopCheck = if (adaptivePartialAggEnabled) { "if (shouldStop()) return;" } else { "" @@ -754,7 +747,7 @@ case class HashAggregateExec( // `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) { + val adaptiveOutputMap = if (adaptivePartialAggEnabled) { s""" |if (!$adaptiveMapOutputDoneTerm) { | $outputMapFuncName(); @@ -764,18 +757,21 @@ case class HashAggregateExec( } else { "" } - val adaptiveResumeBuild = if (adaptivePartialAggConfig.isDefined) { + val adaptiveResumeBuild = if (adaptivePartialAggEnabled) { + val beforeResumedAgg = ctx.freshName("beforeResumedAgg") s""" |if (!$adaptiveChildrenConsumedTerm) { | $adaptiveOutputMap + | long $beforeResumedAgg = System.nanoTime(); | $doAggFuncName(partitionIndex); + | $aggTime.add((System.nanoTime() - $beforeResumedAgg) / $NANOS_PER_MILLIS); | if (shouldStop()) return; |} """.stripMargin } else { "" } - val adaptiveFinalOutput = if (adaptivePartialAggConfig.isDefined) { + val adaptiveFinalOutput = if (adaptivePartialAggEnabled) { adaptiveOutputMap } else { s""" @@ -802,16 +798,16 @@ case class HashAggregateExec( // 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 + override def needStopCheck: Boolean = adaptivePartialAggEnabled // Blocking operators normally do not copy their result because every output row is drained (via // `shouldStop()`) before the next one is produced. Adaptive pass-through breaks that assumption: // when an `Expand` sits below, one input row fans out into several pass-through rows that are all // appended in the same child loop iteration before any drain, and they all alias the single - // result `UnsafeRow`. Copy the result so the buffered rows do not collapse into the last one. - override def needCopyResult: Boolean = - adaptivePartialAggConfig.isDefined || super.needCopyResult + // result `UnsafeRow`. Such children report `needCopyResult` themselves, so propagate their + // requirement rather than copying for every adaptive aggregate. + override def needCopyResult: Boolean = adaptivePartialAggEnabled && + child.asInstanceOf[CodegenSupport].needCopyResult protected override def doConsumeWithKeys(ctx: CodegenContext, input: Seq[ExprCode]): String = { // create grouping key @@ -828,7 +824,7 @@ case class HashAggregateExec( // 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) { + val emptyAggBufferCode = if (adaptivePartialAggEnabled) { GenerateUnsafeProjection.createCode(ctx, declFunctions.flatMap(f => f.initialValues)) } else { null @@ -889,46 +885,45 @@ case class HashAggregateExec( |} """.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. + if (adaptivePartialAggEnabled) { + // The compaction ratio is measured at the operator level: all processed rows against the + // keys held by both maps, so two-level-map routing does not change the decision. The same + // predicate decides both check points -- periodically every `minRows` rows, and right + // before the map would spill (in which case the spill is skipped entirely). + val totalKeys = if (isFastHashMapEnabled) { + s"($fastHashMapTerm.getNumKeys() + $hashMapTerm.getNumKeys())" + } else { + s"$hashMapTerm.getNumKeys()" + } + // After a spill the map starts a new in-memory epoch, so the counters restart and the + // ratio of that epoch alone decides whether the remaining rows are passed through. + val resetEpoch = + s""" + |$processedRowsTerm = 0; + |$adaptiveNextCheckRowTerm = ${adaptiveMinRows}L; + """.stripMargin + val ineffective = + s"$processedRowsTerm < (double) $totalKeys * ${adaptiveMinCompaction}D" s""" |// generate grouping key |${unsafeRowKeyCode.code} |if (!$adaptivePassThroughTerm) { | $probeRegularMap | if ($unsafeRowBuffer == null) { - | if ($sorterTerm == null && $regularMapRowCountTerm > 0 && - | (double) $hashMapTerm.getNumKeys() >= - | $regularMapRowCountTerm * ${cfg.spillReductionRatioThreshold}D) { + | if ($processedRowsTerm > 0 && $ineffective) { | $adaptivePassThroughTerm = true; | } else { | $spillMap + | $resetEpoch | } | } | if ($unsafeRowBuffer != null) { - | $regularMapRowCountTerm += 1; - | if ($sorterTerm == null && - | $regularMapRowCountTerm == $adaptiveNextSampleRowTerm) { - | if ((double) $hashMapTerm.getNumKeys() >= - | $regularMapRowCountTerm * ${cfg.noSpillReductionRatioThreshold}D) { + | $processedRowsTerm += 1; + | if ($processedRowsTerm == $adaptiveNextCheckRowTerm) { + | if ($ineffective) { | $adaptivePassThroughTerm = true; | } else { - | $adaptiveNextSampleRowTerm = $adaptiveNextSampleRowTerm * 2; + | $adaptiveNextCheckRowTerm += ${adaptiveMinRows}L; | } | } | } @@ -962,7 +957,7 @@ case class HashAggregateExec( | ${fastRowKeys.map(_.value).mkString(", ")}); |} """.stripMargin - val guardedFastMapProbe = if (adaptivePartialAggConfig.isDefined) { + val guardedFastMapProbe = if (adaptivePartialAggEnabled) { s""" |if (!$adaptivePassThroughTerm) { | $fastMapProbe @@ -984,10 +979,10 @@ case class HashAggregateExec( // 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) { + // (guarded above), so both buffers stay null. The periodic-check 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 (adaptivePartialAggEnabled) { // 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. @@ -1154,7 +1149,7 @@ case class HashAggregateExec( } else { s"UnsafeRow $unsafeRowBuffer = null;" } - val declareBypassed = if (adaptivePartialAggConfig.isDefined) { + val declareBypassed = if (adaptivePartialAggEnabled) { s"boolean $adaptiveRowBypassedTerm = false;" } else { "" @@ -1173,7 +1168,7 @@ case class HashAggregateExec( // 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 emitPassThroughRow = if (adaptivePartialAggEnabled) { val numBypassingRows = metricTerm(ctx, "numBypassingRows") s""" |if ($adaptiveRowBypassedTerm) { @@ -1219,24 +1214,3 @@ 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 450578984d6b1..4464f05ed87ce 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 @@ -96,8 +96,10 @@ class TungstenAggregationIterator( spillSize: SQLMetric, avgHashProbe: SQLMetric, numTasksFallBacked: SQLMetric, - numBypassingRows: SQLMetric = null, - adaptivePartialAggConfig: Option[AdaptivePartialAggregationConfig] = None) + numBypassingRows: SQLMetric, + adaptivePartialAggEnabled: Boolean, + adaptiveMinRows: Long, + adaptiveMinCompaction: Double) extends AggregationIterator( partIndex, groupingExpressions, @@ -182,20 +184,19 @@ class TungstenAggregationIterator( // 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. + // When adaptive partial aggregation is enabled (`adaptivePartialAggEnabled`), 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 + // `nextPassThroughOutput`). One predicate decides both check points -- the aggregation is + // ineffective when `processedRows < distinctKeys * minCompaction`, i.e. it does not collapse + // `minCompaction` rows into one key: + // - periodically, every `minRows` processed rows; + // - when the map is about to spill, in which case the spill is skipped entirely and the row + // that could not be inserted becomes the first pass-through row. + // A spill starts a new in-memory map epoch: the row counters restart so that epoch is judged on + // its own rows, which lets an input whose cardinality only turns unfavorable later still be + // passed through. Pass-through may therefore coexist with earlier spills; the output stage + // drains the sort-based (or map) output first and streams the remaining rows afterwards. 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. @@ -209,17 +210,15 @@ class TungstenAggregationIterator( } else { var i = 0 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) + val minRows = adaptiveMinRows + // The processed-row count at which the compaction ratio is evaluated next. It advances by + // `minRows` after every check, and restarts after a spill so the new in-memory map epoch is + // judged on its own rows. + var nextCheckRow = minRows + // The partial aggregation is ineffective when it does not collapse `minCompaction` rows into + // one key. There is no fast map on this path, so the map's keys are all the operator holds. + def ineffective(): Boolean = + processedRows < hashMap.getNumKeys().toDouble * adaptiveMinCompaction while (inputIter.hasNext && !passThrough) { val newInput = inputIter.next() val groupingKey = groupingProjection.apply(newInput) @@ -228,11 +227,10 @@ class TungstenAggregationIterator( buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) } if (buffer == null) { - // 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) { + // The map is full and would normally 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 (adaptivePartialAggEnabled && processedRows > 0 && ineffective()) { 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`. @@ -245,6 +243,10 @@ class TungstenAggregationIterator( externalSorter.merge(sorter) } i = 0 + // The map starts a new in-memory epoch, so the counters restart and the ratio of that + // epoch alone decides whether the remaining rows are passed through. + processedRows = 0L + nextCheckRow = minRows buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) if (buffer == null) { // failed to allocate the first page @@ -256,15 +258,13 @@ class TungstenAggregationIterator( 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) { + // Periodic check: if the aggregation is not collapsing enough rows, bypass it for the + // rest of the input. + if (adaptivePartialAggEnabled && processedRows == nextCheckRow) { + if (ineffective()) { passThrough = true } else { - nextSampleRow = nextSampleRow * 2 + nextCheckRow += minRows } } } @@ -413,12 +413,12 @@ class TungstenAggregationIterator( /////////////////////////////////////////////////////////////////////////// // 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. + // passed through as single-row partial buffers. Set in `processInputs` by either check point. + // It may coexist with earlier spills, so the output order is: sort-based (or map) output first, + // then the pass-through rows. private[this] var passThrough: Boolean = false - // The row that could not be inserted at the on-spill tier trigger point. It is stashed here (as + // The row that could not be inserted at the spill check. It is stashed here (as // a copy) so it becomes the first pass-through row rather than being lost. private[this] var pendingPassThroughRow: InternalRow = null @@ -445,9 +445,7 @@ class TungstenAggregationIterator( // Reset the buffer to initial values, then update it with this single row. passThroughAggregationBuffer.copyFrom(initialAggregationBuffer) processRow(passThroughAggregationBuffer, row) - if (numBypassingRows != null) { - numBypassingRows += 1 - } + numBypassingRows += 1 generateOutput(groupingKey, passThroughAggregationBuffer) } @@ -501,7 +499,7 @@ class TungstenAggregationIterator( override final def next(): UnsafeRow = { if (hasNext) { - val res = if (sortBased) { + val res = if (sortBased && sortedInputHasNewGroup) { // Process the current group. processCurrentSortedGroup() // Generate output row for the current group. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala index f9c4ecc14e6c7..e1f9682439a31 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala @@ -83,6 +83,10 @@ class VectorizedHashMapGenerator( | buckets = new int[numBuckets]; | java.util.Arrays.fill(buckets, -1); | } + | + | public int getNumKeys() { + | return numRows; + | } """.stripMargin } 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 index 05907997a419d..7a9eb943eed25 100644 --- 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 @@ -37,7 +37,7 @@ import org.apache.spark.sql.test.SharedSparkSession * 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. + * the feature switch and eligibility rules are honored, and both check points work. */ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper { @@ -45,7 +45,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession 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 + // map fall back (spill) periodically, exercising the spill-check decision path in both the // codegen and interpreted aggregation paths. Kept moderate so low-cardinality inputs (which are // never bypassed and therefore really spill) do not open an unbounded number of spill readers. private val forceSpillFallback = "4, 16" @@ -85,8 +85,8 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession 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") ++ + // Small sample so the no-spill (the periodic check) path triggers on modest inputs. + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "8") ++ spillConf ++ fixedPlanConfs): _*) { val msg = s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap forceSpill=$forceSpill" withClue(msg) { @@ -106,10 +106,12 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession * 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. + * counter, so a forced fallback (or downstream shuffle-write spill) can inflate it. + * Asserted only by the periodic check test, which forces no fallback; use + * `tasksFallBacked` otherwise. * - `tasksFallBacked`: the partial aggregate's `numTasksFallBacked`, incremented only when the - * regular map actually falls back into sort-based aggregation. When Tier 2 bypasses at the + * regular map actually falls back into sort-based aggregation. When the spill check bypasses + * at the * spill boundary the sorter is never created, so this stays 0 -- direct, per-operator * evidence the bypass replaced the sort fallback. */ @@ -177,12 +179,12 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession * 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. + * regular map to spill (for the spill check); when 0 the regular map does not spill. */ private def forEachCodegenAndMap( - sampleRows: Int = 8, + minRows: Long = 8, regularFallback: Int = 0, - noSpillThreshold: Double = -1.0)( + minCompaction: Double = -1.0)( body: String => Unit): Unit = { for { wholeStage <- Seq(true, false) @@ -198,9 +200,8 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession 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) + val thresholdConf = if (minCompaction >= 0.0) { + Seq(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key -> minCompaction.toString) } else { Nil } @@ -209,7 +210,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession 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) ++ + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> minRows.toString) ++ fallbackConf ++ thresholdConf ++ fixedPlanConfs): _*) { body(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") } @@ -222,7 +223,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession 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. + // bypassed by the periodic check. checkAdaptiveMatchesReference { () => spark.range(0, 200, 1, 1) .select($"id" as "k", ($"id" * 2) as "v") @@ -359,7 +360,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession // 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. + // periodic check fires and the regression would show up. checkAdaptiveMatchesReference { () => spark.range(0, 1000, 1, 1) .select(($"id" % 10) as "c") @@ -368,7 +369,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } 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 + // A larger sample lets the map accumulate many keys before the periodic check bypasses, so the // early map output (which also frees the map) spans several drain cycles and re-enters the // map-output function; the results must still match the feature-off reference. val query = () => spark.range(0, 400000, 1, 1) @@ -378,7 +379,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession withSQLConf( (Seq( SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", - SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_SAMPLE_ROWS.key -> "200000", + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "200000", SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> "false") ++ fixedPlanConfs): _*) { val reference = withSQLConf( SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { @@ -651,8 +652,8 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } } - 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 + test("the periodic check fires when the sample shows no reduction, without spilling") { + // No forced regular-map spill: only the periodic check can trigger the bypass. Fully // distinct keys over a small sample cross `noSpillReductionRatioThreshold`, so rows bypass and // the regular map never spills. forEachCodegenAndMap() { clue => @@ -662,18 +663,20 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession .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") + assert(c.skipped > 0, + "the periodic check should bypass fully distinct input under the sample") + assert(c.spillBytes == 0, "the periodic check must decide before any spill happens") + assert(c.tasksFallBacked == 0, + "the periodic check must not fall back to sort-based aggregation") } } } - test("Tier 2 (on-spill) fires when the map would spill on high-cardinality input") { + test("the spill check 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 => + // point should bypass via the spill check rather than spilling. Use a sample larger than the + // input so the periodic check cannot fire first and the spill check is the one exercised. + forEachCodegenAndMap(minRows = 100000, regularFallback = 16) { clue => val df = () => spark.range(0, 200, 1, 1) .select($"id" as "k", $"id" as "v") .groupBy($"k") @@ -681,21 +684,57 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession 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 + "the spill check should bypass high-cardinality input at the spill boundary") + // The whole point of the spill check 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") + assert(c.tasksFallBacked == 0, + "the spill check must replace the sort fallback, not trigger it") + } + } + } + + test("a new in-memory map epoch after a spill can still bypass") { + // A spill starts a new in-memory map epoch and restarts the row counters, so an input whose + // cardinality only turns unfavorable after an early spill is still caught. The first 100 rows + // repeat 5 keys and keep the aggregation effective while the forced fallback makes the map + // spill; the remaining 300 rows are fully distinct, so the new epoch is judged ineffective and + // the rest of the input is passed through. Both the spill and the bypass must be observable. + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + val fastCap = if (twoLevelMap) 4 else 1 + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "8", + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> s"$fastCap, 40") ++ + fixedPlanConfs): _*) { + val df = () => spark.range(0, 400, 1, 1) + .select(when($"id" < 100, $"id" % 5).otherwise($"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.tasksFallBacked > 0, + "the low-cardinality prefix should still spill and fall back to sort") + assert(c.skipped > 0, + "the high-cardinality tail after the spill should be passed through") + } } } } 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 + // Sanity check for the the spill check 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 + // aggregation. This proves the spill check's `tasksFallBacked == 0` reflects the bypass and + // not merely // an input that never reached the spill boundary. for { wholeStage <- Seq(true, false) @@ -723,35 +762,34 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } } - 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) => + test("the spill check decides identically at the exact ratio boundary with codegen on and off") { + // The check before a spill 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 keys when the + // map fills at 50 aggregated rows, i.e. a compaction ratio of exactly 1.25: demanding 1.25 the + // aggregation is kept (`50 < 40 * 1.25` is false), demanding 1.3 it is bypassed. + Seq(1.25 -> false, 1.3 -> true).foreach { case (minCompaction, 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, + // A `minRows` larger than the input keeps the periodic check out of the picture. + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "100000", + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key -> minCompaction.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") { + withClue(s"minCompaction=$minCompaction wholeStage=$wholeStage") { numBypassingRows(df) } } } - withClue(s"threshold=$threshold skipped=$skippedPerCodegen") { + withClue(s"minCompaction=$minCompaction skipped=$skippedPerCodegen") { assert(skippedPerCodegen.forall(_ > 0) == shouldBypass, s"expected bypass=$shouldBypass at the ratio boundary") assert(skippedPerCodegen.map(_ > 0).distinct.length == 1, @@ -760,26 +798,26 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } } - 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 => + test("a very high minCompaction always bypasses once minRows has been processed") { + // Demanding an unreachable compaction ratio is the most aggressive setting: even a + // low-cardinality input that the default threshold keeps aggregating is bypassed at the first + // check point. The results must still match the feature-off reference. + forEachCodegenAndMap(minCompaction = 1000000.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") + "an unreachable compaction ratio 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 + // With a sample larger than the whole input and no regular-map spill forced, the periodic check // point is never reached, so nothing bypasses even though the keys are fully distinct. - forEachCodegenAndMap(sampleRows = 100000) { clue => + forEachCodegenAndMap(minRows = 100000) { clue => val df = () => spark.range(0, 200, 1, 1) .select($"id" as "k", $"id" as "v") .groupBy($"k") 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 index 2b33c37ce781f..489b0594e7505 100644 --- 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 @@ -31,11 +31,11 @@ import org.apache.spark.sql.internal.SQLConf * 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. + * - high-cardinality, no spill: the periodic check 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, + * - high-cardinality, forced regular-map spill: the spill check bypasses instead of spilling, * which should win. - * - low-cardinality, forced regular-map spill: the ratio is too low for the on-spill tier to + * - low-cardinality, forced regular-map spill: the ratio is too low for the spill check to * bypass, so both runs spill identically (no regression). * * To run this benchmark: @@ -81,9 +81,9 @@ object AdaptivePartialAggregationBenchmark extends SqlBasedBenchmark { } } - // 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)") { + // Fully distinct keys make partial aggregation useless, so the periodic check bypasses: the + // feature should be faster than the baseline that maintains a map entry per row. + runBenchmark("high-cardinality input, pass-through at the periodic check") { val N = 8L << 20 val benchmark = new Benchmark("adaptive partial agg, high card, no spill", N, output = output) @@ -91,9 +91,9 @@ object AdaptivePartialAggregationBenchmark extends SqlBasedBenchmark { benchmark.run() } - // 1000 distinct keys over a large input: partial aggregation reduces a lot, the no-spill tier + // 1000 distinct keys over a large input: partial aggregation reduces a lot, the periodic check // never fires, and the two runs must match (no regression). - runBenchmark("low-cardinality input, no-spill pass-through (Tier 1)") { + runBenchmark("low-cardinality input, pass-through at the periodic check") { val N = 16L << 20 val benchmark = new Benchmark("adaptive partial agg, low card, no spill", N, output = output) @@ -102,32 +102,32 @@ object AdaptivePartialAggregationBenchmark extends SqlBasedBenchmark { benchmark.run() } - // Force the regular map to spill quickly and disable the no-spill tier (huge sample). With + // Force the regular map to spill quickly and disable the periodic check (huge minRows). 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 + // spill check 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)") { + runBenchmark("high-cardinality input, pass-through at the spill check") { 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, + val spillCheckConf = Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> Long.MaxValue.toString, "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 1048576") - addCodegenAdaptiveCases(benchmark, () => distinctKeyedDf(N), extraConf = tier2Conf) + addCodegenAdaptiveCases(benchmark, () => distinctKeyedDf(N), extraConf = spillCheckConf) 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 + // (1000 distinct keys), so even at the spill boundary the spill check 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)") { + runBenchmark("low-cardinality input, pass-through at the spill check") { 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, + val spillCheckConf = Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> Long.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) + extraConf = spillCheckConf) benchmark.run() } } From d95d6c52655f88ce0e2010ca3aba4589198d04b7 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Wed, 5 Aug 2026 18:04:47 +0800 Subject: [PATCH 4/5] Count fast-map rows in the compaction ratio, allow minRows = 0 Address the remaining item and the design comment from the latest review. The generated path incremented `processedRows` only inside the regular-map branch, while the denominator counted the keys of both maps. Rows the fast map absorbed were therefore missing from the numerator, so a hot-key input the aggregation was collapsing heavily could measure as ineffective and be bypassed. Move the increment and the periodic check after fast/regular map routing, keyed on either buffer being set, so both sides of the ratio cover the same rows. The row that fails to insert at the spill boundary holds no buffer and stays excluded. The interpreted path has no fast map and was already consistent. Allow `minRows = 0` to disable the periodic check while keeping the check made at the spill boundary, so a spill-only mode is expressible. Co-Authored-By: Claude --- docs/sql-performance-tuning.md | 3 +- .../apache/spark/sql/internal/SQLConf.scala | 6 +- .../aggregate/HashAggregateExec.scala | 95 ++++++++++++------- .../TungstenAggregationIterator.scala | 4 +- .../AdaptivePartialAggregationSuite.scala | 66 ++++++++++++- .../AdaptivePartialAggregationBenchmark.scala | 2 +- 6 files changed, 131 insertions(+), 45 deletions(-) diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index 51470e0179ae0..c7d6288f0ad2f 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -216,7 +216,8 @@ skipped entirely -- so a query that only becomes ineffective later in its input The number of rows to process before the compaction ratio is evaluated, so a decision is never made on too few rows. The ratio is re-evaluated every time this many further rows have been - processed. + processed. Setting this to 0 disables the periodic evaluation entirely, leaving + only the check made when the aggregation map is about to spill. 4.4.0 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 8b547ee63fb37..6d39fdfdc604e 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 @@ -4174,11 +4174,13 @@ object SQLConf { .doc("The number of rows to process before adaptive partial aggregation (see " + s"'${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}') evaluates the compaction ratio. The " + "ratio is evaluated once this many rows have been processed since the previous " + - "evaluation, so a decision is never made on too few rows.") + "evaluation, so a decision is never made on too few rows. A value of 0 disables the " + + "periodic evaluation entirely, leaving only the check made when the aggregation map is " + + "about to spill.") .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .longConf - .checkValue(_ > 0, "The minimum row count must be positive.") + .checkValue(_ >= 0, "The minimum row count must not be negative.") .createWithDefault(100000) val ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION = 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 41094eed14e03..1e8ae555fba5e 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 @@ -851,6 +851,27 @@ case class HashAggregateExec( case _ => ("true", "", "") } + // The compaction ratio is measured at the operator level: all processed rows against the keys + // held by both maps, so two-level-map routing does not change the decision. The same predicate + // decides both check points -- periodically every `minRows` rows, and right before the map + // would spill (in which case the spill is skipped entirely). `minRows = 0` disables the + // periodic check: the row count is only ever compared after being incremented past 0, so it + // never matches and only the spill check remains. + val adaptiveTotalKeys = if (isFastHashMapEnabled) { + s"($fastHashMapTerm.getNumKeys() + $hashMapTerm.getNumKeys())" + } else { + s"$hashMapTerm.getNumKeys()" + } + val adaptiveIneffective = + s"$processedRowsTerm < (double) $adaptiveTotalKeys * ${adaptiveMinCompaction}D" + // After a spill the map starts a new in-memory epoch, so the counters restart and the ratio of + // that epoch alone decides whether the remaining rows are passed through. + val adaptiveResetEpoch = + s""" + |$processedRowsTerm = 0; + |$adaptiveNextCheckRowTerm = ${adaptiveMinRows}L; + """.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 @@ -886,45 +907,17 @@ case class HashAggregateExec( """.stripMargin if (adaptivePartialAggEnabled) { - // The compaction ratio is measured at the operator level: all processed rows against the - // keys held by both maps, so two-level-map routing does not change the decision. The same - // predicate decides both check points -- periodically every `minRows` rows, and right - // before the map would spill (in which case the spill is skipped entirely). - val totalKeys = if (isFastHashMapEnabled) { - s"($fastHashMapTerm.getNumKeys() + $hashMapTerm.getNumKeys())" - } else { - s"$hashMapTerm.getNumKeys()" - } - // After a spill the map starts a new in-memory epoch, so the counters restart and the - // ratio of that epoch alone decides whether the remaining rows are passed through. - val resetEpoch = - s""" - |$processedRowsTerm = 0; - |$adaptiveNextCheckRowTerm = ${adaptiveMinRows}L; - """.stripMargin - val ineffective = - s"$processedRowsTerm < (double) $totalKeys * ${adaptiveMinCompaction}D" s""" |// generate grouping key |${unsafeRowKeyCode.code} |if (!$adaptivePassThroughTerm) { | $probeRegularMap | if ($unsafeRowBuffer == null) { - | if ($processedRowsTerm > 0 && $ineffective) { + | if ($processedRowsTerm > 0 && $adaptiveIneffective) { | $adaptivePassThroughTerm = true; | } else { | $spillMap - | $resetEpoch - | } - | } - | if ($unsafeRowBuffer != null) { - | $processedRowsTerm += 1; - | if ($processedRowsTerm == $adaptiveNextCheckRowTerm) { - | if ($ineffective) { - | $adaptivePassThroughTerm = true; - | } else { - | $adaptiveNextCheckRowTerm += ${adaptiveMinRows}L; - | } + | $adaptiveResetEpoch | } | } |} @@ -977,17 +970,26 @@ case class HashAggregateExec( findOrInsertRegularHashMap } + // True when an aggregation map accepted this row, from either map. The fast map serves a row + // without it ever reaching the regular map, so both buffers have to be consulted to tell an + // aggregated row from one that must be streamed through. + val heldByAMap = if (isFastHashMapEnabled) { + s"($fastRowBuffer != null || $unsafeRowBuffer != null)" + } else { + s"($unsafeRowBuffer != null)" + } + // 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 periodic-check 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. + // `rowBypassed` marks exactly those rows: both probes are skipped (guarded above), so + // neither buffer is set. A row a map did accept is excluded -- including the row that + // flipped pass-through at the periodic check, which is already aggregated in the map that + // took it and must not be re-emitted. val createPassThroughBuffer = if (adaptivePartialAggEnabled) { // 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) { + |if ($adaptivePassThroughTerm && !$heldByAMap) { | $adaptiveRowBypassedTerm = true; | ${emptyAggBufferCode.code} | $unsafeRowBuffer = ${emptyAggBufferCode.value}; @@ -997,8 +999,31 @@ case class HashAggregateExec( "" } + // Count every row an aggregation map accepted so the numerator matches the operator-level + // denominator. Counting inside the regular-map branch alone would drop the rows the fast map + // absorbed from the ratio and bypass an aggregation that is in fact reducing. The row that + // fails to insert at the spill boundary sets pass-through and holds no buffer, so it is + // excluded here and streamed instead. + val adaptivePeriodicCheck = if (adaptivePartialAggEnabled) { + s""" + |if (!$adaptivePassThroughTerm && $heldByAMap) { + | $processedRowsTerm += 1; + | if ($processedRowsTerm == $adaptiveNextCheckRowTerm) { + | if ($adaptiveIneffective) { + | $adaptivePassThroughTerm = true; + | } else { + | $adaptiveNextCheckRowTerm += ${adaptiveMinRows}L; + | } + | } + |} + """.stripMargin + } else { + "" + } + s""" |$findCode + |$adaptivePeriodicCheck |$createPassThroughBuffer """.stripMargin } 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 4464f05ed87ce..ebc7a4aa43eb6 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 @@ -213,7 +213,9 @@ class TungstenAggregationIterator( val minRows = adaptiveMinRows // The processed-row count at which the compaction ratio is evaluated next. It advances by // `minRows` after every check, and restarts after a spill so the new in-memory map epoch is - // judged on its own rows. + // judged on its own rows. `minRows = 0` disables the periodic check: the count is only ever + // compared after being incremented past 0, so it never matches and only the spill check + // below remains. var nextCheckRow = minRows // The partial aggregation is ineffective when it does not collapse `minCompaction` rows into // one key. There is no fast map on this path, so the map's keys are all the operator holds. 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 index 7a9eb943eed25..68c5979b0e98a 100644 --- 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 @@ -111,8 +111,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession * `tasksFallBacked` otherwise. * - `tasksFallBacked`: the partial aggregate's `numTasksFallBacked`, incremented only when the * regular map actually falls back into sort-based aggregation. When the spill check bypasses - * at the - * spill boundary the sorter is never created, so this stays 0 -- direct, per-operator + * 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( @@ -652,6 +651,29 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } } + test("rows absorbed by the fast map count toward the compaction ratio") { + // The compaction ratio is measured at the operator level, so the rows the fast map absorbs + // must count in the numerator just as its keys count in the denominator. This input is + // dominated by a hot-key prefix that the fast map serves without ever reaching the regular + // map, followed by a short distinct tail that does reach it. Counting only the regular map's + // traffic would see the tail alone -- a ratio near 1 -- and bypass an aggregation that is in + // fact collapsing rows heavily. + forEachCodegenAndMap() { clue => + // With the fast map on it holds 4 keys, so `k < 4` is served there and `k >= 4` falls + // through. 400 hot-key rows against 4 hot keys plus 20 tail keys is a ratio of about 17.5, + // far above the default 1.1, so nothing may bypass. + val df = () => spark.range(0, 420, 1, 1) + .select(when($"id" < 400, $"id" % 4).otherwise($"id" - 396) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(df) == 0, + "an aggregation collapsing ~17 rows per key must not bypass; fast-map hits are " + + "missing from the numerator if it does") + } + } + } + test("the periodic check fires when the sample shows no reduction, without spilling") { // No forced regular-map spill: only the periodic check can trigger the bypass. Fully // distinct keys over a small sample cross `noSpillReductionRatioThreshold`, so rows bypass and @@ -731,11 +753,10 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } test("without the feature the same input really does fall back to sort") { - // Sanity check for the the spill check assertion above: with adaptive disabled, the identical + // Sanity check for the spill check assertion above: with adaptive disabled, the identical // high-cardinality input under the same forced fallback genuinely falls back to sort-based // aggregation. This proves the spill check's `tasksFallBacked == 0` reflects the bypass and - // not merely - // an input that never reached the spill boundary. + // not merely an input that never reached the spill boundary. for { wholeStage <- Seq(true, false) twoLevelMap <- Seq(true, false) @@ -814,6 +835,41 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } } + test("minRows = 0 disables the periodic check but keeps the spill check") { + // `minRows = 0` is a sentinel for "never evaluate periodically". Without a forced regular-map + // spill there is no check point at all, so fully distinct input -- which the default settings + // bypass immediately -- must be aggregated all the way through. This is what proves the + // periodic check is genuinely off rather than merely deferred. + forEachCodegenAndMap(minRows = 0) { 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, + "minRows = 0 must disable the periodic check, so nothing may bypass without a spill") + } + } + } + + test("minRows = 0 still bypasses at the spill boundary") { + // The other half of the sentinel: with the periodic check off, the spill check alone still + // bypasses instead of paying the spill I/O, which is the spill-only operating mode. + forEachCodegenAndMap(minRows = 0, 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, + "the spill check must still bypass when the periodic check is disabled") + assert(c.tasksFallBacked == 0, + "the spill check must replace the sort fallback, not trigger it") + } + } + } + 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 periodic check // point is never reached, so nothing bypasses even though the keys are fully distinct. 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 index 489b0594e7505..fa8b9510cfe94 100644 --- 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 @@ -92,7 +92,7 @@ object AdaptivePartialAggregationBenchmark extends SqlBasedBenchmark { } // 1000 distinct keys over a large input: partial aggregation reduces a lot, the periodic check - // never fires, and the two runs must match (no regression). + // never activates pass-through, and the two runs must match (no regression). runBenchmark("low-cardinality input, pass-through at the periodic check") { val N = 16L << 20 val benchmark = new Benchmark("adaptive partial agg, low card, no spill", N, From f35a80c69abd963aad96cc04720ddf63ad14564e Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Wed, 5 Aug 2026 19:42:13 +0800 Subject: [PATCH 5/5] Guard the adaptive codegen fragments and merge the per-row block `adaptiveIneffective` and `adaptiveResetEpoch` read terms that are only named when the feature applies, so they interpolated `null` into the generated source otherwise. The strings were never consumed outside the adaptive branches, but nothing structural said so. Build them under the same guard as their use sites. Merge the periodic check and the pass-through buffer into one block. Their conditions are mutually exclusive on whether a map accepted the row, so branching on that once makes the exclusivity structural and evaluates the test a single time. Co-Authored-By: Claude --- .../aggregate/HashAggregateExec.scala | 90 +++++++++---------- 1 file changed, 44 insertions(+), 46 deletions(-) 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 1e8ae555fba5e..be8985eedba91 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 @@ -857,20 +857,27 @@ case class HashAggregateExec( // would spill (in which case the spill is skipped entirely). `minRows = 0` disables the // periodic check: the row count is only ever compared after being incremented past 0, so it // never matches and only the spill check remains. - val adaptiveTotalKeys = if (isFastHashMapEnabled) { - s"($fastHashMapTerm.getNumKeys() + $hashMapTerm.getNumKeys())" + val adaptiveIneffective = if (adaptivePartialAggEnabled) { + val totalKeys = if (isFastHashMapEnabled) { + s"($fastHashMapTerm.getNumKeys() + $hashMapTerm.getNumKeys())" + } else { + s"$hashMapTerm.getNumKeys()" + } + s"$processedRowsTerm < (double) $totalKeys * ${adaptiveMinCompaction}D" } else { - s"$hashMapTerm.getNumKeys()" + "" } - val adaptiveIneffective = - s"$processedRowsTerm < (double) $adaptiveTotalKeys * ${adaptiveMinCompaction}D" + // After a spill the map starts a new in-memory epoch, so the counters restart and the ratio of // that epoch alone decides whether the remaining rows are passed through. - val adaptiveResetEpoch = + val adaptiveResetEpoch = if (adaptivePartialAggEnabled) { s""" - |$processedRowsTerm = 0; + |$processedRowsTerm = 0L; |$adaptiveNextCheckRowTerm = ${adaptiveMinRows}L; """.stripMargin + } else { + "" + } val findOrInsertRegularHashMap: String = { // Assumes the grouping key projection (`unsafeRowKeyCode.code`) has already run for this row, @@ -970,26 +977,40 @@ case class HashAggregateExec( findOrInsertRegularHashMap } - // True when an aggregation map accepted this row, from either map. The fast map serves a row - // without it ever reaching the regular map, so both buffers have to be consulted to tell an - // aggregated row from one that must be streamed through. - val heldByAMap = if (isFastHashMapEnabled) { - s"($fastRowBuffer != null || $unsafeRowBuffer != null)" - } else { - s"($unsafeRowBuffer != null)" - } - - // When pass-through is active, a row that no map holds must be streamed through. - // `rowBypassed` marks exactly those rows: both probes are skipped (guarded above), so - // neither buffer is set. A row a map did accept is excluded -- including the row that - // flipped pass-through at the periodic check, which is already aggregated in the map that + // Every row is either accepted by an aggregation map or streamed through -- the fast map + // serves a row without it ever reaching the regular map, so both buffers are consulted to + // tell the two apart. + // + // An accepted row counts toward the compaction ratio, so the numerator matches the + // operator-level denominator. Counting inside the regular-map branch alone would drop the + // rows the fast map absorbed from the ratio and bypass an aggregation that is in fact + // reducing. A row no map holds is streamed once pass-through is active: both probes are + // skipped (guarded above), so neither buffer is set, and `rowBypassed` marks exactly those + // rows. The row that fails to insert at the spill boundary lands here too, while the row + // that merely flipped pass-through at the check point is already aggregated in the map that // took it and must not be re-emitted. - val createPassThroughBuffer = if (adaptivePartialAggEnabled) { + val countOrPassThroughRow = if (adaptivePartialAggEnabled) { + val heldByAMap = if (isFastHashMapEnabled) { + s"($fastRowBuffer != null || $unsafeRowBuffer != null)" + } else { + s"($unsafeRowBuffer != null)" + } // The grouping key was already projected in `findOrInsertRegularHashMap` // (`unsafeRowKeyCode.code`), so `unsafeRowKeyCode.value` holds this row's key. Only build // the single-row partial buffer here. s""" - |if ($adaptivePassThroughTerm && !$heldByAMap) { + |if ($heldByAMap) { + | if (!$adaptivePassThroughTerm) { + | $processedRowsTerm += 1; + | if ($processedRowsTerm == $adaptiveNextCheckRowTerm) { + | if ($adaptiveIneffective) { + | $adaptivePassThroughTerm = true; + | } else { + | $adaptiveNextCheckRowTerm += ${adaptiveMinRows}L; + | } + | } + | } + |} else if ($adaptivePassThroughTerm) { | $adaptiveRowBypassedTerm = true; | ${emptyAggBufferCode.code} | $unsafeRowBuffer = ${emptyAggBufferCode.value}; @@ -999,32 +1020,9 @@ case class HashAggregateExec( "" } - // Count every row an aggregation map accepted so the numerator matches the operator-level - // denominator. Counting inside the regular-map branch alone would drop the rows the fast map - // absorbed from the ratio and bypass an aggregation that is in fact reducing. The row that - // fails to insert at the spill boundary sets pass-through and holds no buffer, so it is - // excluded here and streamed instead. - val adaptivePeriodicCheck = if (adaptivePartialAggEnabled) { - s""" - |if (!$adaptivePassThroughTerm && $heldByAMap) { - | $processedRowsTerm += 1; - | if ($processedRowsTerm == $adaptiveNextCheckRowTerm) { - | if ($adaptiveIneffective) { - | $adaptivePassThroughTerm = true; - | } else { - | $adaptiveNextCheckRowTerm += ${adaptiveMinRows}L; - | } - | } - |} - """.stripMargin - } else { - "" - } - s""" |$findCode - |$adaptivePeriodicCheck - |$createPassThroughBuffer + |$countOrPassThroughRow """.stripMargin }