viirya commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4371853661
##########
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:
Is there a real plan where a PartialMerge over Spark's binary collect
buffers sees input sorted on only some of its keys? A Spark partial's output
reaches native code through a `ScanExec` that reports no ordering, so it would
run as `Linear` on `SingleHashAggregateStream`. A sort aggregate sorts on every
key, which is `Sorted`, and the legacy stream doesn't spill there anyway. If
there isn't one, could the comment say this case isn't expected in practice,
rather than that it can still fail this way? Otherwise a reader may go looking
for a failure that can't happen.
##########
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:
`input_sorted_on_a_key` is computed for every aggregate, before the
`proto_mode == ProtoAggregateMode::PartialMerge` check. So every Partial and
Final aggregate now runs `find_longest_permutation` and gains a new `?` error
path, only for the result to be thrown away. Could it be evaluated only inside
the `PartialMerge` condition, for example as the second operand of the `&&`?
##########
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:
This shape needs `EliminateSorts` excluded, because the optimizer otherwise
drops the `sortWithinPartitions` under the aggregate. I'd expect users to hit
it after a sort-merge join on `g`. The join's `HashPartitioning(g)` satisfies
the PartialMerge's clustered distribution on `(g, k)`, so the join, the partial
and the PartialMerge run in one native plan, and the PartialMerge sees its
input sorted on `g`. Could the test use that shape with the default rules, or
add it as a second case?
##########
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:
Planning this as `Final` also gives the consumer the
`OrderedFinalAggregateStream` name, which `limited_to_its_share` in
`fair_pool.rs` exempts by prefix whether the input is fully or partially
sorted. So this PartialMerge is no longer checked against its fair limit, the
limit the issue's error came from. Could the comment here, and the PR
description, say that this is part of why the task now succeeds? Could the
exemption paragraph in `memory_management.md` and the `limited_to_its_share`
doc mention this case too? Both describe the exemption as being for a
`CometSortAggregateExec` over fully sorted input.
##########
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:
With `0L AS g`, the input is sorted on a key that never changes. So the
ordered final stream never emits a finished group before it spills, and the
test never exercises the rule that every group waits for the replay once a run
has spilled. Could the data use a few `g` values, for example `id % 4 AS g`,
and check one row per `g`? A key emitted twice across a prefix boundary would
then show up in the answer, and that's the failure #6857 was about.
--
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]
