diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index 3ba3a6c749ca7..c7d6288f0ad2f 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -181,6 +181,59 @@ 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 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. + + + + + + + + + + + + + + + + + + + + + +
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.4.0
spark.sql.execution.aggregate.adaptivePartialAggregation.minRows100000 + 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. 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
spark.sql.execution.aggregate.adaptivePartialAggregation.minCompaction1.1 + 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.4.0
+ ## Optimizing the Join Strategy ### Automatically Broadcasting Joins diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index beb8d5ee14581..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 @@ -4156,6 +4156,48 @@ 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.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(true) + + 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. 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 not be negative.") + .createWithDefault(100000) + + 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(_ >= 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") .doc("Whether to ignore null fields when generating JSON objects in JSON data source and " + @@ -8903,6 +8945,15 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def bypassPartialAggregation: Boolean = getConf(BYPASS_PARTIAL_AGGREGATION) + def adaptivePartialAggregationEnabled: Boolean = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_ENABLED) + + def adaptivePartialAggregationMinRows: Long = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS) + + def adaptivePartialAggregationMinCompaction: Double = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION) + 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/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 62c4f896f2ee4..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 @@ -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,11 @@ case class HashAggregateExec( peakMemory, spillSize, avgHashProbe, - numTasksFallBacked) + numTasksFallBacked, + numBypassingRows, + adaptivePartialAggEnabled, + adaptiveMinRows, + adaptiveMinCompaction) if (!hasInput && groupingExpressions.isEmpty) { numOutputRows += 1 Iterator.single[UnsafeRow](aggregationIterator.outputForEmptyGroupingKeyWithoutInput()) @@ -141,6 +148,39 @@ case class HashAggregateExec( .map(_.asInstanceOf[DeclarativeAggregate]) private val bufferSchema = DataTypeUtils.fromAttributes(aggregateBufferAttributes) + /** + * Whether adaptive partial aggregation applies to this operator. When it does, the aggregation + * may bypass partial aggregation at runtime and pass the remaining input rows through as + * single-row partial buffers (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). It only + * applies to a pre-shuffle partial aggregation with grouping keys: + * - `Partial` mode only: the downstream `Final` aggregation merges the passed-through + * single-row buffers, so the output contract is unchanged. + * `Final`/`Complete`/`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 adaptivePartialAggEnabled: Boolean = { + conf.adaptivePartialAggregationEnabled && + groupingExpressions.nonEmpty && + // Only the pre-shuffle partial aggregation has a downstream `Final` to merge passed-through + // single-row buffers. `requiredChildDistributionExpressions` is `None` exactly for that + // pre-shuffle phase and `Some` for the `Final`/`Complete` phase. This check is what keeps a + // group-by-only aggregate (no aggregate functions, so an empty `aggregateExpressions`) from + // being admitted vacuously: `aggregateExpressions.forall(_.mode == Partial)` alone is true + // for the empty list, which would wrongly make the `Final` phase eligible as well. + requiredChildDistributionExpressions.isEmpty && + aggregateExpressions.forall(a => a.mode == Partial) + } + + // 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 @@ -154,6 +194,30 @@ case class HashAggregateExec( private var hashMapTerm: String = _ private var sorterTerm: String = _ + // 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 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 + // 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 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 = _ + /** * This is called by generated Java class, should be public. */ @@ -436,6 +500,20 @@ case class HashAggregateExec( protected override def doProduceWithKeys(ctx: CodegenContext): String = { val initAgg = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initAgg") + if (adaptivePartialAggEnabled) { + adaptivePassThroughTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptivePassThrough") + 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 = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveMapOutputDone") + adaptiveMapSetupDoneTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveMapSetupDone") + } if (conf.enableTwoLevelAggMap) { enableTwoLevelHashMap() } else if (conf.enableVectorizedHashMap) { @@ -535,19 +613,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 (adaptivePartialAggEnabled) { + 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 +704,81 @@ 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 (adaptivePartialAggEnabled) { + 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 (adaptivePartialAggEnabled) { + "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 (adaptivePartialAggEnabled) { + s""" + |if (!$adaptiveMapOutputDoneTerm) { + | $outputMapFuncName(); + | if (shouldStop()) return; + |} + """.stripMargin + } else { + "" + } + 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 (adaptivePartialAggEnabled) { + adaptiveOutputMap + } else { + s""" + |$outputFromFastHashMap + |$outputFromRegularHashMap + """.stripMargin + } s""" |if (!$initAgg) { | $initAgg = true; @@ -626,13 +788,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 = adaptivePartialAggEnabled + + // Blocking operators normally do not copy their result because every output row is drained (via + // `shouldStop()`) before the next one is produced. Adaptive pass-through breaks that assumption: + // when an `Expand` sits below, one input row fans out into several pass-through rows that are all + // appended in the same child loop iteration before any drain, and they all alias the single + // result `UnsafeRow`. Such children report `needCopyResult` themselves, so propagate their + // requirement rather than copying for every adaptive aggregate. + override def needCopyResult: Boolean = adaptivePartialAggEnabled && + child.asInstanceOf[CodegenSupport].needCopyResult + protected override def doConsumeWithKeys(ctx: CodegenContext, input: Seq[ExprCode]): String = { // create grouping key val unsafeRowKeyCode = GenerateUnsafeProjection.createCode( @@ -644,6 +820,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 (adaptivePartialAggEnabled) { + 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 +851,123 @@ case class HashAggregateExec( case _ => ("true", "", "") } - val findOrInsertRegularHashMap: String = + // 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 adaptiveIneffective = if (adaptivePartialAggEnabled) { + val totalKeys = if (isFastHashMapEnabled) { + s"($fastHashMapTerm.getNumKeys() + $hashMapTerm.getNumKeys())" + } else { + s"$hashMapTerm.getNumKeys()" + } + s"$processedRowsTerm < (double) $totalKeys * ${adaptiveMinCompaction}D" + } else { + "" + } + + // 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 = if (adaptivePartialAggEnabled) { 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(); - | } - |} + |$processedRowsTerm = 0L; + |$adaptiveNextCheckRowTerm = ${adaptiveMinRows}L; """.stripMargin + } else { + "" + } + + val findOrInsertRegularHashMap: String = { + // Assumes the grouping key projection (`unsafeRowKeyCode.code`) has already run for this row, + // so `unsafeRowKeyCode.value` holds the current key. The projection is emitted exactly once + // per regular-map row (see below); emitting it in more than one runtime branch is unsafe + // because the projection's subexpression/writer state assigned in one branch would be read + // stale from another (e.g. the adaptive pass-through path would reuse the last probed key). + val probeRegularMap = + s""" + |int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode(); + |if ($checkFallbackForBytesToBytesMap) { + | // try to get the buffer from hash map + | $unsafeRowBuffer = + | $hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys, $unsafeRowKeyHash); + |} + """.stripMargin + + val spillMap = + s""" + |if ($sorterTerm == null) { + | $sorterTerm = $hashMapTerm.destructAndCreateExternalSorter(); + |} else { + | $sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter()); + |} + |$resetCounter + |// the hash map had been spilled, so it should have enough memory now, + |// try to allocate buffer again. + |$unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow( + | $unsafeRowKeys, $unsafeRowKeyHash); + |if ($unsafeRowBuffer == null) { + | // failed to allocate the first page + | throw QueryExecutionErrors.aggregateOutOfMemoryError(); + |} + """.stripMargin + + if (adaptivePartialAggEnabled) { + s""" + |// generate grouping key + |${unsafeRowKeyCode.code} + |if (!$adaptivePassThroughTerm) { + | $probeRegularMap + | if ($unsafeRowBuffer == null) { + | if ($processedRowsTerm > 0 && $adaptiveIneffective) { + | $adaptivePassThroughTerm = true; + | } else { + | $spillMap + | $adaptiveResetEpoch + | } + | } + |} + """.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 (adaptivePartialAggEnabled) { + 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 +976,54 @@ case class HashAggregateExec( } else { findOrInsertRegularHashMap } + + // 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 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 ($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}; + |} + """.stripMargin + } else { + "" + } + + s""" + |$findCode + |$countOrPassThroughRow + """.stripMargin } val inputAttrs = aggregateBufferAttributes ++ inputAttributes @@ -845,29 +1158,56 @@ case class HashAggregateExec( } } - val declareRowBuffer: String = if (isFastHashMapEnabled) { - val fastRowType = if (isVectorizedHashMapEnabled) { - classOf[MutableColumnarRow].getName + val declareRowBuffer: String = { + val declareBuffers = if (isFastHashMapEnabled) { + val fastRowType = if (isVectorizedHashMapEnabled) { + classOf[MutableColumnarRow].getName + } else { + "UnsafeRow" + } + s""" + |UnsafeRow $unsafeRowBuffer = null; + |$fastRowType $fastRowBuffer = null; + """.stripMargin + } else { + s"UnsafeRow $unsafeRowBuffer = null;" + } + val declareBypassed = if (adaptivePartialAggEnabled) { + s"boolean $adaptiveRowBypassedTerm = false;" } else { - "UnsafeRow" + "" } s""" - |UnsafeRow $unsafeRowBuffer = null; - |$fastRowType $fastRowBuffer = null; + |$declareBuffers + |$declareBypassed """.stripMargin - } else { - s"UnsafeRow $unsafeRowBuffer = null;" } // We try to do hash map based in-memory aggregation first. If there is not enough memory (the // hash map will return null for new key), we spill the hash map to disk to free memory, then // continue to do in-memory aggregation and spilling until all the rows had been processed. // Finally, sort the spilled aggregate buffers by key, and merge them together for same key. + // + // With adaptive partial aggregation, once pass-through is active `updateRowInHashMap` fills the + // single-row buffer built above; we then emit `key ++ buffer` straight to the parent so the row + // skips both the fast map and the regular map. + val emitPassThroughRow = if (adaptivePartialAggEnabled) { + val numBypassingRows = metricTerm(ctx, "numBypassingRows") + s""" + |if ($adaptiveRowBypassedTerm) { + | $numBypassingRows.add(1); + | $outputFunc(${unsafeRowKeyCode.value}, $unsafeRowBuffer); + |} + """.stripMargin + } else { + "" + } s""" |$declareRowBuffer |$findOrInsertHashMap |$incCounter |$updateRowInHashMap + |$emitPassThroughRow """.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 00d18a2f79a81..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 @@ -95,7 +95,11 @@ class TungstenAggregationIterator( peakMemory: SQLMetric, spillSize: SQLMetric, avgHashProbe: SQLMetric, - numTasksFallBacked: SQLMetric) + numTasksFallBacked: SQLMetric, + numBypassingRows: SQLMetric, + adaptivePartialAggEnabled: Boolean, + adaptiveMinRows: Long, + adaptiveMinCompaction: Double) extends AggregationIterator( partIndex, groupingExpressions, @@ -179,6 +183,20 @@ 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 (`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. @@ -191,7 +209,19 @@ class TungstenAggregationIterator( } } else { var i = 0 - while (inputIter.hasNext) { + var processedRows = 0L + val minRows = adaptiveMinRows + // The processed-row count at which the compaction ratio is evaluated next. It advances by + // `minRows` after every check, and restarts after a spill so the new in-memory map epoch is + // judged on its own rows. `minRows = 0` disables the periodic check: the count is only ever + // compared after being incremented past 0, so it never matches and only the spill check + // below remains. + var nextCheckRow = minRows + // The partial aggregation is ineffective when it does not collapse `minCompaction` rows into + // one key. There is no fast map on this path, so the map's keys are all the operator holds. + def ineffective(): Boolean = + processedRows < hashMap.getNumKeys().toDouble * adaptiveMinCompaction + while (inputIter.hasNext && !passThrough) { val newInput = inputIter.next() val groupingKey = groupingProjection.apply(newInput) var buffer: UnsafeRow = null @@ -199,21 +229,47 @@ 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. 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`. + pendingPassThroughRow = newInput.copy() } else { - externalSorter.merge(sorter) + val sorter = hashMap.destructAndCreateExternalSorter() + if (externalSorter == null) { + externalSorter = sorter + } else { + 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 + 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 + // 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 { + nextCheckRow += minRows + } } } - processRow(buffer, newInput) - i += 1 } if (externalSorter != null) { @@ -354,6 +410,47 @@ class TungstenAggregationIterator( } } + /////////////////////////////////////////////////////////////////////////// + // Part 5b: Methods and fields used by adaptive partial aggregation pass-through. + /////////////////////////////////////////////////////////////////////////// + + // Indicates that partial aggregation has been bypassed and the remaining input rows should be + // passed through as single-row partial buffers. Set in `processInputs` by either check point. + // It may coexist with earlier spills, so the output order is: sort-based (or map) output first, + // then the pass-through rows. + private[this] var passThrough: Boolean = false + + // 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 + + // 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) + numBypassingRows += 1 + generateOutput(groupingKey, passThroughAggregationBuffer) + } + /////////////////////////////////////////////////////////////////////////// // Part 6: Loads input rows and setup aggregationBufferMapIterator if we // have not switched to sort-based aggregation. @@ -398,12 +495,13 @@ class TungstenAggregationIterator( /////////////////////////////////////////////////////////////////////////// override final def hasNext: Boolean = { - (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) + (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) || + passThroughHasNext } 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. @@ -412,7 +510,7 @@ class TungstenAggregationIterator( sortBasedAggregationBuffer.copyFrom(initialAggregationBuffer) outputRow - } else { + } else if (mapIteratorHasNext) { // We did not fall back to sort-based aggregation. val result = generateOutput( @@ -426,13 +524,17 @@ class TungstenAggregationIterator( if (!mapIteratorHasNext) { // If there is no input from aggregationBufferMapIterator, we copy current result. val resultCopy = result.copy() - // Then, we free the map. + // Then, we free the map. Pass-through (if any) does not use the map. hashMap.free() resultCopy } else { result } + } else { + // Adaptive partial aggregation bypassed partial aggregation: emit the remaining input + // rows as single-row partial buffers. + nextPassThroughOutput() } numOutputRows += 1 diff --git a/sql/core/src/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 new file mode 100644 index 0000000000000..68c5979b0e98a --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala @@ -0,0 +1,927 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.aggregate + +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.aggregate.Partial +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests for runtime adaptive partial aggregation + * (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). When a partial aggregate is not reducing + * rows, the operator stops aggregating and streams the remaining rows through as single-row partial + * buffers for the Final aggregate to merge. It must never change results. + * + * The suite has two halves: + * 1. Correctness: output is identical to the reference (feature-off) run across the full matrix + * of codegen on/off, two-level map on/off, and spill/no-spill, over a range of aggregate + * shapes, key types, and `Expand`-bearing plans (ROLLUP / CUBE / GROUPING SETS / + * multi-distinct). + * 2. Triggering: the `numBypassingRows` metric proves the bypass actually fires when (and only + * when) it should -- high-cardinality input bypasses, low-cardinality input keeps aggregating, + * the feature switch and eligibility rules are honored, and both check points work. + */ +class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession + with AdaptiveSparkPlanHelper { + + import testImplicits._ + + // A `testFallbackStartsAt` setting ("fastMapCounter, regularMapCounter") that makes the regular + // map fall back (spill) periodically, exercising the spill-check decision path in both the + // codegen and interpreted aggregation paths. Kept moderate so low-cardinality inputs (which are + // never bypassed and therefore really spill) do not open an unbounded number of spill readers. + private val forceSpillFallback = "4, 16" + + // The upstream `CombineAdjacentAggregation` and `ReplaceHashWithSortAgg` rules would change the + // plan of these small single-partition queries away from a Partial+Final `HashAggregateExec`: + // the former merges the two adjacent phases (no shuffle in between) into a single `Complete` + // aggregate, and the latter converts a hash aggregate to a sort aggregate when the input is + // already sorted by the grouping key (a `Range` over an ascending `id` key). The adaptive + // feature lives in the partial hash aggregation, so both rules are disabled to keep that + // structure in the tests. + private val fixedPlanConfs = Seq( + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false", + SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "false") + + /** + * Runs `df` with adaptive partial aggregation disabled (the reference) and then across the full + * configuration matrix with it enabled, asserting every enabled run matches the reference. + */ + private def checkAdaptiveMatchesReference(build: () => DataFrame): Unit = { + val reference = withSQLConf( + (SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") +: fixedPlanConfs: _*) { + build().collect().toSeq + } + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + forceSpill <- Seq(true, false) + } { + val spillConf = if (forceSpill) { + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> forceSpillFallback) + } else { + Nil + } + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + // Small 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) { + checkAnswer(build(), reference) + } + } + } + } + + /** + * The observable per-run counters we assert on, all read from the partial `HashAggregateExec` in + * a single execution so the metrics are not double-counted: + * - `skipped`: our self-reported `numBypassingRows` metric. + * - `partialOutputRows`: the partial aggregate's own `numOutputRows`. An independent, + * pre-existing counter driven by the normal output path, so it is the ground truth for + * whether rows were streamed through -- it equals the distinct key count when aggregation is + * effective and climbs toward the input row count once the operator bypasses. + * - `spillBytes`: the partial aggregate's `spillSize`. Reliable only when no fallback is + * forced: on the interpreted path this is derived from the task-cumulative memory-spill + * counter, so a forced fallback (or downstream shuffle-write spill) can inflate it. + * Asserted only by the periodic check test, which forces no fallback; use + * `tasksFallBacked` otherwise. + * - `tasksFallBacked`: the partial aggregate's `numTasksFallBacked`, incremented only when the + * regular map actually falls back into sort-based aggregation. When the spill check bypasses + * at the spill boundary the sorter is never created, so this stays 0 -- direct, per-operator + * evidence the bypass replaced the sort fallback. + */ + private case class AggCounters( + skipped: Long, + partialOutputRows: Long, + spillBytes: Long, + tasksFallBacked: Long) + + // Verifies `df` (an already-collected bypassing run) produces the same results as the feature-off + // reference. `build` is re-run for the reference so it gets a genuinely non-adaptive plan rather + // than reusing the bypassing run's cached one. + private def checkAgainstReference(df: DataFrame, build: () => DataFrame): Unit = { + val reference = withSQLConf( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + build().collect().toSeq + } + checkAnswer(df, reference) + } + + private def runAndReadCounters(build: () => DataFrame): AggCounters = { + // The triggering tests assert on metrics, so also verify the bypassing run produces the same + // results as the feature-off reference. + val df = build() + df.collect() + val partialAggs = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => agg + } + // A partial aggregate is always present for the grouped queries these tests use. + assert(partialAggs.nonEmpty, "expected a partial HashAggregateExec in the plan") + val counters = AggCounters( + 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 spill check); when 0 the regular map does not spill. + */ + private def forEachCodegenAndMap( + minRows: Long = 8, + regularFallback: Int = 0, + minCompaction: Double = -1.0)( + body: String => Unit): Unit = { + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + // Shrink the fast map to 4 keys when it is on so rows reach the regular map. The second field + // controls regular-map spilling; 0 means "never" (a large sentinel). + val fallbackConf = if (twoLevelMap || regularFallback > 0) { + val fastCap = if (twoLevelMap) 4 else 1 + val regular = if (regularFallback > 0) regularFallback else Int.MaxValue + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> s"$fastCap, $regular") + } else { + Nil + } + // A negative value means "leave the threshold at its default". + val thresholdConf = if (minCompaction >= 0.0) { + Seq(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key -> minCompaction.toString) + } else { + Nil + } + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> minRows.toString) ++ + fallbackConf ++ thresholdConf ++ fixedPlanConfs): _*) { + body(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") + } + } + } + + ///////////////////////////////////////////////////////////////////////////// + // Part 1: Correctness -- results identical to the feature-off reference. + ///////////////////////////////////////////////////////////////////////////// + + test("results unchanged for high-cardinality input that bypasses partial aggregation") { + // Every grouping key is distinct, so partial aggregation reduces nothing and should be + // bypassed by the periodic check. + checkAdaptiveMatchesReference { () => + spark.range(0, 200, 1, 1) + .select($"id" as "k", ($"id" * 2) as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", max($"v") as "m") + } + } + + test("results unchanged for low-cardinality input that keeps partial aggregation") { + // Few distinct keys, high reduction: partial aggregation is effective and should be kept. + checkAdaptiveMatchesReference { () => + spark.range(0, 600, 1, 1) + .select(($"id" % 5) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", min($"v") as "mn", max($"v") as "mx") + } + } + + test("results unchanged for medium-cardinality input near the reduction threshold") { + // Roughly half the rows are distinct keys; exercises the boundary of the ratio checks. + checkAdaptiveMatchesReference { () => + spark.range(0, 1000, 1, 1) + .select(($"id" % 500) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with multiple grouping keys and string keys") { + checkAdaptiveMatchesReference { () => + spark.range(0, 500, 1, 1) + .select( + concat(lit("g"), ($"id" % 300).cast("string")) as "k1", + ($"id" % 7) as "k2", + $"id" as "v") + .groupBy($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with nullable grouping keys") { + 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 + // periodic check fires and the regression would show up. + checkAdaptiveMatchesReference { () => + spark.range(0, 1000, 1, 1) + .select(($"id" % 10) as "c") + .distinct() + } + } + + test("results unchanged when a large frozen map is output before pass-through streaming") { + // A larger 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) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "200000", + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> "false") ++ fixedPlanConfs): _*) { + val reference = withSQLConf( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + query().collect().toSeq + } + checkAnswer(query(), reference) + } + } + + test("distinct aggregation stays correct") { + checkAdaptiveMatchesReference { () => + spark.range(0, 300, 1, 1) + .select($"id" as "k", ($"id" % 50) as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd", sum($"v") as "s") + } + } + + test("distinct aggregation bypasses on high-cardinality input") { + // The `PartialMerge` phase of the multi-phase distinct plan always aggregates (it is not + // `Partial` mode and requires a distribution), so the rows reaching the distinct `Partial` + // phase are de-duplicated and pass-through carries exactly one distinct value each. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 1000, 1, 1) + .select(($"id" % 100) as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + assert(numBypassingRows(df) > 0, + "expected a distinct partial aggregation to bypass for high-cardinality input") + } + } + } + + test("count distinct: the de-duplication partial aggregate bypasses") { + // `count(DISTINCT v) GROUP BY k` plans two `Partial` phases: the de-duplication partial groups + // on (k, v) and the distinct partial groups on (k). Fully distinct (k, v) pairs make the + // de-duplication partial (2 grouping keys) reduce nothing, so it must bypass. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select(($"id" % 4) as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the (k, v) de-duplication partial to bypass, got $byKeyCount") + } + } + } + + test("count distinct: the distinct partial aggregate bypasses") { + // Mirror of the test above for the other phase: with many distinct keys but few distinct + // values per key, the (k, v) de-duplication partial reduces well while the distinct partial + // (1 grouping key) sees a fresh key per row and must bypass. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select($"id" as "k", ($"id" % 2) as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (grouping on k) to bypass, got $byKeyCount") + } + } + } + + test("count distinct: both partial aggregates bypass and results stay correct") { + // Fully distinct keys and fully distinct values: neither partial phase reduces anything, so + // both bypass in the same execution. The de-duplication partial keeps the (k, v) pairs unique + // and the distinct partial counts them, so the result must still match the reference. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the (k, v) de-duplication partial to bypass, got $byKeyCount") + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (grouping on k) to bypass, got $byKeyCount") + } + } + } + + test("global aggregation (no grouping keys) is never bypassed and stays correct") { + withSQLConf(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true") { + checkAnswer( + spark.range(0, 100, 1, 1).agg(sum($"id") as "s", count(lit(1)) as "c"), + Row(4950L, 100L)) + } + } + + test("results unchanged with an empty input") { + checkAdaptiveMatchesReference { () => + spark.range(0, 0, 1, 1) + .select($"id" as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + // The following four tests cover plans where an `ExpandExec` sits below the partial aggregate + // (ROLLUP / CUBE / GROUPING SETS / multi-distinct). PR apache/spark#28804 statically disabled its + // skip-partial-aggregate optimization whenever an Expand was present, but that was a performance + // heuristic guarding its *static* row sampling, not a correctness requirement. Our decision is + // made at runtime from the observed 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("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 + // 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, + "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("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 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") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(df) + assert(c.skipped > 0, + "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, + "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 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. + 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("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 `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"minCompaction=$minCompaction wholeStage=$wholeStage") { + numBypassingRows(df) + } + } + } + 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, + "codegen and interpreted paths must reach the same decision at the boundary") + } + } + } + + 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, + "an unreachable compaction ratio must bypass even a low-cardinality input") + } + } + } + + 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. + forEachCodegenAndMap(minRows = 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..fa8b9510cfe94 --- /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 periodic check bypasses, which should win. + * - low-cardinality, no spill: nothing bypasses, which must not regress. + * - 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 spill check 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 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) + addCodegenAdaptiveCases(benchmark, () => distinctKeyedDf(N)) + benchmark.run() + } + + // 1000 distinct keys over a large input: partial aggregation reduces a lot, the periodic check + // 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, + 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 periodic check (huge minRows). With + // fully distinct keys the reduction ratio is 1.0, so at the spill boundary the on-spill + // spill check bypasses instead of spilling; the baseline spills repeatedly and falls back + // to sort-based aggregation. + 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 spillCheckConf = Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> Long.MaxValue.toString, + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 1048576") + 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 spill check correctly does not + // bypass: both runs spill and fall back to sort-based aggregation identically (no regression). + 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 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 = spillCheckConf) + benchmark.run() + } + } + + private def distinctKeyedDf(N: Long): DataFrame = + spark.range(N).selectExpr("id as k", "id as v").groupBy("k").agg("v" -> "sum") +}