mbutrovich commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4391435329
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -247,6 +247,37 @@ 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. Two things
make the task succeed
+ // there: the new stream's spill path, and the fair pool exempting
every consumer named
+ // OrderedFinalAggregateStream from its fair limit
(`limited_to_its_share`), the limit the
+ // legacy stream's spill failed against. So this PartialMerge is
checked only against the
+ // pool and Spark's grant.
+ // 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. A PartialMerge over Spark's binary
collect buffers keeps
+ // Single, but that input arrives through a ScanExec with no ordering,
so it is not
+ // expected to be sorted on a key.
+ // DataFusion 56 runs Single over this input with a
ClusteredSingleAggregateStream, and
+ // renames this stream ClusteredFinalAggregateStream, which is also
its consumer name
+ // (apache/datafusion#24697). So on that upgrade, drop this switch and
update the
+ // consumer name prefix in `spill_replay.rs` that the fair pool
exemption uses (#6410).
+ let mode = if proto_mode == ProtoAggregateMode::PartialMerge
+ && input_is_keys_then_states(agg, &aggr_expr, &schema)?
+ && !child
+ .native_plan
+ .equivalence_properties()
+ .find_longest_permutation(&group_by.input_exprs())?
+ .1
+ .is_empty()
+ {
+ DFAggregateMode::Final
+ } else {
+ mode
+ };
Review Comment:
Should this switch apply only when the input is sorted on some but not all
of the grouping keys?
As I read it, the guard also matches a native sort-based PartialMerge.
Spark's `SortAggregateExec` requires its input sorted on all of its grouping
keys
([`SortAggregateExec.scala#L52-L54`](https://github.com/apache/spark/blob/1a07005f0c2cf151c9d85963a99910c4ddb8e896/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala#L52-L54)).
So a `CometSortAggregateExec` PartialMerge over a native sort sees `Sorted`
input, and its states arrive in native types through a Comet shuffle. The
distinct query in
[`min_max.sql#L65-L68`](https://github.com/apache/datafusion-comet/blob/5bf7b3617876f927472bd950726ad5f9617fb202/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql#L65-L68)
and [`SortAggregate runs a distinct collect_list chain
natively`](https://github.com/apache/datafusion-comet/blob/5bf7b3617876f927472bd950726ad5f9617fb202/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala#L4533)
look like this shape to me. The new p
lanner test also expects `Final` for `Sorted` input
([`planner.rs#L4856-L4858`](https://github.com/apache/datafusion-comet/blob/5bf7b3617876f927472bd950726ad5f9617fb202/native/core/src/execution/planner.rs#L4856-L4858)).
On fully sorted input, `Final` doesn't add a spill path. The legacy stream
reports the error for `GroupOrdering::Full`
([`grouped_hash_stream.rs#L492-L510`](https://github.com/apache/datafusion/blob/500744d5c25499055e3b2ca46b022bc425989203/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs#L492-L510)),
and `OrderedFinalAggregateStream` spills only for `PartiallySorted` input
([`ordered_final_stream.rs#L337-L338`](https://github.com/apache/datafusion/blob/500744d5c25499055e3b2ca46b022bc425989203/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs#L337-L338)).
What changes is how the fair pool counts the aggregate. The legacy stream
registers as unable to spill in that case
([`grouped_hash_stream.rs#L518`](https://github.com/apache/datafusion/blob/500744d5c25499055e3b2ca46b022bc425989203/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs#L518)).
`OrderedFinalAggregateStream` always registers as able to
([`ordered_final_stream.rs#L288-L294`](https://github.com/apache/datafusion/blob/500744d5c25499055e3b2ca46b022bc425989203/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs#L288-L294)).
So the aggregate joins `num_spill`, which divides the task's share
([`fair_pool.rs#L249`](https://github.com/apache/datafusion-comet/blob/5bf7b3617876f927472bd950726ad5f9617fb202/native/core/src/execution/memory_pools/fair_pool.rs#L249)).
When the sort below is the only other consumer that can spill, its limit drops
from the task's share, less what the aggregate holds, to half of the share.
[`memory_management.md#L397-L398`](https:
//github.com/apache/datafusion-comet/blob/5bf7b3617876f927472bd950726ad5f9617fb202/docs/source/contributor-guide/memory_management.md?plain=1#L397-L398)
describes that halving for the `Final` sort aggregate.
Is that intended for these PartialMerges? If not, could the guard also
require fewer matched indices from `find_longest_permutation` than there are
grouping keys, which is how DataFusion tells `PartiallySorted` from `Sorted`
([`mod.rs#L1011-L1020`](https://github.com/apache/datafusion/blob/500744d5c25499055e3b2ca46b022bc425989203/datafusion/physical-plan/src/aggregates/mod.rs#L1011-L1020))?
The `Sorted` planner case would then expect `Single`. If it is intended, could
the PR description say so, and could
[`fair_pool.rs#L166-L169`](https://github.com/apache/datafusion-comet/blob/5bf7b3617876f927472bd950726ad5f9617fb202/native/core/src/execution/memory_pools/fair_pool.rs#L166-L169)
and
[`memory_management.md#L399-L402`](https://github.com/apache/datafusion-comet/blob/5bf7b3617876f927472bd950726ad5f9617fb202/docs/source/contributor-guide/memory_management.md?plain=1#L399-L402)
cover this case? Both say the exempted PartialMerge's input is only partially
sorted.
--
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]
