viirya commented on code in PR #26819:
URL: https://github.com/apache/datafusion/pull/26819#discussion_r4407116233
##########
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:
Merged into the shorter comment you suggested in 0ddbaa824.
##########
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:
Agreed. Removed the test and its two imports, and added the constraint
assertion to `group_clustering_preserves_sort_options`. Without the fix, the 8
`Partial` cases of that test fail and the 8 `Single` cases pass.
##########
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:
Removed the loop and its imports; the snapshot is the assertion now.
--
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]
