peter-toth commented on code in PR #58106:
URL: https://github.com/apache/spark/pull/58106#discussion_r3814887092
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala:
##########
@@ -198,16 +208,16 @@ 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) =>
Review Comment:
**Finding 2.** With the predicate widened, the pure-`PartialMerge`
de-duplication phase now matches `isPartialPhase` and lands in the same
grouping-key bucket as the leading `Partial` phase -- both group on `(k, v)`.
So the one phase the new `exists(_.mode == Partial)` guard exists to exclude
has no observable count of its own: if it ever became eligible, its bypasses
would be summed into `byKeyCount(2)` and the `> 0` assertions would still pass.
`checkAgainstReference` would catch the resulting over-counted DISTINCT, but
nothing pins the eligibility rule itself.
The helper comment above already states that such a phase registers no
metric, so asserting it is a couple of lines here:
```scala
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")
```
It holds on this head: that phase reports no `numBypassingRows` metric while
the `PartialMerge ++ Partial` phase reports 392 in the new test.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala:
##########
@@ -477,6 +487,62 @@ 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)
Review Comment:
**Finding 1.** A one-slice `RangeExec` reports `SinglePartition`, which
satisfies every `ClusteredDistribution`, so `EnsureRequirements` inserts no
`Exchange` anywhere and the whole four-phase DISTINCT plan fuses into one
stage. Dumped on this head:
```
*(1) HashAggregate(keys=[k], functions=[sum(v), avg(v), count(distinct v)])
+- *(1) HashAggregate(keys=[k], functions=[merge_sum(v), merge_avg(v),
partial_count(distinct v)])
+- *(1) HashAggregate(keys=[k, v], functions=[merge_sum(v), merge_avg(v)])
+- *(1) HashAggregate(keys=[k, v], functions=[partial_sum(v),
partial_avg(v) FILTER (WHERE (v > 25))])
+- *(1) Project [cast(id as string) AS k, id AS v]
+- *(1) Range (0, 400, step=1, splits=1)
```
Fused, this operator's `outputFunc` feeds the next aggregate's `doConsume`,
so nothing reaches `BufferedRowIterator.currentRows`, `shouldStop()` stays
false for the whole build, and the first bypassed row drains the entire frozen
map in a single `outputMapAndFlush()` call. The interesting half of the
machinery -- one map row per queued row, the queue flush after the map drains,
and the `adaptiveResumeBuild` re-entry -- only runs with an `Exchange` above
the operator, which is also the shape where the merge-order work on #57742
lived. Neither new test reaches it, and neither carries an order-sensitive
non-distinct aggregate, which is the only place the pass-through's
merge-into-an-empty-buffer step is observable.
`distinct aggregation stays correct` (`:481`) does reach the split shape for
this phase at `parts = 2` -- I confirmed both exchanges appear and the
`PartialMerge ++ Partial` phase bypasses 292 rows there -- but its
`expectBypass` assertion sums `numBypassingRows` over every matching phase, so
the leading `Partial` phase alone satisfies it.
Suggested shape: vary the partition count and add an order-sensitive member.
```scala
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")
...
}
}
```
Deriving both `k` and `v` matters: with `v = id` the `Project` keeps the
`Range`'s partitioning as `RangePartitioning(v)`, which satisfies
`ClusteredDistribution(k, v)`, and the exchange is skipped again even at `parts
= 2`.
I ran that query and it passes today (splits 1 and 2, codegen on and off,
two-level map on and off), so this is coverage rather than a live bug. One
caveat if you add `first`/`last`: keep the forced-`testFallbackStartsAt` cells
out of the reference comparison. The sort-based fallback reorders
`first`/`last` in the *reference* arm too -- feature off, that query returns
`(f, l) = (210, 110)` for `k = 10` with `testFallbackStartsAt = "4, 16"` and
`(10, 310)` without it -- so a forced-spill cell would compare two legitimately
different orders.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala:
##########
@@ -477,6 +487,62 @@ 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("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`
Review Comment:
**Finding 3.** "it is not `Partial` mode" stops being a reason with this PR
-- `PartialMerge` modes are eligible now. And for this test's query
(`countDistinct($"v")` with no non-distinct aggregate) that phase has no
aggregate expressions at all, so it is a group-by-only aggregate kept out
purely by its required distribution. Suggest dropping the mode half:
```scala
// 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.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
