jayzhan211 commented on code in PR #24497:
URL: https://github.com/apache/datafusion/pull/24497#discussion_r4499474864
##########
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:
This test is still needed. I restricted the new `Full` selection to
`AggregateMode::Single` locally and all 2959 `datafusion-physical-plan` lib
tests still passed, so nothing covers `Partial` or `SinglePartitioned` over
grouped input. `SinglePartitioned` is the plan a key-partitioned source gets,
which is the shape #24438 is after. Could you add this test in this PR? It
passes on this head.
<details>
<summary>Test covering Partial → Final and SinglePartitioned over grouped
input</summary>
```rs
#[tokio::test]
async fn grouped_input_uses_clustered_partial_and_single_partitioned() ->
Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("key", DataType::Int32, false),
Field::new("value", DataType::Int64, false),
]));
let batch = |keys: Vec<i32>, values: Vec<i64>| {
RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int32Array::from(keys)),
Arc::new(Int64Array::from(values)),
],
)
};
// Groups are contiguous within each partition but not sorted. Key 1
spans
// a batch boundary, and keys 2 and 3 occur in both partitions.
let partitions = vec![
vec![
batch(vec![2, 1], vec![1, 2])?,
batch(vec![1, 3], vec![3, 4])?,
],
vec![batch(vec![3, 3, 2], vec![5, 6, 7])?],
];
let key = col("key", &schema)?;
let input: Arc<dyn ExecutionPlan> = Arc::new(
TestMemoryExec::try_new(&partitions, Arc::clone(&schema), None)?
.try_with_grouping_information(vec![vec![Arc::clone(&key)]])?,
);
let group_by = PhysicalGroupBy::new_single(vec![(key,
"key".to_string())]);
let aggr_expr = Arc::new(
AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?])
.schema(Arc::clone(&schema))
.alias("SUM(value)")
.build()?,
);
let task_ctx = new_migrated_hash_ctx(1024);
let single_partitioned = AggregateExec::try_new(
AggregateMode::SinglePartitioned,
group_by.clone(),
vec![Arc::clone(&aggr_expr)],
vec![None],
Arc::clone(&input),
Arc::clone(&schema),
)?;
assert_eq!(
single_partitioned.group_clustering_mode(),
&GroupClusteringMode::Full
);
assert!(matches!(
single_partitioned.execute_typed(0, &task_ctx)?,
StreamType::ClusteredSingleAggregate(_)
));
let partial = Arc::new(AggregateExec::try_new(
AggregateMode::Partial,
group_by.clone(),
vec![Arc::clone(&aggr_expr)],
vec![None],
input,
Arc::clone(&schema),
)?);
assert_eq!(partial.group_clustering_mode(), &GroupClusteringMode::Full);
assert!(matches!(
partial.execute_typed(0, &task_ctx)?,
StreamType::ClusteredPartialAggregate(_)
));
let final_aggregate = AggregateExec::try_new(
AggregateMode::Final,
group_by.as_final(),
vec![aggr_expr],
vec![None],
Arc::new(CoalescePartitionsExec::new(partial)),
schema,
)?;
let output = collect(final_aggregate.execute(0, task_ctx)?).await?;
assert_snapshot!(batches_to_sort_string(&output), @r"
+-----+------------+
| key | SUM(value) |
+-----+------------+
| 1 | 5 |
| 2 | 8 |
| 3 | 15 |
+-----+------------+
");
Ok(())
}
```
</details>
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1025,13 +1024,18 @@ impl AggregateExec {
})
.collect();
- let num_non_constant_groupby_exprs = groupby_exprs
+ let has_non_constant_groupby_exprs = groupby_exprs
.iter()
- .filter(|expr|
input_eq_properties.is_expr_constant(expr).is_none())
- .count();
- let mut group_clustering_mode = if indices.len() ==
num_non_constant_groupby_exprs
- && !indices.is_empty()
- && group_by.groups.len() == 1
+ .any(|expr| input_eq_properties.is_expr_constant(expr).is_none());
+ // 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
Review Comment:
Optional cleanup: this guard still lets through a single mask that nulls a
key when `has_grouping_set` is false. For example, `vec![vec![true, false]]`
over grouped `(a, b)` rows `(1, 1), (2, 2), (3, 1)` in separate batches gets
`Full` and returns `(NULL, 1)` twice, where hash mode merges them. Only a
hand-built `PhysicalGroupBy` can reach this, since the physical planner and
proto both carry the flag. Comparing against the all-false mask closes it, and
the `aggregates` tests still pass. To cover it, the mask above should give
`GroupClusteringMode::None`.
```suggestion
|| group_by.groups != [vec![false; group_by.expr.len()]]
```
--
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]
