You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ul...@apache.org on 2026/08/21 01:28:55 UTC
(spark) branch master updated: [SPARK-58860][SQL] Support PartialMerge mode in adaptive partial aggregation
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 25b0f249f646 [SPARK-58860][SQL] Support PartialMerge mode in adaptive partial aggregation
25b0f249f646 is described below
commit 25b0f249f6461296ede31015ed1a636c3477d47f
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Fri Aug 21 09:28:39 2026 +0800
[SPARK-58860][SQL] Support PartialMerge mode in adaptive partial aggregation
### What changes were proposed in this pull request?
Extend runtime adaptive partial aggregation (`spark.sql.execution.aggregate.adaptivePartialAggregation.enabled`) to support the DISTINCT intermediate phase, whose aggregate modes are `PartialMerge ++ Partial` (non-distinct aggregates in `PartialMerge`, distinct aggregates in `Partial`). Previously the feature only bypassed partial aggregates whose functions were all in `Partial` mode, so this phase was excluded.
### Why are the changes needed?
Bypassing a `PartialMerge` member is safe because the pass-through merges the incoming buffer into an empty buffer, which leaves it unchanged (the existing first-insertion of a key already relies on `merge(initial, x) == x`), and the downstream `Final` re-merges it. This lets high-cardinality DISTINCT queries benefit from the bypass. A pure `PartialMerge` de-duplication phase must still not bypass (it would over-count DISTINCT); it is kept out by its required distribution, with the `e [...]
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added tests to `AdaptivePartialAggregationSuite` covering DISTINCT with plain and filtered non-distinct aggregates, an order-sensitive non-distinct aggregate across partitions, and an imperative aggregate in the DISTINCT intermediate phase.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code (Anthropic).
Closes #58106 from ulysses-you/worktree-bypass.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../execution/aggregate/HashAggregateExec.scala | 30 +++--
.../AdaptivePartialAggregationSuite.scala | 134 ++++++++++++++++++---
2 files changed, 140 insertions(+), 24 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala
index 6c927dd859c3..59b40ecdfbd0 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
@@ -162,12 +162,20 @@ case class HashAggregateExec(
* 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. `Final`/`Complete` produce the result themselves and have no such
- * downstream. `PartialMerge` does have one and could be
- * supported by passing its incoming buffer through unchanged, but that is left for later.
- * The mode check is also what excludes the intermediate phase of a DISTINCT plan, whose
- * modes are `PartialMerge ++ Partial`.
+ * - `Partial` and `PartialMerge` modes only: the downstream `Final` aggregation merges the
+ * passed-through single-row buffers. `Final`/`Complete` produce the result themselves and
+ * have no such downstream. A `PartialMerge` member is the non-distinct aggregate of the
+ * DISTINCT intermediate phase (`AggUtils.planAggregateWithOneDistinct`): its input row is
+ * already a partial buffer, so the pass-through applies the merge to an empty buffer, which
+ * leaves the incoming buffer unchanged, and the downstream `Final` re-merges it. A pure
+ * `PartialMerge` phase (the de-duplication on keys ++ distinct columns) must not bypass, or
+ * duplicate (key, distinct column) rows would over-count DISTINCT. The built-in planner
+ * never emits such a phase without a required distribution, so the
+ * `requiredChildDistributionExpressions` check below already keeps it out. The
+ * `exists(_.mode == Partial)` check is a defensive guard on top of it, and only for a
+ * de-duplication phase that carries non-distinct aggregates: one with none at all has an
+ * empty `aggregateExpressions`, is admitted by the `isEmpty` disjunct, and still relies on
+ * the distribution check alone.
* - grouping keys present: a global aggregation produces a single output row, so partial
* aggregation achieves the maximum reduction and must never be bypassed.
* - no required distribution: with no aggregate functions (a group-by-only aggregate) the
@@ -176,9 +184,10 @@ case class HashAggregateExec(
* not. It is not a general pre-shuffle test, because `AggUtils.planAggregateWithOneDistinct`
* leaves it `None` on a post-shuffle aggregate. DISTINCT aggregate functions are allowed:
* the phase that de-duplicates on (keys ++ distinct columns) requires a distribution, so
- * this check keeps it out, while the distinct `Partial` phase that groups on the keys
- * alone is eligible even though it sits after a shuffle: another `Exchange` and a `Final`
- * follow it, so its passed-through buffers are still merged.
+ * this check keeps it out, while the distinct partial phase that groups on the keys alone
+ * (carrying the non-distinct aggregates as `PartialMerge`) is eligible even though it sits
+ * after a shuffle: another `Exchange` and a `Final` follow it, so its passed-through
+ * buffers are still merged.
* - batch only: a streaming partial aggregate keeps state across batches, and it is built
* with all-`Partial` modes and no required distribution, so it would otherwise qualify.
* A batch `session_window` grouping likewise qualifies, but its partial aggregate feeds a
@@ -191,7 +200,8 @@ case class HashAggregateExec(
groupingExpressions.nonEmpty &&
!isStreaming &&
!groupingExpressions.exists(_.metadata.contains(SessionWindow.marker)) &&
- aggregateExpressions.forall(a => a.mode == Partial) &&
+ aggregateExpressions.forall(a => a.mode == Partial || a.mode == PartialMerge) &&
+ (aggregateExpressions.exists(_.mode == Partial) || aggregateExpressions.isEmpty) &&
requiredChildDistributionExpressions.isEmpty
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala
index baafe42e4058..5eb9b4a8af6f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala
@@ -18,7 +18,7 @@
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.catalyst.expressions.aggregate.{Partial, PartialMerge}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -67,6 +67,16 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession
SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false",
SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "false")
+ // Whether `agg` is a partial or partial-merge phase. This matches every `HashAggregateExec`
+ // whose modes are all `Partial`/`PartialMerge`, including the DISTINCT intermediate phase
+ // (`PartialMerge ++ Partial`, the non-distinct aggregates in `PartialMerge` and the distinct
+ // ones in `Partial`) that the feature now bypasses. It is a superset of the phases the feature
+ // actually applies to: a pure `PartialMerge` de-duplication phase (excluded by its required
+ // distribution) and an ineligible group-by-only `Final` (vacuously all-`Partial`) also match,
+ // but neither registers a `numBypassingRows` metric, so they add 0 to the bypass count below.
+ private def isPartialPhase(agg: HashAggregateExec): Boolean =
+ agg.aggregateExpressions.forall(a => a.mode == Partial || a.mode == PartialMerge)
+
/**
* Runs `build` with adaptive partial aggregation disabled (the reference) and then across the
* full configuration matrix with it enabled, asserting every enabled run matches the reference.
@@ -121,12 +131,12 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession
s"twoLevelMap=$twoLevelMap forceSpill=$forceSpill"
withClue(msg) {
// Collect once so the metrics are populated, then check whether the bypass fired for
- // this cell. The metric lives on the `Partial`-mode `HashAggregateExec`, so that is the
- // operator the assertion reads.
+ // this cell. The metric lives on the partial `HashAggregateExec` phases, so those are the
+ // operators the assertion reads.
val df = build(inputPartitions)
df.collect()
val skipped = collect(df.queryExecution.executedPlan) {
- case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) =>
+ case agg: HashAggregateExec if isPartialPhase(agg) =>
agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L)
}.sum
if (expectBypass) {
@@ -183,7 +193,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession
val df = build()
df.collect()
val partialAggs = collect(df.queryExecution.executedPlan) {
- case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => agg
+ case agg: HashAggregateExec if isPartialPhase(agg) => agg
}
// A partial aggregate is always present for the grouped queries these tests use.
assert(partialAggs.nonEmpty, "expected a partial HashAggregateExec in the plan")
@@ -200,19 +210,32 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession
private def numBypassingRows(build: () => DataFrame): Long = runAndReadCounters(build).skipped
- // Returns the bypassed-row count per Partial-mode `HashAggregateExec`, keyed by the number of
+ // Returns the bypassed-row count per partial `HashAggregateExec` phase, 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.
+ // group-by has two such phases -- the de-duplication partial (grouping on key + distinct
+ // columns) and the distinct partial (grouping on the keys only, whose non-distinct aggregates run
+ // in `PartialMerge`) -- 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) =>
+ case agg: HashAggregateExec if isPartialPhase(agg) =>
agg.groupingExpressions.length ->
agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L)
}.groupBy(_._1).map { case (n, pairs) => n -> pairs.map(_._2).sum }
+ // The pure `PartialMerge` de-duplication phase (grouping on key + distinct columns) is the one
+ // phase the `exists(_.mode == Partial)` guard exists to exclude, and only when it carries
+ // non-distinct aggregates: with none at all its `aggregateExpressions` is empty, so the guard's
+ // `isEmpty` disjunct admits it and only its required distribution keeps it out, and
+ // `dedupPhases` below (which requires a non-empty `aggregateExpressions`) does not collect it.
+ // It shares the 2-key bucket above with the leading `Partial` phase, so pin its ineligibility
+ // directly.
+ val dedupPhases = collect(df.queryExecution.executedPlan) {
+ case agg: HashAggregateExec if agg.aggregateExpressions.nonEmpty &&
+ agg.aggregateExpressions.forall(_.mode == PartialMerge) => agg
+ }
+ assert(dedupPhases.forall(!_.metrics.contains("numBypassingRows")),
+ "the pure-PartialMerge de-duplication phase must stay ineligible")
checkAgainstReference(df, build)
byKeyCount
}
@@ -457,7 +480,7 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession
val df = query()
df.collect()
val skipped = collect(df.queryExecution.executedPlan) {
- case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) =>
+ case agg: HashAggregateExec if isPartialPhase(agg) =>
agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L)
}.sum
assert(skipped > 0,
@@ -479,10 +502,93 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession
}
}
+ test("distinct with plain and filtered non-distinct aggregates") {
+ // One query carries all three shapes through the DISTINCT intermediate phase
+ // (`PartialMerge ++ Partial`): the distinct aggregate (`count(DISTINCT v)`), a plain
+ // non-distinct aggregate (`sum(v)`), and a filtered non-distinct aggregate
+ // (`avg(v) FILTER (...)`, whose `FILTER` is applied in the leading `Partial` phase only).
+ // Fully distinct keys and values make neither partial phase reduce anything, so both bypass in
+ // the same execution: asserting the 2-key phase (de-duplication, all `Partial`) and the 1-key
+ // phase (distinct partial, `PartialMerge ++ Partial`) together proves the two bypasses coexist,
+ // the plain and filtered non-distinct buffers pass through correctly, and the results still
+ // match the feature-off reference.
+ withTempView("t") {
+ spark.range(0, 400, 1, 1)
+ .select($"id".cast("string") as "k", $"id" as "v")
+ .createOrReplaceTempView("t")
+ forEachCodegenAndMap() { clue =>
+ val df = () => spark.sql(
+ """SELECT k,
+ | count(DISTINCT v) AS cd,
+ | sum(v) AS s,
+ | avg(v) FILTER (WHERE v > 25) AS a_gt25
+ |FROM t GROUP BY k""".stripMargin)
+ withClue(clue) {
+ val byKeyCount = bypassRowsByGroupingKeyCount(df)
+ assert(byKeyCount.get(2).exists(_ > 0),
+ s"expected the de-duplication partial (grouping on k, v) to bypass, got $byKeyCount")
+ assert(byKeyCount.get(1).exists(_ > 0),
+ s"expected the distinct partial (PartialMerge++Partial, grouping on k) to bypass, " +
+ s"got $byKeyCount")
+ }
+ }
+ }
+ }
+
+ test("distinct with an order-sensitive non-distinct aggregate across partitions") {
+ // A single-partition `range` fuses the whole four-phase DISTINCT plan into one stage, so the
+ // split-topology path (the frozen map draining one row per queued row, the queue flush, and
+ // the `shouldStop()` re-entry) never runs for a `PartialMerge` member, including an
+ // order-sensitive one. Deriving both grouping columns and varying the partition count forces
+ // the exchanges; carrying `first`/`last` makes the pass-through's merge-into-an-empty-buffer
+ // step observable. The forced-spill cells are omitted because the sort-based fallback reorders
+ // `first`/`last` in the feature-off reference arm too.
+ forEachCodegenAndMap() { clue =>
+ Seq(1, 2).foreach { parts =>
+ val df = () => spark.range(0, 400, 1, parts)
+ .select(($"id" % 100).cast("string") as "k", ($"id" % 7) as "v", $"id" as "w")
+ .groupBy($"k")
+ .agg(countDistinct($"v") as "cd", sum($"w") as "s",
+ first($"w") as "f", last($"w") as "l")
+ withClue(s"$clue parts=$parts") {
+ val byKeyCount = bypassRowsByGroupingKeyCount(df)
+ assert(byKeyCount.get(2).exists(_ > 0),
+ s"expected the de-duplication partial (grouping on k, v) to bypass, got $byKeyCount")
+ assert(byKeyCount.get(1).exists(_ > 0),
+ s"expected the distinct partial (PartialMerge++Partial, grouping on k) to bypass, " +
+ s"got $byKeyCount")
+ }
+ }
+ }
+ }
+
+ test("an imperative aggregate stays correct in the distinct intermediate phase") {
+ // `approx_count_distinct` uses `HyperLogLogPlusPlus`, an `ImperativeAggregate` whose buffer is
+ // written by `initialize`/`merge` rather than a projection, so the distinct intermediate phase
+ // runs on `TungstenAggregationIterator` (supportCodegen = false). Asserting both phases bypass
+ // -- the de-duplication partial (grouping on `k` + `v`) and the distinct partial (grouping on
+ // `k`, whose `PartialMerge` member is imperative) -- proves the imperative buffer is reset then
+ // merged with the incoming buffer on pass-through, and the results still match the reference.
+ forEachCodegenAndMap() { clue =>
+ val df = () => spark.range(0, 400, 1, 1)
+ .select(($"id" % 50).cast("string") as "k", ($"id" % 20) as "v", $"id" as "x")
+ .groupBy($"k")
+ .agg(approx_count_distinct($"x") as "acd", countDistinct($"v") as "cd")
+ withClue(clue) {
+ val byKeyCount = bypassRowsByGroupingKeyCount(df)
+ assert(byKeyCount.get(2).exists(_ > 0),
+ s"expected the de-duplication partial (grouping on k, v) to bypass, got $byKeyCount")
+ assert(byKeyCount.get(1).exists(_ > 0),
+ s"expected the distinct partial (imperative PartialMerge, grouping on k) to bypass, " +
+ s"got $byKeyCount")
+ }
+ }
+ }
+
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.
+ // The `PartialMerge` phase of the multi-phase distinct plan always aggregates (it requires a
+ // distribution, so it is never eligible), so the rows reaching the distinct `Partial` phase
+ // are de-duplicated and pass-through carries exactly one distinct value each.
forEachCodegenAndMap() { clue =>
val df = () => spark.range(0, 1000, 1, 1)
.select(($"id" % 100).cast("string") as "k", $"id" as "v")
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org