jayzhan211 commented on code in PR #24497:
URL: https://github.com/apache/datafusion/pull/24497#discussion_r4480661590
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -5961,71 +6040,104 @@ mod tests {
}
#[tokio::test]
- async fn unsorted_contiguous_groups_use_final_emission() -> Result<()> {
+ async fn unsorted_contiguous_groups_use_incremental_emission() ->
Result<()> {
Review Comment:
Could you also cover the partial stage? Multi-partition plans run
`ClusteredPartialAggregateStream` followed by a final aggregate, but only
`Single` is exercised here. Something like this, using the same grouped input:
```rs
let partial = Arc::new(AggregateExec::try_new(
AggregateMode::Partial,
group_by.clone(),
vec![Arc::clone(&aggr_expr)],
vec![None],
Arc::clone(&input) as Arc<dyn ExecutionPlan>,
Arc::clone(&schema),
)?);
assert_eq!(partial.group_clustering_mode(), &GroupClusteringMode::Full);
assert!(matches!(
partial.execute_typed(0, &task_ctx)?,
StreamType::ClusteredPartialAggregate(_)
));
let final_agg = AggregateExec::try_new(
AggregateMode::Final,
group_by.as_final(),
vec![aggr_expr],
vec![None],
Arc::new(CoalescePartitionsExec::new(partial)),
schema,
)?;
let output = collect(final_agg.execute(0, task_ctx)?).await?;
// same snapshot as the single-stage result
```
The input needs at least two partitions so the final stage actually merges,
and a `TestMemoryExec` input with grouping information instead of `BarrierExec`
keeps it simple.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1044,6 +1044,10 @@ impl AggregateExec {
// states without using group boundaries to recognize completed groups.
if group_by.has_grouping_set() || mode == AggregateMode::PartialReduce
{
group_clustering_mode = GroupClusteringMode::None;
+ } else if num_non_constant_groupby_exprs > 0
Review Comment:
Optional cleanup: `grouping_satisfy` already covers the ordering-based
`Full` case, since it drops constants and then checks
`find_longest_permutation`. The mode can therefore be picked in one place. This
also puts back the `groups.len() == 1` check that the old `Full` branch had.
Without it, a hand-built `PhysicalGroupBy` with two groups and
`has_grouping_set = false` over grouped input currently gets `Full` and emits
duplicate groups. I ran this locally and all `aggregates` tests pass:
```diff
- let mut group_clustering_mode = if indices.len() ==
num_non_constant_groupby_exprs
- && !indices.is_empty()
- && group_by.groups.len() == 1
- {
- GroupClusteringMode::Full
- } else if !indices.is_empty() {
- GroupClusteringMode::Partial(indices)
- } else {
- GroupClusteringMode::None
- };
-
- // Grouping sets can change group keys. PartialReduce combines
intermediate
- // states without using group boundaries to recognize completed
groups.
- if group_by.has_grouping_set() || mode ==
AggregateMode::PartialReduce {
- group_clustering_mode = GroupClusteringMode::None;
- } else if num_non_constant_groupby_exprs > 0
- &&
input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())?
- {
- group_clustering_mode = GroupClusteringMode::Full;
- }
+ // Grouping sets can change group keys. PartialReduce combines
intermediate
+ // states without using group boundaries to recognize completed
groups.
+ let group_clustering_mode = if group_by.has_grouping_set()
+ || group_by.groups.len() != 1
+ || mode == AggregateMode::PartialReduce
+ {
+ GroupClusteringMode::None
+ } else if num_non_constant_groupby_exprs > 0
+ &&
input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())?
+ {
+ GroupClusteringMode::Full
+ } else if !indices.is_empty() {
+ GroupClusteringMode::Partial(indices)
+ } else {
+ GroupClusteringMode::None
+ };
```
--
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]
