feat: stream aggregates for group-contiguous input

Author: xavlee

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::<AggregateExec>()
+        .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();

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<dyn ExecutionPlan>;
+
+    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())];

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<dyn PhysicalExpr>] {
+        &[]
+    }
+
     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())
     }
 

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<Arc<dyn PhysicalExpr>>,
+        ) -> 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<dyn ExecutionPlan> = (&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![

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<FFI_PhysicalExpr>,
 }
 
 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<FFI_PhysicalExpr> {
+    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<FFI_PlanProperties> 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(<Arc<dyn datafusion_physical_expr::PhysicalExpr>>::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<PlanProperties> {

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!(

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)
 }
 

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<dyn ExecutionPlan> = (&plan).try_into()?;
+        assert!(plan.is::<ForeignExecutionPlan>());
+
+        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> {

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)