diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0943c780e117b..6d69b7de01a07 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -37,8 +37,9 @@ //! //! ### Group clustering optimization //! -//! When the input is ordered by group keys, rows are clustered by those keys. -//! The clustered paths use this guarantee to emit completed groups early. +//! When the input is ordered by group keys, or its equivalence properties declare +//! them grouped, rows are clustered by those keys. The clustered paths use this +//! guarantee to emit completed groups early. //! //! ```text //! AggregateExec (final, clustered) @@ -1009,7 +1010,6 @@ impl AggregateExec { LexRequirement::new(new_requirements).map(OrderingRequirements::new_soft); // Constant expressions never change, so they cannot mark a completed group. - // Exclude them from both the ordering indices and the group expression count. // If our aggregation has grouping sets then our base grouping exprs will // be expanded based on the flags in `group_by.groups` where for each // group we swap the grouping expr for `null` if the flag is `true` @@ -1025,13 +1025,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 != [vec![false; group_by.expr.len()]] + || mode == AggregateMode::PartialReduce + { + GroupClusteringMode::None + } else if has_non_constant_groupby_exprs + && input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())? { GroupClusteringMode::Full } else if !indices.is_empty() { @@ -1040,12 +1045,6 @@ impl AggregateExec { 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; - } - // construct a map from the input expression to the output expression of the Aggregation group by let group_expr_mapping = ProjectionMapping::try_new(group_by.expr.clone(), &input.schema())?; @@ -2457,7 +2456,7 @@ impl ExecutionPlan for AggregateExec { metrics: _, // Derived at construction from the input ordering and `group_by`. required_input_ordering: _, - // Derived at construction from the input ordering and `group_by`. + // Derived at construction from the input properties and `group_by`. group_clustering_mode: _, // Derived at construction by `Self::compute_properties`. cache: _, @@ -3372,7 +3371,8 @@ mod tests { use crate::test::TestMemoryExec; use crate::test::assert_is_pending; use crate::test::exec::{ - BlockingExec, PanicExec, StatisticsExec, assert_strong_count_converges_to_zero, + BarrierExec, BlockingExec, PanicExec, StatisticsExec, + assert_strong_count_converges_to_zero, }; use arrow::array::AsArray; @@ -3381,7 +3381,7 @@ mod tests { Int64Array, NullArray, StringArray, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; - use arrow::datatypes::{Float64Type, Int32Type, Int64Type, UInt32Type}; + use arrow::datatypes::{Float64Type, Int32Type, Int64Type, TimeUnit, UInt32Type}; use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_common::{DataFusionError, assert_contains, internal_err}; use datafusion_execution::config::SessionConfig; @@ -3394,6 +3394,7 @@ mod tests { Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, Operator, Signature, Volatility, }; + use datafusion_functions::datetime::date_bin; use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; use datafusion_functions_aggregate::array_agg::array_agg_udaf; use datafusion_functions_aggregate::average::avg_udaf; @@ -3402,10 +3403,9 @@ mod tests { use datafusion_functions_aggregate::median::median_udaf; use datafusion_functions_aggregate::min_max::min_udaf; use datafusion_functions_aggregate::sum::sum_udaf; - use datafusion_physical_expr::Partitioning; - use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::{Literal, NotExpr, binary}; + use datafusion_physical_expr::{Partitioning, PhysicalSortExpr, ScalarFunctionExpr}; use crate::projection::ProjectionExec; use crate::repartition::RepartitionExec; @@ -5565,7 +5565,7 @@ mod tests { } #[test] - fn partial_reduce_does_not_advertise_input_ordering() -> Result<()> { + fn partial_reduce_does_not_use_group_clustering() -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Float64, false), @@ -5574,39 +5574,251 @@ mod tests { Column::new("a", 0), ))]) .unwrap(); - let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? - .try_with_sort_information(vec![ordering])?; - let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); - assert!( - input.properties().output_ordering().is_some(), - "test setup: the input is ordered by the group key" - ); + let ordered_input = + TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let grouped_input = + TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_grouping_information(vec![vec![col("a", &schema)?]])?; - let partial_reduce = AggregateExec::try_new( - AggregateMode::PartialReduce, - PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), - vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("SUM(b)") - .build()?, - )], - vec![None], - input, + for input in [ordered_input, grouped_input] { + assert!( + input + .properties() + .equivalence_properties() + .grouping_satisfy([col("a", &schema)?])? + ); + let partial_reduce = AggregateExec::try_new( + AggregateMode::PartialReduce, + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )], + vec![None], + Arc::new(input), + Arc::clone(&schema), + )?; + + assert_eq!( + partial_reduce.group_clustering_mode(), + &GroupClusteringMode::None + ); + assert_eq!(partial_reduce.maintains_input_order(), vec![false]); + assert!( + partial_reduce.properties().output_ordering().is_none(), + "partial reduce advertised an ordering it does not maintain: {:?}", + partial_reduce.properties().output_ordering() + ); + + let stream = partial_reduce.execute_typed(0, &new_migrated_hash_ctx(1024))?; + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + } + + Ok(()) + } + + #[tokio::test] + async fn multiple_group_masks_do_not_use_group_clustering() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let a = col("a", &schema)?; + let b = col("b", &schema)?; + let batch = RecordBatch::try_new( Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 2, 1])), + Arc::new(Int32Array::from(vec![20, 20, 10])), + Arc::new(Int64Array::from(vec![10, 20, 40])), + ], )?; + let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?; + let grouped_input = input + .clone() + .try_with_grouping_information(vec![vec![Arc::clone(&a), Arc::clone(&b)]])?; + // Repeating a mask repeats the keys, so input grouping does not prove + // contiguity after the masks are expanded. + let group_by = PhysicalGroupBy::new( + vec![(a, "a".to_string()), (b, "b".to_string())], + vec![], + vec![vec![false, false]; 2], + false, + ); + let aggr_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + ); - assert_eq!( - partial_reduce.group_clustering_mode(), - &GroupClusteringMode::None + for input in [input, grouped_input] { + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + Arc::new(input), + Arc::clone(&schema), + )?; + let output = + collect(aggregate.execute(0, new_migrated_hash_ctx(1024))?).await?; + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&output), @r" + +---+----+------------+ + | a | b | SUM(value) | + +---+----+------------+ + | 1 | 10 | 80 | + | 2 | 20 | 60 | + +---+----+------------+ + "); + } + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + } + Ok(()) + } + + #[tokio::test] + async fn nulling_group_mask_does_not_use_group_clustering() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let a = col("a", &schema)?; + let b = col("b", &schema)?; + let batches = [(1, 1, 10), (2, 2, 20), (3, 1, 30)] + .into_iter() + .map(|(a, b, value)| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![a])), + Arc::new(Int32Array::from(vec![b])), + Arc::new(Int64Array::from(vec![value])), + ], + ) + }) + .collect::, _>>()?; + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)?; + let grouped_input = input + .clone() + .try_with_grouping_information(vec![vec![Arc::clone(&a), Arc::clone(&b)]])?; + // Nulling a turns distinct input tuples into (NULL, 1), (NULL, 2), + // (NULL, 1), which are no longer contiguous. + let group_by = PhysicalGroupBy::new( + vec![(a, "a".to_string()), (b, "b".to_string())], + vec![ + (lit(ScalarValue::Int32(None)), "a".to_string()), + (lit(ScalarValue::Int32(None)), "b".to_string()), + ], + vec![vec![true, false]], + false, ); - assert_eq!(partial_reduce.maintains_input_order(), vec![false]); - assert!( - partial_reduce.properties().output_ordering().is_none(), - "partial reduce advertised an ordering it does not maintain: {:?}", - partial_reduce.properties().output_ordering() + let aggr_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, ); + for input in [input, grouped_input] { + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + Arc::new(input), + Arc::clone(&schema), + )?; + let output = + collect(aggregate.execute(0, new_migrated_hash_ctx(1024))?).await?; + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&output), @r" + +---+---+------------+ + | a | b | SUM(value) | + +---+---+------------+ + | | 1 | 40 | + | | 2 | 20 | + +---+---+------------+ + "); + } + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + } + Ok(()) + } + + #[tokio::test] + async fn constant_only_grouping_uses_final_emission() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int64Array::from(vec![10, 20])), + ], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![batch.clone(), batch]], + Arc::clone(&schema), + None, + )?; + let predicate = binary(col("key", &schema)?, Operator::Eq, lit(1i32), &schema)?; + let input: Arc = + Arc::new(FilterExecBuilder::new(predicate, input).build()?); + + // Cover both a literal and a column known to be constant after filtering. + for key in [lit(1i32), col("key", &schema)?] { + assert!( + input + .equivalence_properties() + .grouping_satisfy([Arc::clone(&key)])? + ); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![(key, "key".to_string())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + )], + vec![None], + Arc::clone(&input), + Arc::clone(&schema), + )?; + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + assert_eq!(aggregate.cache().emission_type, EmissionType::Final); + + let stream = aggregate.execute_typed(0, &new_migrated_hash_ctx(1024))?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let output = collect(stream.into()).await?; + allow_duplicates! { + assert_snapshot!(batches_to_string(&output), @r" + +-----+------------+ + | key | SUM(value) | + +-----+------------+ + | 1 | 60 | + +-----+------------+ + "); + } + } Ok(()) } @@ -5961,37 +6173,38 @@ mod tests { } #[tokio::test] - async fn unsorted_contiguous_groups_use_final_emission() -> Result<()> { + async fn unsorted_contiguous_groups_use_incremental_emission() -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Int32, false), Field::new("time_bin", DataType::Int64, false), Field::new("value", DataType::Int64, false), ])); - // Two sorted logical runs are emitted as batches in one DataFusion - // partition. Every distinct grouping tuple occupies one contiguous range, - // but tuple order resets at the batch boundary, so (key, time_bin) is not - // globally sorted. + // Every grouping tuple occupies one contiguous range in a single + // partition, with (2, 20) spanning both batches. Tuple order resets + // within the second batch, so (key, time_bin) is not globally sorted. let input_batches = vec![ RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(vec![1, 1, 2, 2])), - Arc::new(Int64Array::from(vec![20, 20, 20, 20])), - Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + Arc::new(Int32Array::from(vec![1, 1, 2])), + Arc::new(Int64Array::from(vec![20, 20, 20])), + Arc::new(Int64Array::from(vec![10, 20, 30])), ], )?, RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(vec![1, 1, 2, 2])), - Arc::new(Int64Array::from(vec![0, 0, 0, 0])), - Arc::new(Int64Array::from(vec![50, 60, 70, 80])), + Arc::new(Int32Array::from(vec![2, 1, 1, 2, 2])), + Arc::new(Int64Array::from(vec![20, 0, 0, 0, 0])), + Arc::new(Int64Array::from(vec![40, 50, 60, 70, 80])), ], )?, ]; + let key = col("key", &schema)?; + let time_bin = col("time_bin", &schema)?; let group_by = PhysicalGroupBy::new_single(vec![ - (col("key", &schema)?, "key".to_string()), - (col("time_bin", &schema)?, "time_bin".to_string()), + (Arc::clone(&key), "key".to_string()), + (Arc::clone(&time_bin), "time_bin".to_string()), ]); let aggr_expr = Arc::new( AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) @@ -5999,33 +6212,65 @@ mod tests { .alias("SUM(value)") .build()?, ); - let input: Arc = - TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; - assert_eq!(input.output_partitioning().partition_count(), 1); + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + eq_properties.add_groupings([vec![key, time_bin]]); + let input = Arc::new( + BarrierExec::new(vec![input_batches], Arc::clone(&schema)) + .without_start_barrier() + .with_batch_barrier() + .with_equivalence_properties(eq_properties) + .with_log(false), + ); + assert_eq!( + input.properties().output_partitioning().partition_count(), + 1 + ); let aggregate = AggregateExec::try_new( AggregateMode::Single, group_by, vec![aggr_expr], vec![None], - input, + Arc::clone(&input) as Arc, schema, )?; assert_eq!( aggregate.group_clustering_mode(), - &GroupClusteringMode::None + &GroupClusteringMode::Full ); - // This captures the behavior before #24438. When the source can declare - // `(key, time_bin)` group-contiguous, the corresponding case can use - // `EmissionType::Incremental`. - assert_eq!(aggregate.cache().emission_type, EmissionType::Final); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + assert!( + aggregate + .cache() + .equivalence_properties() + .geq_class() + .is_empty() + ); + assert!(aggregate.cache().output_ordering().is_none()); let task_ctx = new_migrated_hash_ctx(1024); let stream = aggregate.execute_typed(0, &task_ctx)?; - assert!(matches!(stream, StreamType::SingleHash(_))); - let stream: SendableRecordBatchStream = stream.into(); - let output = collect(stream).await?; + assert!(matches!(stream, StreamType::ClusteredSingleAggregate(_))); + let mut stream: SendableRecordBatchStream = stream.into(); + input.wait_batch().await; + // The second batch and EOF remain blocked until the first completed + // group has been emitted. + let first = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("completed group was not emitted before the next input batch") + .expect("stream ended before the next input batch")?; + assert_snapshot!(batches_to_string(std::slice::from_ref(&first)), @r" ++-----+----------+------------+ +| key | time_bin | SUM(value) | ++-----+----------+------------+ +| 1 | 20 | 30 | ++-----+----------+------------+ +"); + input.wait_batch().await; + let mut output = vec![first]; + output.extend(collect(stream).await?); assert_snapshot!(batches_to_sort_string(&output), @r" +-----+----------+------------+ | key | time_bin | SUM(value) | @@ -6040,6 +6285,158 @@ mod tests { Ok(()) } + #[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, values: Vec| { + 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 = 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(()) + } + + #[test] + fn grouped_date_bin_projects_to_aggregate() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time", DataType::Timestamp(TimeUnit::Second, None), false), + ])); + let time_bin_expr = || -> Result> { + Ok(Arc::new(ScalarFunctionExpr::try_new( + date_bin(), + vec![ + lit(ScalarValue::new_interval_dt(0, 10_000)), + col("time", &schema)?, + ], + &schema, + Arc::new(ConfigOptions::default()), + )?)) + }; + let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_grouping_information(vec![vec![ + col("key", &schema)?, + time_bin_expr()?, + ]])?; + let projection = ProjectionExec::try_new( + [ + ProjectionExpr::new(col("key", &schema)?, "key"), + // Build this expression independently from the source + // assertion so the test exercises semantic expression matching. + ProjectionExpr::new(time_bin_expr()?, "time_bin"), + ], + Arc::new(input), + )?; + + let projected_schema = projection.schema(); + let key = col("key", &projected_schema)?; + let time_bin = col("time_bin", &projected_schema)?; + assert!( + projection + .properties() + .equivalence_properties() + .grouping_satisfy([Arc::clone(&key), Arc::clone(&time_bin)])? + ); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + (key, "key".to_string()), + (time_bin, "time_bin".to_string()), + ]), + vec![], + vec![], + Arc::new(projection), + projected_schema, + )?; + + assert!(aggregate.input().output_ordering().is_none()); + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::Full + ); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + Ok(()) + } + /// Ensures for ordered input, `ClusteredPartialAggregateStream` is used. #[tokio::test] async fn ordered_partial_aggregate_planning() -> Result<()> { diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 2f77f903c9ad1..f8648f6cdf0a6 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -318,6 +318,9 @@ pub struct BarrierExec { /// all streams wait on this barrier to produce start_data_barrier: Option>, + /// each stream waits on this barrier before sending each batch + batch_barrier: Option>, + /// the stream wait for this to return Poll::Ready(None) finish_barrier: Option>, @@ -336,6 +339,7 @@ impl BarrierExec { data, schema, start_data_barrier: barrier, + batch_barrier: None, cache: Arc::new(cache), finish_barrier: None, log: true, @@ -352,6 +356,30 @@ impl BarrierExec { self } + /// Wait for [`Self::wait_batch`] before sending each batch. + pub fn with_batch_barrier(mut self) -> Self { + self.batch_barrier = Some(Arc::new(Barrier::new(self.data.len() + 1))); + self + } + + /// Release the next batch in each partition. + pub async fn wait_batch(&self) { + self.batch_barrier + .as_ref() + .expect("Must only be called when having a batch barrier") + .wait() + .await; + } + + /// Declare the equivalence properties of the test data. + pub fn with_equivalence_properties( + mut self, + properties: EquivalenceProperties, + ) -> Self { + self.cache = Arc::new(self.cache.as_ref().clone().with_eq_properties(properties)); + self + } + pub fn with_finish_barrier(mut self) -> Self { let barrier = Arc::new(( // wait for all streams and the input @@ -486,6 +514,7 @@ impl ExecutionPlan for BarrierExec { // task simply sends data in order after barrier is reached let data = self.data[partition].clone(); let start_barrier = self.start_data_barrier.as_ref().map(Arc::clone); + let batch_barrier = self.batch_barrier.as_ref().map(Arc::clone); let finish_barrier = self.finish_barrier.as_ref().map(Arc::clone); let log = self.log; let tx = builder.tx(); @@ -497,6 +526,9 @@ impl ExecutionPlan for BarrierExec { barrier.wait().await; } for batch in data { + if let Some(barrier) = &batch_barrier { + barrier.wait().await; + } if log { println!("Partition {partition} sending batch"); } diff --git a/docs/source/library-user-guide/upgrading/57.0.0.md b/docs/source/library-user-guide/upgrading/57.0.0.md index 2f63a9d39a392..7b9b94fa039f5 100644 --- a/docs/source/library-user-guide/upgrading/57.0.0.md +++ b/docs/source/library-user-guide/upgrading/57.0.0.md @@ -38,3 +38,14 @@ parents and siblings unchanged. The values themselves are shared through `Arc`. Update logical-plan snapshots if needed: generated lambda parameter names are local to lexical scopes, so sibling lambdas may reuse names. + +### `GroupClusteringMode::Full` no longer implies sorted input + +`AggregateExec` can now select `GroupClusteringMode::Full` when the input's +equivalence properties declare the complete `GROUP BY` tuple as grouped +(contiguous within each partition), even if the input is not sorted. + +`Full` describes group contiguity only. Code that previously used +`group_clustering_mode() == GroupClusteringMode::Full` to infer sorted input or +output should use `properties().output_ordering()` on the input or aggregate +plan, respectively.