mbutrovich commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4383770943
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -247,6 +247,35 @@ 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 replaces the legacy stream with an
OrderedSingleAggregateStream, so
+ // recheck this on that upgrade (#6410).
Review Comment:
As I read DataFusion `main`, there is no `OrderedSingleAggregateStream`
there anymore. apache/datafusion#24697 renamed the ordered streams to
[`ClusteredSingleAggregateStream`](https://github.com/apache/datafusion/blob/c18e5d974c49382ba892fa9f89ce71d3aa0ab8ae/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs#L205)
and
[`ClusteredFinalAggregateStream`](https://github.com/apache/datafusion/blob/c18e5d974c49382ba892fa9f89ce71d3aa0ab8ae/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs#L121),
and those are also their memory consumer names. This fix depends on the
`OrderedFinalAggregateStream[` prefix in
[`spill_replay.rs`](https://github.com/apache/datafusion-comet/blob/9c4e4b72d9bb9fa35a8c8a40b74ed889a432a0ad/native/core/src/execution/memory_pools/spill_replay.rs#L185)
and the fair pool exemption that reuses it
([`fair_pool.rs`](https://github.com/apache/datafusion-comet/blob/9c4e4b72d9bb9fa35a8c8a40b74ed889a432a0ad/native/core/src/execution/memory_po
ols/fair_pool.rs#L175)). On 56, neither would match the stream this
PartialMerge runs on. #6404 pins a DataFusion revision from before the rename,
so its CI hasn't seen this yet.
Could the comment name the stream as it is on `main`? And since this PR
closes #6880, could #6410 get an item for the 56 upgrade that covers both
removing this `Final` switch and updating the consumer name prefix? Its current
#6880 entry only says to recheck the failure.
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -309,6 +338,37 @@ impl OperatorBuilder for HashAggBuilder {
}
}
+/// Whether a PartialMerge's input is its grouping keys, in order, followed by
the states of
+/// its MergeAsPartial aggregates in their native types, which DataFusion's
Final assumes. A
+/// Spark partial aggregate's collect buffers, for example, arrive as binary.
+fn input_is_keys_then_states(
+ agg: &datafusion_comet_proto::spark_operator::HashAggregate,
+ aggr_expr: &[Arc<AggregateFunctionExpr>],
+ input_schema: &SchemaRef,
+) -> Result<bool, ExecutionError> {
+ let num_keys = agg.grouping_exprs.len();
+ let keys_lead = agg.initial_input_buffer_offset as usize == num_keys
+ && agg.grouping_exprs.iter().enumerate().all(|(idx, expr)| {
+ matches!(&expr.expr_struct, Some(ExprStruct::Bound(bound)) if
bound.index as usize == idx)
+ });
+ if !keys_lead {
+ return Ok(false);
+ }
+ let mut input_types = input_schema
+ .fields()
+ .iter()
+ .skip(num_keys)
+ .map(|f| f.data_type());
+ for expr in aggr_expr {
+ for field in expr.state_fields()? {
+ if input_types.next() != Some(field.data_type()) {
+ return Ok(false);
+ }
+ }
+ }
+ Ok(input_types.next().is_none())
Review Comment:
What do you think about adding cases to
`partial_merge_over_sorted_input_runs_as_final` for the inputs this guard
rejects? I mutated the guard on the head commit and ran `cargo test -p
datafusion-comet --lib execution::`. Replacing `if !keys_lead` with a check
that never returns, or replacing the last line with `Ok(true)`, leaves every
test passing. Only the type comparison is checked today, by the `collect_list`
case. Each rejected input is one where `Final` would read the wrong columns or
spill with a schema that doesn't match its groups, so I think each needs a case
that expects `Single`. For example, the grouping keys in swapped order
(`key(1), key(0)`), `initial_input_buffer_offset` of 1, and one extra `int64`
input column after the states.
--
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]
