feat: add narrow group-contiguous source property

Author: xavlee

datafusion/datasource/src/source.rs

@@ -165,6 +165,19 @@ pub trait DataSource: Any + Send + Sync + Debug {
 
     fn output_partitioning(&self) -> Partitioning;
     fn eq_properties(&self) -> EquivalenceProperties;
+
+    /// Expressions whose complete tuple is contiguous within each output
+    /// partition.
+    ///
+    /// See [`ExecutionPlan::group_contiguous_exprs`] for the full correctness
+    /// contract. Expressions must refer to the schema returned by
+    /// [`Self::eq_properties`]. A source rewrite such as [`Self::repartitioned`]
+    /// or [`Self::try_swapping_with_projection`] must preserve, remap, or clear
+    /// this assertion as appropriate for the rewritten output streams.
+    fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+        &[]
+    }
+
     fn scheduling_type(&self) -> SchedulingType {
         SchedulingType::NonCooperative
     }
@@ -397,6 +410,10 @@ impl ExecutionPlan for DataSourceExec {
         &self.cache
     }
 
+    fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+        self.data_source.group_contiguous_exprs()
+    }
+
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         Vec::new()
     }
@@ -705,3 +722,90 @@ where
         Self::new(Arc::new(source))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+    use datafusion_physical_expr::expressions::Column;
+    use datafusion_physical_plan::EmptyRecordBatchStream;
+
+    #[derive(Debug)]
+    struct GroupContiguousSource {
+        schema: SchemaRef,
+        group_contiguous_exprs: Vec<Arc<dyn PhysicalExpr>>,
+    }
+
+    impl DataSource for GroupContiguousSource {
+        fn open(
+            &self,
+            _partition: usize,
+            _context: Arc<TaskContext>,
+        ) -> Result<SendableRecordBatchStream> {
+            Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(
+                &self.schema,
+            ))))
+        }
+
+        fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
+            write!(f, "GroupContiguousSource")
+        }
+
+        fn output_partitioning(&self) -> Partitioning {
+            Partitioning::UnknownPartitioning(1)
+        }
+
+        fn eq_properties(&self) -> EquivalenceProperties {
+            EquivalenceProperties::new(Arc::clone(&self.schema))
+        }
+
+        fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+            &self.group_contiguous_exprs
+        }
+
+        fn partition_statistics(
+            &self,
+            _partition: Option<usize>,
+        ) -> Result<Arc<Statistics>> {
+            Ok(Arc::new(Statistics::new_unknown(&self.schema)))
+        }
+
+        fn with_fetch(&self, _limit: Option<usize>) -> Option<Arc<dyn DataSource>> {
+            None
+        }
+
+        fn fetch(&self) -> Option<usize> {
+            None
+        }
+
+        fn try_swapping_with_projection(
+            &self,
+            _projection: &ProjectionExprs,
+        ) -> Result<Option<Arc<dyn DataSource>>> {
+            Ok(None)
+        }
+
+        fn apply_expressions(
+            &self,
+            _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+        ) -> Result<TreeNodeRecursion> {
+            Ok(TreeNodeRecursion::Continue)
+        }
+    }
+
+    #[test]
+    fn data_source_exec_exposes_group_contiguous_exprs() -> Result<()> {
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
+        let key = Arc::new(Column::new("key", 0)) as Arc<dyn PhysicalExpr>;
+        let exec = DataSourceExec::new(Arc::new(GroupContiguousSource {
+            schema,
+            group_contiguous_exprs: vec![Arc::clone(&key)],
+        }));
+
+        assert_eq!(exec.group_contiguous_exprs().len(), 1);
+        assert!(exec.group_contiguous_exprs()[0].eq(&key));
+        Ok(())
+    }
+}

datafusion/physical-plan/src/coop.rs

@@ -261,6 +261,10 @@ impl ExecutionPlan for CooperativeExec {
         &self.properties
     }
 
+    fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+        self.input.group_contiguous_exprs()
+    }
+
     fn maintains_input_order(&self) -> Vec<bool> {
         vec![true; self.children().len()]
     }
@@ -477,7 +481,9 @@ pub fn make_cooperative(stream: SendableRecordBatchStream) -> SendableRecordBatc
 mod tests {
     use super::*;
 
-    use arrow_schema::SchemaRef;
+    use crate::test::TestMemoryExec;
+    use arrow_schema::{DataType, Field, Schema, SchemaRef};
+    use datafusion_physical_expr::expressions::col;
 
     use futures::stream;
 
@@ -497,6 +503,19 @@ mod tests {
         Box::pin(RecordBatchStreamAdapter::new(schema, s))
     }
 
+    #[test]
+    fn cooperative_exec_preserves_group_contiguity() -> Result<()> {
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
+        let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?
+            .try_with_group_contiguous_exprs(vec![col("key", &schema)?])?;
+        let cooperative = CooperativeExec::new(Arc::new(source));
+
+        assert_eq!(cooperative.group_contiguous_exprs().len(), 1);
+        assert!(cooperative.group_contiguous_exprs()[0].eq(&col("key", &schema)?));
+        Ok(())
+    }
+
     #[tokio::test]
     async fn yield_less_than_threshold() -> Result<()> {
         let count = TASK_BUDGET - 10;

datafusion/physical-plan/src/execution_plan.rs

@@ -159,6 +159,32 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync {
     /// trait, which is implemented for all `ExecutionPlan`s.
     fn properties(&self) -> &Arc<PlanProperties>;
 
+    /// Expressions whose complete tuple is contiguous within each output
+    /// partition.
+    ///
+    /// For every distinct tuple of expression values, all rows with that tuple
+    /// must occur in at most one contiguous range in each output stream. Once a
+    /// stream produces a different tuple, the previous tuple must never occur
+    /// again. Tuple equality follows `GROUP BY` semantics, and tuple values do
+    /// not need to be sorted.
+    ///
+    /// This property does not imply any particular output ordering or
+    /// distribution. It is also deliberately fail-closed: the default is no
+    /// guarantee, and operators must explicitly preserve it. Currently,
+    /// [`ProjectionExec`] preserves the complete tuple when every expression
+    /// can be projected, and [`crate::coop::CooperativeExec`] preserves it
+    /// because it does not change rows. Other operators use the default empty
+    /// value.
+    ///
+    /// # Correctness
+    ///
+    /// This is a correctness contract. An invalid declaration can cause a
+    /// streaming aggregate to emit a group before all of its rows have been
+    /// observed, producing incorrect results.
+    fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+        &[]
+    }
+
     /// Returns an error if this individual node does not conform to its invariants.
     /// These invariants are typically only checked in debug mode.
     ///

datafusion/physical-plan/src/projection.rs

@@ -83,6 +83,8 @@ pub struct ProjectionExec {
     metrics: ExecutionPlanMetricsSet,
     /// Cache holding plan properties like equivalences, output partitioning etc.
     cache: Arc<PlanProperties>,
+    /// Complete group-contiguous tuple projected from the input, if any.
+    group_contiguous_exprs: Vec<Arc<dyn PhysicalExpr>>,
 }
 
 impl ProjectionExec {
@@ -208,6 +210,8 @@ impl ProjectionExec {
         // Construct a map from the input expressions to the output expression of the Projection
         let projection_mapping =
             projector.projection().projection_mapping(&input.schema())?;
+        let group_contiguous_exprs =
+            Self::project_group_contiguous_exprs(&input, &projection_mapping);
         let cache = Self::compute_properties(
             &input,
             &projection_mapping,
@@ -219,9 +223,26 @@ impl ProjectionExec {
             input,
             metrics: ExecutionPlanMetricsSet::new(),
             cache: Arc::new(cache),
+            group_contiguous_exprs,
         })
     }
 
+    /// Projects the complete group-contiguous tuple, dropping it when any
+    /// component cannot be mapped to the output schema.
+    fn project_group_contiguous_exprs(
+        input: &Arc<dyn ExecutionPlan>,
+        projection_mapping: &ProjectionMapping,
+    ) -> Vec<Arc<dyn PhysicalExpr>> {
+        input
+            .equivalence_properties()
+            .project_expressions(
+                input.group_contiguous_exprs().iter(),
+                projection_mapping,
+            )
+            .collect::<Option<Vec<_>>>()
+            .unwrap_or_default()
+    }
+
     /// The projection expressions stored as tuples of (expression, output column name)
     pub fn expr(&self) -> &[ProjectionExpr] {
         self.projector.projection().as_ref()
@@ -345,6 +366,10 @@ impl ExecutionPlan for ProjectionExec {
         &self.cache
     }
 
+    fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+        &self.group_contiguous_exprs
+    }
+
     fn maintains_input_order(&self) -> Vec<bool> {
         // Tell optimizer this operator doesn't reorder its input
         vec![true]
@@ -386,11 +411,21 @@ impl ExecutionPlan for ProjectionExec {
     ) -> Result<Arc<dyn ExecutionPlan>> {
         validate_child_count!(self, children);
         match options.children_properties {
-            ChildrenPropertiesMode::Keep => Ok(Arc::new(Self {
-                input: children.swap_remove(0),
-                metrics: ExecutionPlanMetricsSet::new(),
-                ..Self::clone(&*self)
-            })),
+            ChildrenPropertiesMode::Keep => {
+                let input = children.swap_remove(0);
+                let projection_mapping = self
+                    .projector
+                    .projection()
+                    .projection_mapping(&input.schema())?;
+                let group_contiguous_exprs =
+                    Self::project_group_contiguous_exprs(&input, &projection_mapping);
+                Ok(Arc::new(Self {
+                    input,
+                    metrics: ExecutionPlanMetricsSet::new(),
+                    group_contiguous_exprs,
+                    ..Self::clone(&*self)
+                }))
+            }
             ChildrenPropertiesMode::Recompute => {
                 // `Keep` above requires the child's properties to be unchanged
                 // outright. A rule that introduces a sort below this projection
@@ -631,6 +666,8 @@ impl ExecutionPlan for ProjectionExec {
             metrics: _,
             // Derived plan properties, recomputed on decode.
             cache: _,
+            // Derived from the input assertion and projection expressions.
+            group_contiguous_exprs: _,
         } = self;
         let projection_exprs = projector.projection().as_ref();
         let input = ctx.encode_child(input)?;
@@ -1515,6 +1552,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};
@@ -1526,6 +1564,70 @@ 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_exprs(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))
+        );
+
+        // A strict subset is not sufficient: contiguity of `(key, time_bin)`
+        // does not imply 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());
+
+        // Row-preserving operators do not inherit the assertion unless they
+        // opt in explicitly.
+        let filter = FilterExec::try_new(lit(true), Arc::new(source.clone()))?;
+        assert!(filter.group_contiguous_exprs().is_empty());
+
+        // The assertion deliberately lives outside PlanProperties. Exercise
+        // the child-replacement fast path with identical cached properties and
+        // verify that ProjectionExec still recomputes it.
+        assert!(Arc::ptr_eq(source.properties(), plain_source.properties()));
+        let projection: Arc<dyn ExecutionPlan> = 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(