This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 77f946a815b0 [SPARK-58860][SQL] Support PartialMerge mode in adaptive
partial aggregation
77f946a815b0 is described below
commit 77f946a815b01a63b9e5dedda0651ef34531072d
Author: Xiduo You <[email protected]>
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 <[email protected]>
Signed-off-by: Xiduo You <[email protected]>
(cherry picked from commit 25b0f249f6461296ede31015ed1a636c3477d47f)
Signed-off-by: Xiduo You <[email protected]>
---
.../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: [email protected]
For additional commands, e-mail: [email protected]
