comphead commented on code in PR #26819:
URL: https://github.com/apache/datafusion/pull/26819#discussion_r4403637861
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -5570,6 +5582,114 @@ mod tests {
Ok(())
}
+ /// Partial aggregation output is not unique on the group keys across
+ /// partitions, so it must not declare `Unique(group keys)`. Otherwise an
+ /// ordering on the group keys alone satisfies `[k, sum]` after a merge,
and
+ /// `SortExec` skips sorting rows that share a key.
+ #[tokio::test]
+ async fn partial_aggregate_does_not_declare_unique_group_keys() ->
Result<()> {
Review Comment:
This test overlaps with
`test_keep_sort_above_repartitioned_partial_aggregate`. Both build the same two
sorted partitions and the same Partial `sum`, and both check that the sort on
`[k, sum]` survives.
Once the constraint is gone, `ordering_satisfy` can't use it, so asserting
the constraint directly is enough here. The existing
`group_clustering_preserves_sort_options` already runs `Partial` and `Single`
over sorted input, so one assertion there covers both modes:
```rust
let expected_constraints = match mode.output_mode() {
AggregateOutputMode::Final => {
Constraints::new_unverified(vec![Constraint::Unique(vec![0, 1])])
}
AggregateOutputMode::Partial => Constraints::default(),
};
assert_eq!(
aggregate.properties().equivalence_properties().constraints(),
&expected_constraints
);
```
Then this test, including the `SortPreservingMergeExec` and `SortExec`
execution and the two new imports, can go. The `enforce_sorting.rs` test still
covers the removed-sort path end to end.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1481,20 +1481,30 @@ impl AggregateExec {
// Group by expression will be a distinct value after the aggregation.
// Add it into the constraint set.
- let mut constraints = eq_properties.constraints().to_vec();
- let new_constraint = Constraint::Unique(
- group_expr_mapping
- .iter()
- .flat_map(|(_, target_cols)| {
- target_cols.iter().filter_map(|(expr, _)| {
- expr.downcast_ref::<Column>().map(|c| c.index())
+ //
+ // Only do this for modes that produce final values. Partial state is
+ // not unique on the group keys: a partial stage can emit the same
group
+ // more than once (early emission under memory pressure, skipping
partial
+ // aggregation), and operators that merge partitions (e.g.
+ // `RepartitionExec`, `SortPreservingMergeExec`) keep constraints. A
stale
+ // `Unique` lets an ordering on the group keys wrongly satisfy a longer
+ // ordering, which removes a required sort.
Review Comment:
The old and new comments now read as one paragraph. Could they be merged and
shortened? For example:
```rust
// Only final output is unique on the group keys. Partial output can repeat a
// group within a partition (early emission, skipped partial aggregation) and
// across partitions once they are merged.
```
##########
datafusion/core/tests/physical_optimizer/enforce_sorting.rs:
##########
@@ -3635,3 +3635,96 @@ async fn
test_passthrough_wrapper_projection_keeps_ordering() -> Result<()> {
Ok(())
}
+
+/// A parent needs `Hash(k)` and `[k, sum]` within each partition, above a
+/// Partial aggregate over input sorted on `k`. With `prefer_existing_sort`,
the
+/// optimizer replaces the sort with an order-preserving repartition only when
+/// the repartitioned output satisfies `[k, sum]`. Partial aggregate output is
+/// not unique on `k` across input partitions, so the sort must stay.
+#[tokio::test]
+async fn test_keep_sort_above_repartitioned_partial_aggregate() -> Result<()> {
+ use arrow::array::AsArray;
+ use arrow::datatypes::{Int32Type, Int64Type};
+ use datafusion_functions_aggregate::sum::sum_udaf;
+ use datafusion_physical_expr::aggregate::AggregateExprBuilder;
+ use datafusion_physical_expr_common::sort_expr::LexRequirement;
+ use datafusion_physical_optimizer::output_requirements::OutputRequirements;
+ use datafusion_physical_plan::aggregates::{
+ AggregateExec, AggregateMode, PhysicalGroupBy,
+ };
+ use datafusion_physical_plan::common::collect;
+ use datafusion_physical_plan::test::TestMemoryExec;
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ // Both partitions are sorted on `k` and contain keys 1 and 2.
+ let partitions = [
+ vec![record_batch!(("k", Int32, [1, 2]), ("v", Int64, [9, 5]))?],
+ vec![record_batch!(("k", Int32, [1, 2]), ("v", Int64, [1, 7]))?],
+ ];
+ let input = TestMemoryExec::try_new(&partitions, Arc::clone(&schema),
None)?
+ .try_with_sort_information(vec![[sort_expr("k", &schema)].into()])?;
+ let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input)));
+ let partial: Arc<dyn ExecutionPlan> = Arc::new(AggregateExec::try_new(
+ AggregateMode::Partial,
+ PhysicalGroupBy::new_single(vec![(col("k", &schema)?,
"k".to_string())]),
+ vec![Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("sum_v")
+ .build()?,
+ )],
+ vec![None],
+ input,
+ Arc::clone(&schema),
+ )?);
+
+ let k: Arc<dyn PhysicalExpr> = Arc::new(Column::new("k", 0));
+ let sum: Arc<dyn PhysicalExpr> =
+ Arc::new(Column::new(partial.schema().field(1).name(), 1));
+ let requirement = LexRequirement::new([
+ PhysicalSortRequirement::new(Arc::clone(&k),
Some(SortOptions::default())),
+ PhysicalSortRequirement::new(sum, Some(SortOptions::default())),
+ ])
+ .unwrap();
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(OutputRequirementExec::new(
+ partial,
+ Some(OrderingRequirements::new(requirement)),
+ Distribution::KeyPartitioned(vec![k]),
+ None,
+ ));
+
+ let mut config = ConfigOptions::new();
+ config.execution.target_partitions = 2;
+ config.optimizer.prefer_existing_sort = true;
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ let optimized = OutputRequirements::new_remove_mode().optimize(optimized,
&config)?;
+ assert_snapshot!(displayable(optimized.as_ref()).indent(true), @r"
+ SortExec: expr=[k@0 ASC, sum_v[sum]@1 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([k@0], 2), input_partitions=2
+ AggregateExec: mode=Partial, gby=[k@0 as k], aggr=[sum_v],
group_clustering_mode=Full
+ DataSourceExec: partitions=2, partition_sizes=[1, 1],
output_ordering=k@0 ASC
+ ");
+
+ let task_ctx = Arc::new(TaskContext::default());
Review Comment:
The snapshot above already shows the `SortExec` is kept. Executing the plan
only re-tests `SortExec` itself, so this loop and the `AsArray`, `Int32Type`,
`Int64Type` and `collect` imports could be dropped.
--
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]
