diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 45dabaac90f59..0e3a353621261 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/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 = 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 = Arc::new(Column::new("k", 0)); + let sum: Arc = + 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 = 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()); + for partition in 0..optimized.output_partitioning().partition_count() { + let batches = + collect(optimized.execute(partition, Arc::clone(&task_ctx))?).await?; + let rows: Vec<(i32, i64)> = batches + .iter() + .flat_map(|batch| { + let k = batch.column(0).as_primitive::(); + let sum = batch.column(1).as_primitive::(); + (0..batch.num_rows()).map(move |i| (k.value(i), sum.value(i))) + }) + .collect(); + assert!( + rows.is_sorted(), + "partition {partition} is not sorted: {rows:?}" + ); + } + + Ok(()) +} diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index f2b5e70f80003..46f4a1e07c06e 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/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::().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. + if mode.output_mode() == AggregateOutputMode::Final { + 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::().map(|c| c.index()) + }) }) - }) - .collect(), - ); - constraints.push(new_constraint); - eq_properties = - eq_properties.with_constraints(Constraints::new_unverified(constraints)); + .collect(), + ); + constraints.push(new_constraint); + eq_properties = + eq_properties.with_constraints(Constraints::new_unverified(constraints)); + } // Get output partitioning: let input_partitioning = input.output_partitioning().clone(); @@ -3416,6 +3426,8 @@ mod tests { use crate::projection::ProjectionExec; use crate::repartition::RepartitionExec; + use crate::sorts::sort::SortExec; + use crate::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_expr::projection::ProjectionExpr; use futures::{FutureExt, Stream, StreamExt}; use insta::{allow_duplicates, assert_snapshot}; @@ -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<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int64, false), + ])); + let batch = |k: Vec, v: Vec| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(k)), Arc::new(Int64Array::from(v))], + ) + }; + // Both partitions are sorted on `k` and contain keys 1 and 2. + let partitions = [ + vec![batch(vec![1, 2], vec![9, 5])?], + vec![batch(vec![1, 2], vec![1, 7])?], + ]; + let k_ordering = + LexOrdering::new([PhysicalSortExpr::new_default(col("k", &schema)?)]) + .unwrap(); + let input = TestMemoryExec::try_new(&partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![k_ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + let group_by = + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".to_string())]); + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum_v") + .build()?, + )]; + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregates.clone(), + vec![None], + input, + Arc::clone(&schema), + )?); + assert_eq!(partial.group_clustering_mode(), &GroupClusteringMode::Full); + assert!( + partial + .properties() + .equivalence_properties() + .constraints() + .is_empty(), + "partial aggregate declared {:?}", + partial.properties().equivalence_properties().constraints() + ); + + let k_sum_ordering = LexOrdering::new([ + PhysicalSortExpr::new_default(Arc::new(Column::new("k", 0))), + PhysicalSortExpr::new_default(Arc::new(Column::new( + partial.schema().field(1).name(), + 1, + ))), + ]) + .unwrap(); + let merged: Arc = Arc::new(SortPreservingMergeExec::new( + LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(Column::new( + "k", 0, + )))]) + .unwrap(), + Arc::clone(&partial) as _, + )); + assert!( + !merged + .equivalence_properties() + .ordering_satisfy(k_sum_ordering.clone())? + ); + let sort = Arc::new(SortExec::new(k_sum_ordering, merged)); + let task_ctx = Arc::new(TaskContext::default()); + let output = collect(sort.execute(0, task_ctx)?).await?; + assert_snapshot!(batches_to_string(&output), @r" ++---+------------+ +| k | sum_v[sum] | ++---+------------+ +| 1 | 1 | +| 1 | 9 | +| 2 | 5 | +| 2 | 7 | ++---+------------+ +"); + + // The final aggregate's output is unique on the group keys. + let final_agg = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggregates, + vec![None], + Arc::new(CoalescePartitionsExec::new(partial)), + Arc::clone(&schema), + )?; + assert_eq!( + final_agg + .properties() + .equivalence_properties() + .constraints(), + &Constraints::new_unverified(vec![Constraint::Unique(vec![0])]) + ); + + Ok(()) + } + #[tokio::test] async fn constant_grouping_expr_is_not_a_completion_boundary() -> Result<()> { let schema = Arc::new(Schema::new(vec![