diff --git a/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs b/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs index 9e63c341c92d9..c0b7956d31b36 100644 --- a/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs +++ b/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs @@ -35,12 +35,14 @@ use datafusion_physical_expr::expressions::{col, lit}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::combine_partial_final_agg::CombinePartialFinalAggregate; -use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, }; use datafusion_physical_plan::displayable; +use datafusion_physical_plan::execution_plan::EmissionType; use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; /// Runs the CombinePartialFinalAggregate optimizer and asserts the plan against the expected macro_rules! assert_optimized { @@ -233,6 +235,50 @@ fn aggregations_with_group_combined() -> datafusion_common::Result<()> { Ok(()) } +#[test] +fn partition_disjoint_aggregation_combines_and_streams() -> datafusion_common::Result<()> +{ + let schema = schema(); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![col("c", &schema)?])?; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("Sum(b)") + .build()?, + )]; + let partial = partial_aggregate_exec( + Arc::new(source), + PhysicalGroupBy::new_single(vec![(col("c", &schema)?, "c".to_string())]), + aggr_expr.clone(), + ); + let final_group_by = PhysicalGroupBy::new_single(vec![( + col("c", &partial.schema())?, + "c".to_string(), + )]); + let plan = final_aggregate_exec(partial, final_group_by, aggr_expr); + + let optimized = + CombinePartialFinalAggregate::default().optimize(plan, &ConfigOptions::new())?; + let aggregate = optimized + .downcast_ref::() + .expect("adjacent partial and final aggregates should combine"); + assert_eq!(aggregate.mode(), &AggregateMode::Single); + assert_eq!( + aggregate.input_order_mode(), + &datafusion_physical_plan::InputOrderMode::Linear + ); + assert_eq!(optimized.pipeline_behavior(), EmissionType::Incremental); + + let display = displayable(optimized.as_ref()).indent(true).to_string(); + assert_snapshot!(display.trim(), @r" +AggregateExec: mode=Single, gby=[c@2 as c], aggr=[Sum(b)], group_completion_mode=Full + DataSourceExec: partitions=1, partition_sizes=[0], group_contiguous=[c@2] +"); + + Ok(()) +} + #[test] fn aggregations_with_limit_combined() -> datafusion_common::Result<()> { let schema = schema(); diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 2dbacf1d898ac..35bcb23122747 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -72,6 +72,7 @@ use datafusion_physical_plan::joins::utils::JoinOn; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::test::TestMemoryExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, @@ -2617,6 +2618,29 @@ fn added_repartition_to_single_partition() -> Result<()> { Ok(()) } +#[test] +fn partition_disjoint_aggregate_preserves_source_partitioning() -> Result<()> { + let schema = schema(); + let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![col("a", &schema)?])?; + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), + vec![], + vec![], + Arc::new(input), + schema, + )?) as Arc; + + let plan = TestConfig::default().to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + assert_plan!(plan, @r" + AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[], group_completion_mode=Full + DataSourceExec: partitions=1, partition_sizes=[0], group_contiguous=[a@0] + "); + + Ok(()) +} + #[test] fn repartition_deepest_node() -> Result<()> { let alias = vec![("a".to_string(), "a".to_string())]; diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 741010c595197..c947de3957f6f 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -165,6 +165,20 @@ pub trait DataSource: Any + Send + Sync + Debug { fn output_partitioning(&self) -> Partitioning; fn eq_properties(&self) -> EquivalenceProperties; + + /// Components of one composite key whose tuple values occur in one + /// contiguous range within each output stream. + /// + /// See [`ExecutionPlan::group_contiguous_exprs`] for the full contract. + /// Implementations must return expressions in terms of the schema returned + /// by [`Self::eq_properties`]. Any implementation that returns a rewritten + /// source from [`Self::repartitioned`] or [`Self::try_swapping_with_projection`] + /// must remap, preserve, or clear these expressions so the contract remains + /// valid for the rewritten output streams. + fn group_contiguous_exprs(&self) -> &[Arc] { + &[] + } + fn scheduling_type(&self) -> SchedulingType { SchedulingType::NonCooperative } @@ -384,7 +398,20 @@ impl DisplayAs for DataSourceExec { } DisplayFormatType::TreeRender => {} } - self.data_source.fmt_as(t, f) + self.data_source.fmt_as(t, f)?; + if matches!(t, DisplayFormatType::Default | DisplayFormatType::Verbose) + && !self.group_contiguous_exprs().is_empty() + { + write!( + f, + ", group_contiguous=[{}]", + self.group_contiguous_exprs() + .iter() + .map(ToString::to_string) + .join(", ") + )?; + } + Ok(()) } } @@ -674,6 +701,7 @@ impl DataSourceExec { EmissionType::Incremental, Boundedness::Bounded, ) + .with_group_contiguous_exprs(data_source.group_contiguous_exprs().to_vec()) .with_scheduling_type(data_source.scheduling_type()) } diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index d7ee5dace30cc..3bcd00b3482d4 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -615,6 +615,17 @@ pub mod tests { self.dynamic_expressions = dynamic_expressions; self } + + pub fn with_group_contiguous_exprs( + mut self, + group_contiguous_exprs: Vec>, + ) -> Self { + self.props = Arc::new( + PlanProperties::clone(&self.props) + .with_group_contiguous_exprs(group_contiguous_exprs), + ); + self + } } impl DisplayAs for EmptyExec { @@ -787,6 +798,30 @@ pub mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_group_contiguous_exprs() -> Result<()> { + use datafusion_physical_expr::expressions::col; + + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("key", arrow::datatypes::DataType::Int32, false), + ])); + let expression = col("key", &schema)?; + let original_plan = Arc::new( + EmptyExec::new(schema) + .with_group_contiguous_exprs(vec![Arc::clone(&expression)]), + ); + + let mut ffi_plan = FFI_ExecutionPlan::new(original_plan, None); + ffi_plan.library_marker_id = crate::mock_foreign_marker_id; + let foreign_plan: Arc = (&ffi_plan).try_into()?; + + let group_contiguous_exprs = foreign_plan.group_contiguous_exprs(); + assert_eq!(group_contiguous_exprs.len(), 1); + assert!(group_contiguous_exprs[0].eq(&expression)); + + Ok(()) + } + #[test] fn test_ffi_execution_plan_children() -> Result<()> { let schema = Arc::new(arrow::datatypes::Schema::new(vec![ diff --git a/datafusion/ffi/src/plan_properties.rs b/datafusion/ffi/src/plan_properties.rs index 09ef26af32349..9b70fac24aea5 100644 --- a/datafusion/ffi/src/plan_properties.rs +++ b/datafusion/ffi/src/plan_properties.rs @@ -28,6 +28,7 @@ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use stabby::vec::Vec as SVec; use crate::arrow_wrappers::WrappedSchema; +use crate::physical_expr::FFI_PhysicalExpr; use crate::physical_expr::partitioning::FFI_Partitioning; use crate::physical_expr::sort::FFI_PhysicalSortExpr; use crate::util::FFI_Option; @@ -63,6 +64,13 @@ pub struct FFI_PlanProperties { /// the foreign interface. See [`crate::get_library_marker_id`] and /// the crate's `README.md` for more information. pub library_marker_id: extern "C" fn() -> usize, + + /// Return the components of the plan's group-contiguous composite key. + /// + /// See [`datafusion_physical_plan::ExecutionPlan::group_contiguous_exprs`] + /// for the correctness contract. + pub group_contiguous_exprs: + unsafe extern "C" fn(plan: &Self) -> SVec, } struct PlanPropertiesPrivateData { @@ -109,6 +117,18 @@ unsafe extern "C" fn output_ordering_fn_wrapper( ordering.into() } +unsafe extern "C" fn group_contiguous_exprs_fn_wrapper( + properties: &FFI_PlanProperties, +) -> SVec { + properties + .inner() + .group_contiguous_exprs() + .iter() + .cloned() + .map(FFI_PhysicalExpr::from) + .collect() +} + unsafe extern "C" fn schema_fn_wrapper(properties: &FFI_PlanProperties) -> WrappedSchema { let schema: SchemaRef = Arc::clone(properties.inner().eq_properties.schema()); schema.into() @@ -145,6 +165,7 @@ impl From<&PlanProperties> for FFI_PlanProperties { release: release_fn_wrapper, private_data: Box::into_raw(private_data) as *mut c_void, library_marker_id: crate::get_library_marker_id, + group_contiguous_exprs: group_contiguous_exprs_fn_wrapper, } } } @@ -186,12 +207,16 @@ impl TryFrom for PlanProperties { let boundedness: Boundedness = unsafe { (ffi_props.boundedness)(&ffi_props).into() }; - Ok(PlanProperties::new( - eq_properties, - partitioning, - emission_type, - boundedness, - )) + let group_contiguous_exprs = + unsafe { (ffi_props.group_contiguous_exprs)(&ffi_props) } + .iter() + .map(>::from) + .collect(); + + Ok( + PlanProperties::new(eq_properties, partitioning, emission_type, boundedness) + .with_group_contiguous_exprs(group_contiguous_exprs), + ) } } @@ -273,16 +298,16 @@ mod tests { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); + let column = datafusion::physical_plan::expressions::col("a", &schema)?; let mut eqp = EquivalenceProperties::new(Arc::clone(&schema)); - let _ = eqp.reorder([PhysicalSortExpr::new_default( - datafusion::physical_plan::expressions::col("a", &schema)?, - )]); + let _ = eqp.reorder([PhysicalSortExpr::new_default(Arc::clone(&column))]); Ok(PlanProperties::new( eqp, Partitioning::RoundRobinBatch(3), EmissionType::Incremental, Boundedness::Bounded, - )) + ) + .with_group_contiguous_exprs(vec![column])) } fn create_range_test_props() -> Result { diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 83f842508ab2c..9e93ec1a43d0a 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -976,7 +976,7 @@ mod tests { let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?; assert_eq!( format!("{physical_plan:?}"), - "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None } }" + "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None, group_contiguous_exprs: [] } }" ); assert_eq!( diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index fbc3e83ba49fc..aa06dbdc21a8a 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -178,8 +178,9 @@ extern "C" fn construct_table_provider_factory( pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); - - let plan = Arc::new(EmptyExec::new(schema)); + let expression = datafusion_physical_expr::expressions::col("a", &schema).unwrap(); + let plan = + Arc::new(EmptyExec::new(schema).with_group_contiguous_exprs(vec![expression])); FFI_ExecutionPlan::new(plan, None) } diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 4067d7eb49b2a..02b61ab7fce5e 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -109,6 +109,21 @@ mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_group_contiguous_exprs_cross_library() + -> Result<(), DataFusionError> { + let module = get_module()?; + let plan = (module.create_empty_exec)(); + let plan: Arc = (&plan).try_into()?; + assert!(plan.is::()); + + let group_contiguous_exprs = plan.group_contiguous_exprs(); + assert_eq!(group_contiguous_exprs.len(), 1); + assert_eq!(group_contiguous_exprs[0].to_string(), "a@0"); + + Ok(()) + } + #[test] fn test_ffi_execution_plan_new_sets_runtimes_on_children() -> Result<(), DataFusionError> { diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index 93862df3b4236..b14bbfbfb7b09 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -152,6 +152,42 @@ mod tests { "); } + #[test] + fn test_cooperative_exec_preserves_partition_disjoint_aggregation() -> Result<()> { + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_plan::ExecutionPlanProperties; + use datafusion_physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, + }; + use datafusion_physical_plan::execution_plan::EmissionType; + use datafusion_physical_plan::test::TestMemoryExec; + + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![col("key", &schema)?])?; + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]), + vec![], + vec![], + Arc::new(input), + schema, + )?); + + let optimized = + EnsureCooperative::new().optimize(aggregate, &ConfigOptions::new())?; + assert_eq!(optimized.pipeline_behavior(), EmissionType::Incremental); + assert_snapshot!(displayable(optimized.as_ref()).indent(true).to_string(), @r" + AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[], group_completion_mode=Full + CooperativeExec + DataSourceExec: partitions=1, partition_sizes=[0], group_contiguous=[key@0] + "); + + Ok(()) + } + #[tokio::test] async fn test_optimizer_is_idempotent() { // Comprehensive idempotency test: verify f(f(...f(x))) = f(x) diff --git a/datafusion/physical-plan/benches/partial_ordering.rs b/datafusion/physical-plan/benches/partial_ordering.rs index bdadd6274b75e..f176cbab0a517 100644 --- a/datafusion/physical-plan/benches/partial_ordering.rs +++ b/datafusion/physical-plan/benches/partial_ordering.rs @@ -17,12 +17,27 @@ use std::sync::Arc; -use arrow::array::{ArrayRef, Int32Array}; +use arrow::array::{ArrayRef, Int32Array, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; use datafusion_physical_plan::aggregates::order::GroupOrderingPartial; +use datafusion_physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, collect}; use criterion::{Criterion, criterion_group, criterion_main}; const BATCH_SIZE: usize = 8192; +const LOGICAL_RUNS: usize = 16; +const GROUPS_PER_RUN: usize = 2048; +const ROWS_PER_GROUP: usize = BATCH_SIZE / GROUPS_PER_RUN; fn create_test_arrays(num_columns: usize) -> Vec { (0..num_columns) @@ -56,5 +71,108 @@ fn bench_new_groups(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_new_groups); +fn aggregate_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])) +} + +/// Build locally key-ordered logical runs whose key ranges are disjoint. The +/// runs are concatenated in descending range order so the resulting stream is +/// intentionally not globally sorted. +fn partition_disjoint_batches(schema: &SchemaRef) -> Vec { + (0..LOGICAL_RUNS) + .map(|run| { + let key_offset = (LOGICAL_RUNS - run - 1) * GROUPS_PER_RUN; + let keys = + (0..BATCH_SIZE).map(|row| (key_offset + row / ROWS_PER_GROUP) as i32); + let values = std::iter::repeat_n(1_i64, BATCH_SIZE); + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from_iter_values(keys)), + Arc::new(Int64Array::from_iter_values(values)), + ], + ) + .unwrap() + }) + .collect() +} + +fn aggregate_plan(partition_disjoint: bool) -> Arc { + let schema = aggregate_schema(); + let batches = partition_disjoint_batches(&schema); + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None).unwrap(); + let input: Arc = if partition_disjoint { + Arc::new( + input + .try_with_group_contiguous_keys(vec![col("key", &schema).unwrap()]) + .unwrap(), + ) + } else { + Arc::new(input) + }; + let group_by = PhysicalGroupBy::new_single(vec![( + col("key", &schema).unwrap(), + "key".to_string(), + )]); + let aggregate = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema).unwrap()]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build() + .unwrap(), + ); + + Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![aggregate], + vec![None], + input, + schema, + ) + .unwrap(), + ) +} + +fn bench_partition_disjoint_aggregate(c: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(BATCH_SIZE) + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + let mut group = c.benchmark_group("partition_disjoint_aggregate"); + + for (name, partition_disjoint) in [ + ("unordered_hash", false), + ("partition_disjoint_streaming", true), + ] { + let plan = aggregate_plan(partition_disjoint); + group.bench_function(name, |b| { + b.iter(|| { + let batches = runtime + .block_on(collect(Arc::clone(&plan), Arc::clone(&task_ctx))) + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + LOGICAL_RUNS * GROUPS_PER_RUN + ); + }); + }); + } + + group.finish(); +} + +criterion_group!( + benches, + bench_new_groups, + bench_partition_disjoint_aggregate +); criterion_main!(benches); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 4ced967a0977b..b65d3cc04b5da 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -15,8 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Common utilities for aggregate tables used in aggregations that inputs are ordered -//! by the groups. +//! Common utilities for aggregate tables that can identify completed group ranges. use std::marker::PhantomData; use std::sync::Arc; @@ -28,14 +27,13 @@ use datafusion_common::assert_or_internal_err; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::EmitTo; -use crate::InputOrderMode; use crate::PhysicalExpr; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, }; use crate::aggregates::grouped_hash_stream::create_group_accumulator; -use crate::aggregates::order::GroupOrdering; +use crate::aggregates::order::{GroupCompletionMode, GroupOrdering}; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, @@ -77,11 +75,12 @@ impl OrderedAggregateTableMetrics { /// Aggregate table shared by the ordered single, partial and final paths. /// -/// # Ordering optimization +/// # Group-completion optimization /// /// The table consumes input batches while `GroupOrdering` tracks which groups /// are proven complete. Completed groups can be emitted before the input stream -/// ends, which keeps memory bounded by the active ordered key range. +/// ends, which keeps memory bounded by the active key range. The guarantee can +/// come from physical sort order or from partition-disjoint group keys. /// /// # Single, partial and final variant difference /// @@ -113,7 +112,7 @@ impl OrderedAggregateTableMetrics { /// `OrderedAggrMode` selects the aggregate semantics. For example, /// `OrderedAggregateTable::::new(...)` consumes raw rows /// and emits partial states, while -/// `OrderedAggregateTable::::new_with_input_order(...)` +/// `OrderedAggregateTable::::new_with_group_completion(...)` /// consumes partial states and emits final values. /// /// Shared methods live on `impl`; single/partial/final behavior lives on @@ -184,7 +183,7 @@ impl OrderedAggregateTable { output_schema: SchemaRef, state_schema: SchemaRef, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, aggregate_mode: &AggregateMode, filters: Vec>>, metrics: OrderedAggregateTableMetrics, @@ -194,7 +193,7 @@ impl OrderedAggregateTable { "OrderedAggregateTable requires config batch_size >= 1" ); - let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_ordering = GroupOrdering::try_new_for_mode(group_completion_mode)?; let group_schema = agg.group_by.group_schema(input_schema)?; let group_values = new_group_values(group_schema, &group_ordering)?; let aggregate_arguments = aggregate_expressions( diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index d0d0c99bb5bd8..79fa968984e9b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Aggregate table for final aggregation when partial-state input is ordered. +//! Aggregate table for final aggregation with incrementally completed groups. //! //! See comments in [`super::ordered_partial_table`] for details. @@ -25,8 +25,8 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; -use crate::InputOrderMode; use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::order::GroupCompletionMode; use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase}; use super::common::HashAggregateAccumulator; @@ -42,12 +42,12 @@ use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics} /// /// See comments at [`OrderedAggregateTable`] for details. impl OrderedAggregateTable { - pub(in crate::aggregates) fn new_with_input_order( + pub(in crate::aggregates) fn new_with_group_completion( agg: &AggregateExec, input_schema: &SchemaRef, output_schema: SchemaRef, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, metrics: OrderedAggregateTableMetrics, ) -> Result { Self::new_for_mode( @@ -56,7 +56,7 @@ impl OrderedAggregateTable { output_schema, Arc::clone(input_schema), batch_size, - input_order_mode, + group_completion_mode, &AggregateMode::Final, vec![None; agg.aggr_expr.len()], metrics, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index 6ed93e59f3296..6e4cab7204058 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -15,13 +15,15 @@ // specific language governing permissions and limitations // under the License. -//! Aggregate table for partial aggregation when input is ordered by group keys. +//! Aggregate table for partial aggregation with incrementally completed groups. //! //! See the [`super::common_ordered`] comments for the high-level ideas. //! -//! This operator handles input that is ordered by group keys: +//! This operator handles input whose group-completion keys are contiguous: //! - Fully ordered: `GROUP BY a, b`, input is `ORDER BY a, b` //! - Partially ordered: `GROUP BY a, b`, input is `ORDER BY a` +//! - Partition-disjoint: logical source runs may reset key order, but completed +//! key tuples never recur in the same output stream //! //! When a group key combination is exhausted, this table eagerly flushes the //! completed groups to improve memory efficiency. @@ -68,7 +70,7 @@ impl OrderedAggregateTable { output_schema, state_schema, batch_size, - &agg.input_order_mode, + &agg.group_completion_mode, &AggregateMode::Partial, agg.filter_expr.iter().cloned().collect(), metrics, @@ -90,11 +92,11 @@ impl OrderedAggregateTable { } /// Emits the next batch of partial state rows for groups proven complete by - /// the input ordering. + /// the input's group-contiguity guarantee. /// - /// For example, when the query is `GROUP BY a` and the input is ordered by - /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3` - /// are complete and safe to emit. + /// For example, when the query is `GROUP BY a`, seeing a new `a` range means + /// the prior range is complete and safe to emit when `a` is ordered or + /// otherwise guaranteed contiguous. /// /// Key steps: /// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly. diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs index ce1ce647b46fe..4366b40e5066d 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Aggregate table for single aggregation when raw input is ordered. +//! Aggregate table for single aggregation with incrementally completed groups. //! //! See comments in [`super::ordered_partial_table`] for details. @@ -59,7 +59,7 @@ impl OrderedAggregateTable { output_schema, state_schema, batch_size, - &agg.input_order_mode, + &agg.group_completion_mode, &agg.mode, agg.filter_expr.iter().cloned().collect(), metrics, diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index c0253093c8a7b..23cd61e6e388c 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -504,7 +504,7 @@ impl GroupedHashAggregateStream { .collect::>() .join(", "); let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})"); - let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?; + let group_ordering = GroupOrdering::try_new_for_mode(&agg.group_completion_mode)?; let oom_mode = match (agg.mode, &group_ordering) { // In partial aggregation mode, always prefer to emit incomplete results early. (AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly, @@ -516,10 +516,10 @@ impl GroupedHashAggregateStream { { OutOfMemoryMode::Spill } - // For `GroupOrdering::Full`, the incoming stream is already sorted. This ensures the - // number of incomplete groups can be kept small at all times. If we still hit - // an out-of-memory condition, spilling to disk would not be beneficial since the same - // situation is likely to reoccur when reading back the spilled data. + // For `GroupOrdering::Full`, each group key is contiguous. This keeps the + // number of incomplete groups small even if successive keys are not sorted. + // If we still hit an out-of-memory condition, spilling to disk would not be + // beneficial since the same situation is likely to recur during replay. // Therefore, we fall back to simply reporting the error immediately. // This mode will also be used if the `DiskManager` is not configured to allow spilling // to disk. diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 2df5960188a2b..8fce12ad6d6b0 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -45,6 +45,7 @@ use super::aggregate_hash_table::{ AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker, PartialSkipMarker, }; +use super::order::GroupCompletionMode; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::skip_partial::SkipAggregationProbe; use crate::metrics::{ @@ -326,6 +327,7 @@ impl FinalSpillContext { let mut final_agg = agg.clone(); final_agg.input_order_mode = InputOrderMode::Sorted; + final_agg.group_completion_mode = GroupCompletionMode::Full; Ok(Self { final_agg, @@ -414,7 +416,7 @@ impl FinalSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0f1b718a3c8e3..a25abcf1ba1c5 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -62,6 +62,13 @@ //! See [`OrderedPartialAggregateStream`], [`OrderedFinalAggregateStream`], and //! [`OrderedSingleAggregateStream`] for details. //! +//! A source can also use +//! [`ExecutionPlan::group_contiguous_exprs`] +//! to declare that group-key tuples are contiguous without being globally +//! sorted. The aggregate reuses the ordered early-emission paths for group +//! completion, but does not advertise an output ordering. This guarantee is +//! consumed by the first aggregate after the source and optional projections. +//! //! Related configuration: //! //! - [`datafusion.execution.target_partitions`](datafusion_common::config::ExecutionOptions::target_partitions) @@ -203,6 +210,7 @@ use datafusion_physical_expr_common::sort_expr::{ use datafusion_expr::utils::AggregateOrderSensitivity; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use itertools::Itertools; +use order::GroupCompletionMode; use topk::hash_table::is_supported_hash_key_type; use topk::heap::is_supported_heap_type; @@ -866,6 +874,12 @@ pub struct AggregateExec { required_input_ordering: Option, /// Describes how the input is ordered relative to the group by columns input_order_mode: InputOrderMode, + /// Describes how the executor can detect completed groups. + /// + /// Input ordering establishes a corresponding completion mode. A + /// group-contiguous source key can establish a stronger completion mode + /// without changing `input_order_mode` or implying a global sort order. + group_completion_mode: GroupCompletionMode, cache: Arc, /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]), /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node. @@ -876,6 +890,45 @@ pub struct AggregateExec { dynamic_filter: Option>, } +/// Returns the group-completion mode established by a partition-disjoint key. +/// +/// If the complete disjoint key is present in `groupby_exprs`, each key value +/// identifies one contiguous range of input rows. An exact match establishes +/// full group completion. If the aggregate has additional keys, the disjoint +/// key establishes partial group completion: all groups in the prior key range +/// can be emitted together. +fn group_completion_from_partition_disjointness( + eq_properties: &EquivalenceProperties, + groupby_exprs: &[Arc], + group_contiguous_exprs: &[Arc], +) -> Option { + if group_contiguous_exprs.is_empty() + || group_contiguous_exprs.len() > groupby_exprs.len() + { + return None; + } + + let mut matched_group_exprs = vec![false; groupby_exprs.len()]; + let mut indices = Vec::with_capacity(group_contiguous_exprs.len()); + + for contiguous_expr in group_contiguous_exprs { + let (idx, _) = groupby_exprs.iter().enumerate().find(|(idx, group_expr)| { + !matched_group_exprs[*idx] + && eq_properties + .eq_group() + .exprs_equal(contiguous_expr, group_expr) + })?; + matched_group_exprs[idx] = true; + indices.push(idx); + } + + if indices.len() == groupby_exprs.len() { + Some(GroupCompletionMode::Full) + } else { + Some(GroupCompletionMode::Partial(indices)) + } +} + impl AggregateExec { /// Function used in `OptimizeAggregateOrder` optimizer rule, /// where we need parts of the new value, others cloned from the old one @@ -890,6 +943,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + group_completion_mode: self.group_completion_mode.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -910,6 +964,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + group_completion_mode: self.group_completion_mode.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -1032,6 +1087,20 @@ impl AggregateExec { input_order_mode = InputOrderMode::Linear; } + let group_completion_mode = if input_order_mode == InputOrderMode::Linear + && !group_by.has_grouping_set() + && group_by.is_single() + { + group_completion_from_partition_disjointness( + input_eq_properties, + &groupby_exprs, + input.group_contiguous_exprs(), + ) + .unwrap_or(GroupCompletionMode::None) + } else { + GroupCompletionMode::from(&input_order_mode) + }; + // 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())?; @@ -1039,13 +1108,13 @@ impl AggregateExec { let cache = if group_by.has_grouping_set() { Self::compute_grouping_set_properties(&input, Arc::clone(&schema)) } else { - Self::compute_properties( + Self::compute_properties_with_group_completion( &input, Arc::clone(&schema), &group_expr_mapping, group_by.is_true_no_grouping(), &mode, - &input_order_mode, + &group_completion_mode, aggr_expr.as_ref(), )? }; @@ -1062,6 +1131,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, + group_completion_mode, cache: Arc::new(cache), dynamic_filter: None, }; @@ -1254,7 +1324,7 @@ impl AggregateExec { fn should_use_partial_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::Partial - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() @@ -1265,7 +1335,7 @@ impl AggregateExec { _context: &TaskContext, ) -> bool { self.mode == AggregateMode::Partial - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() @@ -1276,7 +1346,7 @@ impl AggregateExec { self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned ) && self.limit_options_supported_by_hash_stream() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1289,7 +1359,7 @@ impl AggregateExec { self.mode == AggregateMode::PartialReduce && self.limit_options.is_none() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1299,7 +1369,7 @@ impl AggregateExec { self.mode, AggregateMode::Single | AggregateMode::SinglePartitioned ) && self.limit_options.is_none() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1309,7 +1379,7 @@ impl AggregateExec { self.mode, AggregateMode::Single | AggregateMode::SinglePartitioned ) && self.limit_options.is_none() - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1319,7 +1389,7 @@ impl AggregateExec { self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned ) && self.limit_options_supported_by_hash_stream() - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1384,6 +1454,27 @@ impl AggregateExec { mode: &AggregateMode, input_order_mode: &InputOrderMode, aggr_exprs: &[Arc], + ) -> Result { + let group_completion_mode = GroupCompletionMode::from(input_order_mode); + Self::compute_properties_with_group_completion( + input, + schema, + group_expr_mapping, + is_true_no_grouping, + mode, + &group_completion_mode, + aggr_exprs, + ) + } + + fn compute_properties_with_group_completion( + input: &Arc, + schema: SchemaRef, + group_expr_mapping: &ProjectionMapping, + is_true_no_grouping: bool, + mode: &AggregateMode, + group_completion_mode: &GroupCompletionMode, + aggr_exprs: &[Arc], ) -> Result { // Construct equivalence properties: let mut eq_properties = input @@ -1434,7 +1525,7 @@ impl AggregateExec { }; // TODO: Emission type and boundedness information can be enhanced here - let emission_type = if *input_order_mode == InputOrderMode::Linear { + let emission_type = if *group_completion_mode == GroupCompletionMode::None { EmissionType::Final } else { input.pipeline_behavior() @@ -1947,6 +2038,15 @@ impl DisplayAs for AggregateExec { if self.input_order_mode != InputOrderMode::Linear { write!(f, ", ordering_mode={:?}", self.input_order_mode)?; } + if self.group_completion_mode + != GroupCompletionMode::from(&self.input_order_mode) + { + write!( + f, + ", group_completion_mode={:?}", + self.group_completion_mode + )?; + } } DisplayFormatType::TreeRender => { let format_expr_with_alias = @@ -2084,6 +2184,16 @@ impl ExecutionPlan for AggregateExec { vec![self.input_order_mode != InputOrderMode::Linear] } + fn benefits_from_input_partitioning(&self) -> Vec { + // Repartitioning can increase aggregate parallelism, but it destroys a + // group-contiguous key. Preserve the source's deliberate partition + // layout when that key is what makes this aggregate streaming. + vec![ + self.group_completion_mode + == GroupCompletionMode::from(&self.input_order_mode), + ] + } + fn children(&self) -> Vec<&Arc> { vec![&self.input] } @@ -2327,6 +2437,8 @@ impl ExecutionPlan for AggregateExec { required_input_ordering: _, // Derived at construction from the input ordering and `group_by`. input_order_mode: _, + // Derived from input ordering or partition-disjoint source keys. + group_completion_mode: _, // Derived at construction by `Self::compute_properties`. cache: _, dynamic_filter, @@ -3205,6 +3317,7 @@ mod tests { use crate::filter::FilterExecBuilder; use crate::metrics::MetricValue; use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::stream::RecordBatchStreamAdapter; use crate::test::TestMemoryExec; use crate::test::assert_is_pending; use crate::test::exec::{ @@ -3213,10 +3326,11 @@ mod tests { use arrow::array::{ BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, - Int64Array, StructArray, UInt32Array, UInt64Array, + Int64Array, StructArray, TimestampSecondArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; - use arrow::datatypes::Int32Type; + use arrow::datatypes::{Int32Type, TimeUnit}; + use datafusion_common::config::ConfigOptions; use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_common::{DataFusionError, internal_err}; use datafusion_execution::config::SessionConfig; @@ -3227,6 +3341,7 @@ mod tests { Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, 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; @@ -3235,10 +3350,10 @@ 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}; + use datafusion_physical_expr::{Partitioning, ScalarFunctionExpr}; use crate::projection::ProjectionExec; use crate::repartition::RepartitionExec; @@ -3410,6 +3525,106 @@ mod tests { } } + /// A one-partition source that yields one batch and then remains pending. + /// It lets tests distinguish true early emission from output produced only + /// after end-of-input. + #[derive(Debug)] + struct OneBatchThenPendingExec { + batch: RecordBatch, + cache: Arc, + } + + impl OneBatchThenPendingExec { + fn new( + batch: RecordBatch, + group_contiguous_exprs: Vec>, + ) -> Self { + let cache = PlanProperties::new( + EquivalenceProperties::new(batch.schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + ) + .with_group_contiguous_exprs(group_contiguous_exprs); + Self { + batch, + cache: Arc::new(cache), + } + } + } + + impl DisplayAs for OneBatchThenPendingExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "OneBatchThenPendingExec") + } + } + + impl ExecutionPlan for OneBatchThenPendingExec { + fn name(&self) -> &'static str { + "OneBatchThenPendingExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn replace_children( + self: Arc, + _children: Vec>, + _options: ReplaceChildrenOptions, + ) -> Result> { + internal_err!("Children cannot be replaced in {self:?}") + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let stream = futures::stream::iter([Ok(self.batch.clone())]) + .chain(futures::stream::pending()); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.batch.schema(), + stream, + ))) + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))) + } + } + fn new_finite_memory_migrated_hash_ctx( batch_size: usize, max_memory: usize, @@ -4512,6 +4727,405 @@ mod tests { Ok(()) } + #[test] + fn partition_disjoint_composite_key_implication() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let a = col("a", &schema)?; + let b = col("b", &schema)?; + + // A contiguous `a` range lets an aggregate grouped by `(a, b)` emit + // every completed `a` range, even if `b` is unordered inside it. + assert_eq!( + group_completion_from_partition_disjointness( + &eq_properties, + &[Arc::clone(&a), Arc::clone(&b)], + &[Arc::clone(&a)], + ), + Some(GroupCompletionMode::Partial(vec![0])) + ); + + // Disjointness of `(a, b)` says nothing about whether `a` alone can + // recur in a later logical partition. + assert_eq!( + group_completion_from_partition_disjointness( + &eq_properties, + &[Arc::clone(&a)], + &[Arc::clone(&a), Arc::clone(&b)], + ), + None + ); + + // Group-by key order is immaterial for an exact tuple match. + assert_eq!( + group_completion_from_partition_disjointness( + &eq_properties, + &[Arc::clone(&b), Arc::clone(&a)], + &[a, b], + ), + Some(GroupCompletionMode::Full) + ); + + // The property is intentionally not a generic row-order property: + // only ProjectionExec opts into propagation. + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![col("a", &schema)?])?; + let source: Arc = Arc::new(source); + let filter = FilterExecBuilder::new(lit(true), source).build()?; + assert!(filter.group_contiguous_exprs().is_empty()); + + Ok(()) + } + + #[tokio::test] + async fn partition_disjoint_aggregate_emits_before_input_ends() -> 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![2, 2, 3, 3])), + Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + ], + )?; + let key = col("key", &schema)?; + let group_by = + PhysicalGroupBy::new_single(vec![(Arc::clone(&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 input: Arc = Arc::new(OneBatchThenPendingExec::new( + batch.clone(), + vec![Arc::clone(&key)], + )); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + input, + Arc::clone(&schema), + )?; + assert_eq!(aggregate.input_order_mode, InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + assert_eq!(aggregate.benefits_from_input_partitioning(), vec![false]); + assert!(aggregate.group_contiguous_exprs().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::OrderedSingleAggregate(_))); + let mut stream: SendableRecordBatchStream = stream.into(); + let first = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("partition-disjoint aggregate should emit before input ends") + .transpose()? + .expect("partition-disjoint aggregate should emit one completed group"); + assert_snapshot!(batches_to_sort_string(&[first]), @r" ++-----+------------+ +| key | SUM(value) | ++-----+------------+ +| 2 | 30 | ++-----+------------+ +"); + + // The legacy aggregate implementation must honor the same contract + // while the migration flag remains configurable. + let input: Arc = + Arc::new(OneBatchThenPendingExec::new(batch.clone(), vec![key])); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + input, + Arc::clone(&schema), + )?; + let fallback_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", false), + ), + ); + let stream = aggregate.execute_typed(0, &fallback_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + let mut stream: SendableRecordBatchStream = stream.into(); + let first = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("fallback aggregate should emit before input ends") + .transpose()? + .expect("fallback aggregate should emit one completed group"); + assert_snapshot!(batches_to_sort_string(&[first]), @r" ++-----+------------+ +| key | SUM(value) | ++-----+------------+ +| 2 | 30 | ++-----+------------+ +"); + + // The same never-ending input without the property remains blocking. + let input: Arc = + Arc::new(OneBatchThenPendingExec::new(batch, vec![])); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![aggr_expr], + vec![None], + input, + schema, + )?; + assert_eq!(aggregate.cache().emission_type, EmissionType::Final); + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let mut stream: SendableRecordBatchStream = stream.into(); + let mut next = stream.next().boxed(); + assert_is_pending(&mut next); + + Ok(()) + } + + #[tokio::test] + async fn partition_disjoint_aggregate_handles_order_resets() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 2, 3, 3])), + Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + ], + )?, + // A new logical source partition can reset key order, but its key + // values must not overlap any prior logical partition. + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![0, 0, 1, 1])), + Arc::new(Int64Array::from(vec![1, 2, 3, 4])), + ], + )?, + ]; + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![col("key", &schema)?])?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let aggr_expr = vec![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(2); + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggr_expr.clone(), + vec![None], + Arc::new(input.clone()), + Arc::clone(&schema), + )?; + assert_eq!(partial.input_order_mode, InputOrderMode::Linear); + assert_eq!(partial.group_completion_mode, GroupCompletionMode::Full); + assert_eq!(partial.cache().emission_type, EmissionType::Incremental); + assert!(matches!( + partial.execute_typed(0, &task_ctx)?, + StreamType::OrderedPartialAggregate(_) + )); + + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + aggr_expr, + vec![None], + Arc::new(input), + schema, + )?; + + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedSingleAggregate(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+------------+ +| key | SUM(value) | ++-----+------------+ +| 0 | 3 | +| 1 | 7 | +| 2 | 30 | +| 3 | 70 | ++-----+------------+ +"); + + Ok(()) + } + + #[tokio::test] + async fn partition_disjoint_subset_uses_partial_group_ordering() -> 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 batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + // `a` is contiguous, but deliberately resets from 3 to 0. + Arc::new(Int32Array::from(vec![2, 2, 2, 2, 3, 3, 0, 0, 0, 0])), + // `b` is not ordered within an `a` range and can recur there. + Arc::new(Int32Array::from(vec![1, 0, 1, 0, 0, 1, 1, 0, 1, 0])), + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])), + ], + )?; + let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![col("a", &schema)?])?; + let group_by = PhysicalGroupBy::new_single(vec![ + (col("a", &schema)?, "a".to_string()), + (col("b", &schema)?, "b".to_string()), + ]); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + )], + vec![None], + Arc::new(input), + schema, + )?; + assert_eq!( + aggregate.group_completion_mode, + GroupCompletionMode::Partial(vec![0]) + ); + + let stream = aggregate.execute_typed(0, &new_migrated_hash_ctx(2))?; + assert!(matches!(stream, StreamType::OrderedSingleAggregate(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++---+---+------------+ +| a | b | SUM(value) | ++---+---+------------+ +| 0 | 0 | 18 | +| 0 | 1 | 16 | +| 2 | 0 | 6 | +| 2 | 1 | 4 | +| 3 | 0 | 5 | +| 3 | 1 | 6 | ++---+---+------------+ +"); + + Ok(()) + } + + #[tokio::test] + async fn partition_disjoint_derived_key_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), + Field::new("value", DataType::Int64, false), + ])); + let batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1, 1])), + Arc::new(TimestampSecondArray::from(vec![20, 21, 30, 31])), + Arc::new(Int64Array::from(vec![1, 2, 3, 4])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1, 1])), + Arc::new(TimestampSecondArray::from(vec![0, 1, 10, 11])), + Arc::new(Int64Array::from(vec![5, 6, 7, 8])), + ], + )?, + ]; + // The source, rather than ProjectionExec, certifies that logical source + // partitions do not split or overlap the ten-second `date_bin` values. + let time_bin = Arc::new(ScalarFunctionExpr::try_new( + date_bin(), + vec![ + lit(ScalarValue::new_interval_dt(0, 10_000)), + col("time", &schema)?, + ], + &schema, + Arc::new(ConfigOptions::default()), + )?) as Arc; + let source = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![ + col("key", &schema)?, + Arc::clone(&time_bin), + ])?; + let projection = ProjectionExec::try_new( + [ + ProjectionExpr::new(col("key", &schema)?, "key"), + ProjectionExpr::new(time_bin, "time_bin"), + ProjectionExpr::new(col("value", &schema)?, "value"), + ], + Arc::new(source), + )?; + assert_eq!(projection.group_contiguous_exprs().len(), 2); + + let projected_schema = projection.schema(); + let group_by = PhysicalGroupBy::new_single(vec![ + (col("key", &projected_schema)?, "key".to_string()), + (col("time_bin", &projected_schema)?, "time_bin".to_string()), + ]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &projected_schema)?]) + .schema(Arc::clone(&projected_schema)) + .alias("SUM(value)") + .build()?, + )]; + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + aggr_expr, + vec![None], + Arc::new(projection), + projected_schema, + )?; + assert_eq!(aggregate.input_order_mode, InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full); + + let stream = aggregate.execute_typed(0, &new_migrated_hash_ctx(2))?; + assert!(matches!(stream, StreamType::OrderedSingleAggregate(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+---------------------+------------+ +| key | time_bin | SUM(value) | ++-----+---------------------+------------+ +| 1 | 1970-01-01T00:00:00 | 11 | +| 1 | 1970-01-01T00:00:10 | 15 | +| 1 | 1970-01-01T00:00:20 | 3 | +| 1 | 1970-01-01T00:00:30 | 7 | ++-----+---------------------+------------+ +"); + + Ok(()) + } + /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. #[tokio::test] async fn ordered_partial_aggregate_planning() -> Result<()> { @@ -7059,6 +7673,71 @@ mod tests { Ok(()) } + #[tokio::test] + async fn partition_disjoint_partial_completion_spills() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("completion_key", DataType::Int64, false), + Field::new("group_key", DataType::Int64, false), + Field::new("value", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + // The completion key is contiguous but deliberately unsorted. + Arc::new(Int64Array::from(vec![2, 2, 0, 0, 1, 1])), + Arc::new(Int64Array::from(vec![0, 1, 0, 1, 0, 1])), + Arc::new(Int64Array::from(vec![1; 6])), + ], + )?; + let input = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_group_contiguous_keys(vec![col("completion_key", &schema)?])?; + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + ( + col("completion_key", &schema)?, + "completion_key".to_string(), + ), + (col("group_key", &schema)?, "group_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::new(input), + Arc::clone(&schema), + )?); + + assert_eq!(aggregate.input_order_mode, InputOrderMode::Linear); + assert_eq!( + aggregate.group_completion_mode, + GroupCompletionMode::Partial(vec![0]) + ); + assert!(aggregate.cache().output_ordering().is_none()); + + let task_ctx = new_migrated_spill_ctx(1, 600); + let output = collect(aggregate.execute(0, task_ctx)?).await?; + assert_spill_count_metric(true, Arc::clone(&aggregate)); + assert_snapshot!(batches_to_sort_string(&output), @r" ++----------------+-----------+------------+ +| completion_key | group_key | SUM(value) | ++----------------+-----------+------------+ +| 0 | 0 | 1 | +| 0 | 1 | 1 | +| 1 | 0 | 1 | +| 1 | 1 | 1 | +| 2 | 0 | 1 | +| 2 | 1 | 1 | ++----------------+-----------+------------+ +"); + + Ok(()) + } + /// Tests that when the memory pool is too small to accommodate the sort /// reservation during spill, the error is properly propagated as /// ResourcesExhausted rather than silently exceeding memory limits. diff --git a/datafusion/physical-plan/src/aggregates/order/full.rs b/datafusion/physical-plan/src/aggregates/order/full.rs index ca818d6a2d598..7a531b7db275d 100644 --- a/datafusion/physical-plan/src/aggregates/order/full.rs +++ b/datafusion/physical-plan/src/aggregates/order/full.rs @@ -18,11 +18,11 @@ use datafusion_expr::EmitTo; use std::mem::size_of; -/// Tracks grouping state when the data is ordered entirely by its -/// group keys +/// Tracks grouping state when each complete group key forms one contiguous +/// range of input rows. /// -/// When the group values are sorted, as soon as we see group `n+1` we -/// know we will never see any rows for group `n` again and thus they +/// A full sort order is sufficient but not necessary. As soon as a new group is +/// seen, the contiguity guarantee proves that all prior groups are complete and /// can be emitted. /// /// For example, given `SUM(amt) GROUP BY id` if the input is sorted diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 259411b00b697..40c09efa7c9b6 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -28,14 +28,39 @@ use crate::InputOrderMode; pub use full::GroupOrderingFull; pub use partial::GroupOrderingPartial; -/// Ordering information for each group in the hash table +/// Describes how an aggregate can detect completed groups. +/// +/// Unlike [`InputOrderMode`], this mode does not imply that key values are +/// sorted. A source-provided contiguity guarantee can establish `Partial` or +/// `Full` completion even when successive key ranges have no global order. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum GroupCompletionMode { + /// No group can be proven complete before end-of-input. + None, + /// A subset of the group keys identifies contiguous ranges. + Partial(Vec), + /// Every complete group key identifies one contiguous range. + Full, +} + +impl From<&InputOrderMode> for GroupCompletionMode { + fn from(value: &InputOrderMode) -> Self { + match value { + InputOrderMode::Linear => Self::None, + InputOrderMode::PartiallySorted(indices) => Self::Partial(indices.clone()), + InputOrderMode::Sorted => Self::Full, + } + } +} + +/// Tracks which groups in the hash table are complete. #[derive(Debug)] pub enum GroupOrdering { - /// Groups are not ordered + /// No group-completion guarantee is available. None, - /// Groups are ordered by some pre-set of the group keys + /// A subset of the group keys forms contiguous ranges. Partial(GroupOrderingPartial), - /// Groups are entirely contiguous, + /// Complete group keys form contiguous ranges. Full(GroupOrderingFull), } @@ -43,17 +68,27 @@ impl GroupOrdering { /// Create a `GroupOrdering` for the specified ordering pub fn try_new(mode: &InputOrderMode) -> Result { match mode { - InputOrderMode::Linear => Ok(GroupOrdering::None), + InputOrderMode::Linear => Ok(Self::None), InputOrderMode::PartiallySorted(order_indices) => { - GroupOrderingPartial::try_new(order_indices.clone()) - .map(GroupOrdering::Partial) + GroupOrderingPartial::try_new(order_indices.clone()).map(Self::Partial) + } + InputOrderMode::Sorted => Ok(Self::Full(GroupOrderingFull::new())), + } + } + + /// Create a `GroupOrdering` for the specified group-completion mode. + pub(crate) fn try_new_for_mode(mode: &GroupCompletionMode) -> Result { + match mode { + GroupCompletionMode::None => Ok(Self::None), + GroupCompletionMode::Partial(order_indices) => { + GroupOrderingPartial::try_new(order_indices.clone()).map(Self::Partial) } - InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())), + GroupCompletionMode::Full => Ok(Self::Full(GroupOrderingFull::new())), } } - /// Returns how many groups can be emitted while respecting the current - /// ordering guarantees, or `None` if no data can be emitted. + /// Returns how many groups can be emitted under the current completion + /// guarantee, or `None` if no data can be emitted. pub fn emit_to(&self) -> Option { match self { GroupOrdering::None => None, @@ -65,7 +100,7 @@ impl GroupOrdering { /// Returns the emit strategy to use under memory pressure (OOM). /// /// Returns the strategy that must be used when emitting up to `n` groups - /// while respecting the current ordering guarantees. + /// while respecting the current completion guarantee. /// /// Returns `None` if no data can be emitted. pub fn oom_emit_to(&self, n: usize) -> Option { @@ -93,7 +128,7 @@ impl GroupOrdering { } } - /// Resets the ordering state while preserving the configured ordering mode. + /// Resets the completion state while preserving the configured mode. /// /// Ordered partial aggregation uses this after passing intermediate states /// downstream, and ordered final aggregation uses it after spilling a run. @@ -147,7 +182,7 @@ impl GroupOrdering { Ok(()) } - /// Returns the size of memory used by the ordering state, in bytes. + /// Returns the size of memory used by the completion state, in bytes. pub fn size(&self) -> usize { size_of::() + match self { diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 1603bb6d079be..4427037b82ea5 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -27,12 +27,13 @@ use datafusion_common::{Result, ScalarValue}; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::EmitTo; -/// Tracks grouping state when the data is ordered by some subset of -/// the group keys. +/// Tracks grouping state when a subset of the group keys forms contiguous +/// ranges of input rows. /// -/// Once the next *sort key* value is seen, never see groups with that -/// sort key again, so we can emit all groups with the previous sort -/// key and earlier. +/// Once the next completion-key value is seen, the prior value cannot recur, so +/// all groups from that prior range can be emitted. A sort order on the subset +/// is one way to establish this guarantee, but the values may also reset between +/// disjoint logical source runs. /// /// For example, given `SUM(amt) GROUP BY id, state` if the input is /// sorted by `state`, when a new value of `state` is seen, all groups diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 2c26b74da7748..8852354f6bdb0 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Final aggregate stream for ordered partial-state input. +//! Final aggregate stream for partial-state input with contiguous group ranges. use std::ops::ControlFlow; use std::sync::Arc; @@ -35,16 +35,16 @@ use super::AggregateExec; use super::aggregate_hash_table::{ FinalMarker, OrderedAggregateTable, OrderedAggregateTableMetrics, }; +use super::order::GroupCompletionMode; use crate::aggregates::AggregateMode; use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::SpillManager; use crate::stream::EmptyRecordBatchStream; -use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; +use crate::{RecordBatchStream, SendableRecordBatchStream}; -/// Final aggregate stream for `InputOrderMode::Sorted` and -/// `InputOrderMode::PartiallySorted`. +/// Final aggregate stream when completed group ranges can be identified. /// /// See comments at [`super::ordered_partial_stream::OrderedPartialAggregateStream`] for details. /// @@ -52,7 +52,7 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// /// This section is only for implementation notes, for background, see [`super::ordered_partial_stream::OrderedPartialAggregateStream`] /// -/// For partially sorted input, spilling works as follows: +/// For a `Partial` group-completion mode, spilling works as follows: /// /// - Reserve the table footprint plus one `u32` sort index per buffered group. The /// extra index array is used in later sorting before spilling. @@ -132,14 +132,16 @@ impl OrderedFinalSpillContext { context: &Arc, partition: usize, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, spill_schema: &SchemaRef, spill_metrics: SpillMetrics, ) -> Result { let group_schema = agg.group_by.group_schema(spill_schema)?; let output_ordering = agg.cache.output_ordering(); - let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { - return internal_err!("Ordered final spill requires partially ordered input"); + let GroupCompletionMode::Partial(order_indices) = group_completion_mode else { + return internal_err!( + "Ordered final spill requires partial group completion" + ); }; let spill_indices = order_indices.iter().copied().chain( (0..group_schema.fields().len()).filter(|idx| !order_indices.contains(idx)), @@ -249,7 +251,7 @@ impl OrderedFinalSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, @@ -269,10 +271,10 @@ impl OrderedFinalAggregateStream { agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned )); - debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(agg.group_completion_mode, GroupCompletionMode::None); let input = agg.input.execute(partition, Arc::clone(context))?; - Self::new_with_input(agg, context, partition, input, &agg.input_order_mode) + Self::new_with_input(agg, context, partition, input, &agg.group_completion_mode) } pub(in crate::aggregates) fn new_with_input( @@ -280,7 +282,7 @@ impl OrderedFinalAggregateStream { context: &Arc, partition: usize, input: SendableRecordBatchStream, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, ) -> Result { let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let metrics = OrderedAggregateTableMetrics::new(agg, partition); @@ -299,7 +301,7 @@ impl OrderedFinalAggregateStream { context, partition, input, - input_order_mode, + group_completion_mode, baseline_metrics, metrics, Some(spill_metrics), @@ -319,7 +321,7 @@ impl OrderedFinalAggregateStream { context: &Arc, partition: usize, input: SendableRecordBatchStream, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, baseline_metrics: BaselineMetrics, metrics: OrderedAggregateTableMetrics, spill_metrics: Option, @@ -329,13 +331,13 @@ impl OrderedFinalAggregateStream { agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned )); - debug_assert_ne!(*input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(*group_completion_mode, GroupCompletionMode::None); let schema = Arc::clone(&agg.schema); let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); - let can_spill = matches!(input_order_mode, InputOrderMode::PartiallySorted(_)) + let can_spill = matches!(group_completion_mode, GroupCompletionMode::Partial(_)) && context.runtime_env().disk_manager.tmp_files_enabled(); let spill_context = if can_spill { let Some(spill_metrics) = spill_metrics else { @@ -346,7 +348,7 @@ impl OrderedFinalAggregateStream { context, partition, batch_size, - input_order_mode, + group_completion_mode, &input_schema, spill_metrics, )?)) @@ -354,12 +356,12 @@ impl OrderedFinalAggregateStream { None }; - let table = OrderedAggregateTable::::new_with_input_order( + let table = OrderedAggregateTable::::new_with_group_completion( agg, &input_schema, Arc::clone(&schema), batch_size, - input_order_mode, + group_completion_mode, metrics, )?; Ok(Self { diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 9e93a111a6466..550894947a092 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Partial aggregate stream for ordered group input. +//! Partial aggregate stream for input with contiguous group ranges. use std::sync::Arc; @@ -29,13 +29,12 @@ use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; -use crate::aggregates::order::GroupOrdering; +use crate::aggregates::order::{GroupCompletionMode, GroupOrdering}; use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; -use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; +use crate::{SendableRecordBatchStream, metrics}; -/// Partial aggregate stream for `InputOrderMode::Sorted` and -/// `InputOrderMode::PartiallySorted`. +/// Partial aggregate stream when completed group ranges can be identified. /// /// # Example /// @@ -59,18 +58,18 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// Output: results for all groups (for example, `AVG(x)` calculated from the /// state) /// -/// # Order-based Optimization +/// # Group-completion optimization /// /// For the aggregation work, the hash aggregation implementation is reused. /// /// After each input batch, check whether any groups can be emitted eagerly to -/// improve memory efficiency. For example, if the last group key seen is -/// `k = 100`, it is safe to emit all groups with keys less than 100 because the -/// input is ordered. +/// improve memory efficiency. Once a new contiguous key range begins, the prior +/// range is safe to emit. The guarantee may come from sort order or from +/// partition-disjoint source keys. /// /// # Memory Pressure and Spilling /// -/// ## Fully ordered case +/// ## Fully contiguous case /// /// If the input is ordered by every group key, for example: /// @@ -84,7 +83,7 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// If a memory reservation nevertheless fails, the stream returns the error /// directly, indicating an unexpected behavior. /// -/// ## Partially ordered case +/// ## Partial completion-key case /// /// If the input is ordered by only a subset of the group keys, for example: /// @@ -126,7 +125,7 @@ impl OrderedPartialAggregateStream { partition: usize, ) -> Result { debug_assert_eq!(agg.mode, AggregateMode::Partial); - debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(agg.group_completion_mode, GroupCompletionMode::None); let schema = Arc::clone(&agg.schema); let input = agg.input.execute(partition, Arc::clone(context))?; @@ -277,11 +276,11 @@ impl OrderedPartialAggregateStream { /// Update the memory reservation, and: /// - If memory reservation succeed, returns `Ok(None)` /// - If memory reservation failed, - /// - If input is partially ordered, materialize all the output, and + /// - With partial group completion, materialize all the output, and /// directly send them to the final aggregation stage. /// Returns `Ok(Some(batch))` - /// - If input is fully ordered, directly return error. It's not - /// expected to use more than constant memory. + /// - With full group completion, directly return an error. The number + /// of incomplete groups is expected to remain bounded. /// Returns `Err(..)` /// /// # Implementation Note diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index da00b42e5c3ed..9b898ca4836a5 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Single-stage aggregate stream for ordered raw input. +//! Single-stage aggregate stream for raw input with contiguous group ranges. use std::ops::ControlFlow; use std::sync::Arc; @@ -34,6 +34,7 @@ use futures::stream::{Stream, StreamExt}; use super::aggregate_hash_table::{ OrderedAggregateTable, OrderedAggregateTableMetrics, SingleMarker, }; +use super::order::GroupCompletionMode; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; @@ -44,8 +45,7 @@ use crate::spill::spill_manager::SpillManager; use crate::stream::EmptyRecordBatchStream; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; -/// Single aggregate stream for `InputOrderMode::Sorted` and -/// `InputOrderMode::PartiallySorted`. +/// Single aggregate stream when completed group ranges can be identified. /// /// # Example /// @@ -62,18 +62,18 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// Input: raw rows /// Output: final results for all groups (for example, `AVG(x)`) /// -/// # Order-based Optimization +/// # Group-completion optimization /// /// For the aggregation work, the hash aggregation implementation is reused. /// /// After each input batch, check whether any groups can be emitted eagerly to -/// improve memory efficiency. For example, if the last group key seen is -/// `k = 100`, it is safe to emit all groups with keys less than 100 because the -/// input is ordered. +/// improve memory efficiency. Once a new contiguous key range begins, the prior +/// range is safe to emit. The guarantee may come from sort order or from +/// partition-disjoint source keys. /// /// # Memory Pressure and Spilling /// -/// ## Fully ordered case +/// ## Fully contiguous case /// /// If the input is ordered by every group key, for example: /// @@ -87,7 +87,7 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// If a memory reservation nevertheless fails, the stream returns the error /// directly, indicating an unexpected behavior. /// -/// ## Partially ordered case +/// ## Partial completion-key case /// /// If the input is ordered by only a subset of the group keys, for example: /// @@ -175,15 +175,15 @@ impl OrderedSingleSpillContext { context: &Arc, partition: usize, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, spill_schema: &SchemaRef, spill_metrics: SpillMetrics, ) -> Result { let group_schema = agg.group_by.group_schema(&agg.input().schema())?; let output_ordering = agg.cache.output_ordering(); - let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { + let GroupCompletionMode::Partial(order_indices) = group_completion_mode else { return internal_err!( - "Ordered single spill requires partially ordered input" + "Ordered single spill requires partial group completion" ); }; let spill_indices = order_indices.iter().copied().chain( @@ -222,6 +222,7 @@ impl OrderedSingleSpillContext { }; final_agg.group_by = Arc::new(agg.group_by.as_final()); final_agg.input_order_mode = InputOrderMode::Sorted; + final_agg.group_completion_mode = GroupCompletionMode::Full; Ok(Self { final_agg, @@ -309,7 +310,7 @@ impl OrderedSingleSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, @@ -329,7 +330,7 @@ impl OrderedSingleAggregateStream { agg.mode, AggregateMode::Single | AggregateMode::SinglePartitioned )); - debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(agg.group_completion_mode, GroupCompletionMode::None); let schema = Arc::clone(&agg.schema); let input = agg.input.execute(partition, Arc::clone(context))?; @@ -353,7 +354,7 @@ impl OrderedSingleAggregateStream { )?; let can_spill = - matches!(agg.input_order_mode, InputOrderMode::PartiallySorted(_)) + matches!(agg.group_completion_mode, GroupCompletionMode::Partial(_)) && context.runtime_env().disk_manager.tmp_files_enabled(); let spill_context = if can_spill { Some(Box::new(OrderedSingleSpillContext::new( @@ -361,7 +362,7 @@ impl OrderedSingleAggregateStream { context, partition, batch_size, - &agg.input_order_mode, + &agg.group_completion_mode, &state_schema, spill_metrics, )?)) diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 3e306d72a7e82..6e3667bbe273a 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -39,6 +39,7 @@ use futures::stream::{Stream, StreamExt}; use super::aggregate_hash_table::{ AggregateHashTable, OrderedAggregateTableMetrics, SingleMarker, }; +use super::order::GroupCompletionMode; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; @@ -212,6 +213,7 @@ impl SingleSpillContext { }; final_agg.group_by = Arc::new(agg.group_by.as_final()); final_agg.input_order_mode = InputOrderMode::Sorted; + final_agg.group_completion_mode = GroupCompletionMode::Full; Ok(Self { final_agg, @@ -300,7 +302,7 @@ impl SingleSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index a4d081b3d9e75..c43c931a8dfff 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -159,6 +159,50 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// trait, which is implemented for all `ExecutionPlan`s. fn properties(&self) -> &Arc; + /// Expressions whose tuple values are contiguous within each output + /// partition, even when they are not globally sorted. + /// + /// The returned expressions are the components of one composite key, not + /// independent keys or alternative guarantees. For every output partition + /// and every key tuple (including tuples containing null), all rows with + /// that tuple must occur in at most one contiguous range. For example, + /// `(a, 2), (a, 2), (b, 1), (a, 0)` satisfies this property, while + /// `(a, 2), (b, 1), (a, 2)` does not. + /// + /// This property is intended for data sources that concatenate logical + /// source partitions. A source can establish it when: + /// + /// 1. rows within each logical partition are grouped by these expressions + /// (for example, lexicographically sorted), and + /// 2. no tuple value occurs in more than one logical partition. + /// + /// For example, a source that concatenates files sorted by `(key, time)` + /// may declare `(key, date_bin(time))` only after verifying that the file + /// boundaries do not split a `date_bin` value. Declaring `(key, time)` does + /// not automatically establish the property for `(key, date_bin(time))`. + /// + /// This is not an output-distribution guarantee: the same tuple may occur + /// in different output partitions. It cannot by itself satisfy a + /// [`Distribution`] requirement. + /// + /// [`ProjectionExec`] propagates the property only when every expression can + /// be projected. Row-transparent execution wrappers such as + /// [`crate::coop::CooperativeExec`] also preserve it. Other operators, + /// including aggregates, use the default empty value, so the guarantee is + /// limited to a data source, optional projections and wrappers, and the + /// first consumer. + /// + /// # Correctness + /// + /// This is a correctness contract. An invalid declaration may cause a + /// streaming aggregate to emit a group before all of its rows are seen. + /// Implementations should set this property with + /// [`PlanProperties::with_group_contiguous_exprs`] rather than override this + /// method, so child-property caching remains correct. + fn group_contiguous_exprs(&self) -> &[Arc] { + self.properties().group_contiguous_exprs() + } + /// Returns an error if this individual node does not conform to its invariants. /// These invariants are typically only checked in debug mode. /// @@ -1496,6 +1540,8 @@ pub struct PlanProperties { pub scheduling_type: SchedulingType, /// See [ExecutionPlanProperties::output_ordering] output_ordering: Option, + /// See [`ExecutionPlan::group_contiguous_exprs`] + group_contiguous_exprs: Vec>, } impl PlanProperties { @@ -1516,6 +1562,7 @@ impl PlanProperties { evaluation_type: EvaluationType::Lazy, scheduling_type: SchedulingType::NonCooperative, output_ordering, + group_contiguous_exprs: vec![], } } @@ -1567,6 +1614,18 @@ impl PlanProperties { self } + /// Overwrite the group-contiguous composite key. + /// + /// See [`ExecutionPlan::group_contiguous_exprs`] for the correctness + /// contract. + pub fn with_group_contiguous_exprs( + mut self, + group_contiguous_exprs: Vec>, + ) -> Self { + self.group_contiguous_exprs = group_contiguous_exprs; + self + } + /// Set constraints having mut reference. pub fn set_constraints(&mut self, constraints: Constraints) { self.eq_properties.set_constraints(constraints); @@ -1590,6 +1649,11 @@ impl PlanProperties { self.output_ordering.as_ref() } + /// Components of the group-contiguous composite key, if any. + pub fn group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + /// Get schema of the node. pub(crate) fn schema(&self) -> &SchemaRef { self.eq_properties.schema() @@ -1981,9 +2045,9 @@ pub fn reset_plan_states(plan: Arc) -> Result>>() + .unwrap_or_default(); let cache = Self::compute_properties( &input, &projection_mapping, Arc::clone(projector.output_schema()), reuse_from, - )?; + )? + .with_group_contiguous_exprs(group_contiguous_exprs); Ok(Self { projector, input, @@ -1498,6 +1510,7 @@ mod tests { use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; + use crate::test::TestMemoryExec; use crate::test::exec::StatisticsExec; use arrow::datatypes::{DataType, Field, Schema}; @@ -1509,6 +1522,65 @@ mod tests { BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit, }; + #[test] + fn group_contiguous_projection_is_all_or_nothing() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time", DataType::Int32, false), + ])); + let time_bin = binary( + col("time", &schema)?, + Operator::Divide, + lit(ScalarValue::Int32(Some(10))), + &schema, + )?; + let plain_source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?; + let source = plain_source.clone().try_with_group_contiguous_keys(vec![ + col("key", &schema)?, + Arc::clone(&time_bin), + ])?; + + let projection = ProjectionExec::try_new( + [ + ProjectionExpr::new(col("key", &schema)?, "key"), + ProjectionExpr::new(time_bin, "time_bin"), + ], + Arc::new(source.clone()), + )?; + let projected_schema = projection.schema(); + let expected = [ + col("key", &projected_schema)?, + col("time_bin", &projected_schema)?, + ]; + assert_eq!(projection.group_contiguous_exprs().len(), expected.len()); + assert!( + projection + .group_contiguous_exprs() + .iter() + .zip(expected) + .all(|(actual, expected)| actual.eq(&expected)) + ); + + // Projecting only one component must not turn contiguity of the + // `(key, time_bin)` tuple into the stronger (and generally false) + // claim that `key` alone is contiguous. + let partial_projection = ProjectionExec::try_new( + [ProjectionExpr::new(col("key", &schema)?, "key")], + Arc::new(source.clone()), + )?; + assert!(partial_projection.group_contiguous_exprs().is_empty()); + + // Child replacement must notice this metadata and recompute the + // projection's cached properties. + assert!(!Arc::ptr_eq(source.properties(), plain_source.properties())); + let projection: Arc = Arc::new(projection); + let replaced = + replace_children_if_necessary(projection, vec![Arc::new(plain_source)])?; + assert!(replaced.group_contiguous_exprs().is_empty()); + + Ok(()) + } + #[test] fn test_try_new_with_schema_metadata_only_replaces_metadata() -> Result<()> { let input_schema = Arc::new(Schema::new(vec![Field::new( diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index b38a46d160755..5fca12ba451df 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -73,6 +73,8 @@ pub struct TestMemoryExec { projection: Option>, /// Sort information: one or more equivalent orderings sort_information: Vec, + /// Composite key whose values are contiguous within each output stream. + group_contiguous_exprs: Vec>, /// if partition sizes should be displayed show_sizes: bool, /// The maximum number of records to read from this plan. If `None`, @@ -106,16 +108,27 @@ impl DisplayAs for TestMemoryExec { let limit = self .fetch .map_or(String::new(), |limit| format!(", fetch={limit}")); + let group_contiguous_exprs = self.group_contiguous_exprs(); + let group_contiguous = if group_contiguous_exprs.is_empty() { + String::new() + } else { + let exprs = group_contiguous_exprs + .iter() + .map(ToString::to_string) + .collect::>() + .join(", "); + format!(", group_contiguous=[{exprs}]") + }; if self.show_sizes { write!( f, - "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", + "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}{group_contiguous}", partition_sizes.len(), ) } else { write!( f, - "partitions={}{limit}{output_ordering}{constraints}", + "partitions={}{limit}{output_ordering}{constraints}{group_contiguous}", partition_sizes.len(), ) } @@ -226,6 +239,7 @@ impl TestMemoryExec { EmissionType::Incremental, Boundedness::Bounded, ) + .with_group_contiguous_exprs(self.group_contiguous_exprs.clone()) } fn output_partitioning(&self) -> Partitioning { @@ -268,6 +282,7 @@ impl TestMemoryExec { projected_schema, projection, sort_information: vec![], + group_contiguous_exprs: vec![], show_sizes: true, fetch: None, }) @@ -363,6 +378,50 @@ impl TestMemoryExec { Ok(self) } + /// Attach a composite key whose values occur in one contiguous range in + /// each output stream. See [`ExecutionPlan::group_contiguous_exprs`] for + /// the correctness contract. + pub fn try_with_group_contiguous_keys( + mut self, + mut group_contiguous_exprs: Vec>, + ) -> Result { + // All expressions must refer to the original schema. + let fields = self.schema.fields(); + let ambiguous_column = group_contiguous_exprs + .iter() + .flat_map(collect_columns) + .find(|col| { + fields + .get(col.index()) + .map(|field| field.name() != col.name()) + .unwrap_or(true) + }); + assert_or_internal_err!( + ambiguous_column.is_none(), + "Column {:?} is not found in the original schema of the TestMemoryExec", + ambiguous_column.as_ref().unwrap() + ); + + if let Some(projection) = &self.projection { + let base_schema = self.original_schema(); + let proj_exprs = projection.iter().map(|idx| { + let name = base_schema.field(*idx).name(); + (Arc::new(Column::new(name, *idx)) as _, name.to_string()) + }); + let projection_mapping = + ProjectionMapping::try_new(proj_exprs, &base_schema)?; + let base_eqp = EquivalenceProperties::new(base_schema); + group_contiguous_exprs = base_eqp + .project_expressions(group_contiguous_exprs.iter(), &projection_mapping) + .collect::>>() + .unwrap_or_default(); + } + + self.group_contiguous_exprs = group_contiguous_exprs; + self.cache = Arc::new(self.compute_properties()); + Ok(self) + } + /// Arc clone of ref to original schema pub fn original_schema(&self) -> SchemaRef { Arc::clone(&self.schema)