ulysses-you commented on code in PR #58106:
URL: https://github.com/apache/spark/pull/58106#discussion_r3817967992
##########
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:
Thanks for the detailed analysis and the suggested shape. Added `distinct
with an order-sensitive non-distinct aggregate across partitions`, which varies
the partition count (1, 2), derives both `k` and `v` so the exchanges actually
appear, and carries `first`/`last`. It uses `forEachCodegenAndMap()` without
`regularFallback` so the sort-based fallback cannot reorder `first`/`last` in
the reference arm.
##########
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:
Good point. Added the `dedupPhases` assertion to
`bypassRowsByGroupingKeyCount` to pin the pure-`PartialMerge` de-duplication
phase's ineligibility directly.
##########
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:
Done, dropped the mode half and kept only the distribution reason.
--
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]
