andygrove commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4371966378
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -247,6 +247,30 @@ impl OperatorBuilder for HashAggBuilder {
})
.collect();
+ // DataFusion 55 runs Single over input sorted on some of its keys
with the legacy
+ // GroupedHashAggregateStream, whose spill can ask for more than its
share of the pool
+ // (#6880). Its Final runs an OrderedFinalAggregateStream there
instead, which spills, and
+ // MergeAsPartial evaluates to its state under Final too. Final reads
each state from the
+ // columns after the keys and spills its groups in the input's types,
so plan the
+ // PartialMerge as Final only when its input is its keys followed by
native states.
+ // One over Spark's binary collect buffers keeps Single and can still
fail this way.
+ // DataFusion 56 replaces the legacy stream with an
OrderedSingleAggregateStream, so
+ // recheck this on that upgrade (#6410).
+ let input_sorted_on_a_key = !child
Review Comment:
Good catch. The sort check now runs last in the `&&` chain, so only a
`PartialMerge` whose input is its keys followed by native states pays for it.
Partial and Final aggregates skip it.
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -247,6 +247,30 @@ impl OperatorBuilder for HashAggBuilder {
})
.collect();
+ // DataFusion 55 runs Single over input sorted on some of its keys
with the legacy
+ // GroupedHashAggregateStream, whose spill can ask for more than its
share of the pool
+ // (#6880). Its Final runs an OrderedFinalAggregateStream there
instead, which spills, and
+ // MergeAsPartial evaluates to its state under Final too. Final reads
each state from the
+ // columns after the keys and spills its groups in the input's types,
so plan the
+ // PartialMerge as Final only when its input is its keys followed by
native states.
+ // One over Spark's binary collect buffers keeps Single and can still
fail this way.
Review Comment:
I couldn't construct one. A Spark partial's output arrives through a
`ScanExec` that reports no ordering, so it runs as `Linear` on the hash stream.
The comment now says that input isn't expected to be sorted on a key, and the
description no longer says it can still fail.
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -247,6 +247,30 @@ impl OperatorBuilder for HashAggBuilder {
})
.collect();
+ // DataFusion 55 runs Single over input sorted on some of its keys
with the legacy
+ // GroupedHashAggregateStream, whose spill can ask for more than its
share of the pool
+ // (#6880). Its Final runs an OrderedFinalAggregateStream there
instead, which spills, and
+ // MergeAsPartial evaluates to its state under Final too. Final reads
each state from the
+ // columns after the keys and spills its groups in the input's types,
so plan the
+ // PartialMerge as Final only when its input is its keys followed by
native states.
+ // One over Spark's binary collect buffers keeps Single and can still
fail this way.
+ // DataFusion 56 replaces the legacy stream with an
OrderedSingleAggregateStream, so
+ // recheck this on that upgrade (#6410).
+ let input_sorted_on_a_key = !child
+ .native_plan
+ .equivalence_properties()
+ .find_longest_permutation(&group_by.input_exprs())?
+ .1
+ .is_empty();
+ let mode = if proto_mode == ProtoAggregateMode::PartialMerge
Review Comment:
Done. The comment here now says the task succeeds because of the new
stream's spill path and because `limited_to_its_share` exempts the stream from
the fair limit. The `limited_to_its_share` doc, the exemption paragraph in
`memory_management.md` and the PR description say it too.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -301,6 +301,47 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("PartialMerge over input sorted on a prefix of its keys spills (issue
#6880)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark
+ .range(0, 125000, 1, 1)
+ .selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS
k", "id AS v")
Review Comment:
Changed. The data now has four values of `g` (`id % 4`), and the test checks
one row per `g` with `max(v)` and `count_distinct(k)`. Both pool sizes still
spill and fail with the issue's error when I turn the planner change off.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -301,6 +301,47 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("PartialMerge over input sorted on a prefix of its keys spills (issue
#6880)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark
+ .range(0, 125000, 1, 1)
+ .selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS
k", "id AS v")
Review Comment:
Changed. The data now has four values of `g` (`id % 4`), and the test checks
one row per `g` with `max(v)` and `count_distinct(k)`. Both pool sizes still
spill and fail with the issue's error when I turn the planner change off.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -301,6 +301,47 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("PartialMerge over input sorted on a prefix of its keys spills (issue
#6880)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark
+ .range(0, 125000, 1, 1)
+ .selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS
k", "id AS v")
+ .write
+ .parquet(path)
+ // The repartition on g satisfies the PartialMerge keyed by (g, k), so
the sort, the partial
+ // aggregate and the PartialMerge run in one native plan with no
exchange between them, and
+ // the PartialMerge sees its input sorted on g. In 2 and 8 MiB pools,
the PartialMerge's
+ // spill asked for more than its share before it planned as an ordered
final aggregate.
+ for (poolMiB <- Seq(2, 8)) {
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName,
Review Comment:
Added as a second test: a sort-merge join on `g` with the default rules,
which plans the join, the partial and the PartialMerge in one native plan. It
needs AQE off, or it becomes a broadcast join, and a 16 MiB pool, because the
sorts under the join can't get what they need from a 2 or 8 MiB pool's share.
It fails with the issue's error without the planner change.
--
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]
