feat: add grouped equivalence properties

Author: xavlee

datafusion/expr-common/src/sort_properties.rs

@@ -37,16 +37,24 @@ use arrow::datatypes::DataType;
 pub enum SortProperties {
     /// Use the ordinary [`SortOptions`] struct to represent ordered data:
     Ordered(SortOptions),
-    // This alternative represents unordered data:
+    /// Within each partition, all rows with the same value for this expression
+    /// form one contiguous run. The runs may occur in any order.
+    Grouped,
+    /// This alternative represents unordered data:
     #[default]
     Unordered,
-    // Singleton is used for single-valued literal numbers:
+    /// Singleton is used for single-valued literal numbers:
     Singleton,
 }
 
 impl SortProperties {
     pub fn add(&self, rhs: &Self) -> Self {
         match (self, rhs) {
+            // Addition can collapse distinct values (for example through
+            // floating-point rounding), which may join non-adjacent groups.
+            (Self::Grouped, Self::Singleton) | (Self::Singleton, Self::Grouped) => {
+                Self::Unordered
+            }
             (Self::Singleton, _) => *rhs,
             (_, Self::Singleton) => *self,
             (Self::Ordered(lhs), Self::Ordered(rhs))
@@ -65,6 +73,11 @@ impl SortProperties {
     pub fn sub(&self, rhs: &Self) -> Self {
         match (self, rhs) {
             (Self::Singleton, Self::Singleton) => Self::Singleton,
+            // Subtraction can collapse distinct values (for example through
+            // floating-point rounding), which may join non-adjacent groups.
+            (Self::Grouped, Self::Singleton) | (Self::Singleton, Self::Grouped) => {
+                Self::Unordered
+            }
             (Self::Singleton, Self::Ordered(rhs)) => Self::Ordered(SortOptions {
                 descending: !rhs.descending,
                 nulls_first: rhs.nulls_first,
@@ -89,6 +102,9 @@ impl SortProperties {
                 descending: !rhs.descending,
                 nulls_first: rhs.nulls_first,
             }),
+            // Comparisons can map several non-adjacent grouped values to the
+            // same boolean value, so they do not preserve grouping.
+            (Self::Grouped, Self::Singleton) => Self::Unordered,
             (_, Self::Singleton) => *self,
             (Self::Ordered(lhs), Self::Ordered(rhs))
                 if lhs.descending != rhs.descending
@@ -147,6 +163,7 @@ mod sort_properties_test {
     const ASC_NL: SortProperties = ordered(false, false);
     const DESC_NF: SortProperties = ordered(true, true);
     const DESC_NL: SortProperties = ordered(true, false);
+    const GROUPED: SortProperties = SortProperties::Grouped;
     const UNORDERED: SortProperties = SortProperties::Unordered;
     const SINGLETON: SortProperties = SortProperties::Singleton;
 
@@ -197,6 +214,27 @@ mod sort_properties_test {
                 SINGLETON,
                 SINGLETON,
             ),
+            (
+                "add: may collapse grouped values",
+                SortProperties::add,
+                GROUPED,
+                SINGLETON,
+                UNORDERED,
+            ),
+            (
+                "sub: may collapse grouped values",
+                SortProperties::sub,
+                GROUPED,
+                SINGLETON,
+                UNORDERED,
+            ),
+            (
+                "comparison does not preserve grouping",
+                SortProperties::gt_or_gteq,
+                GROUPED,
+                SINGLETON,
+                UNORDERED,
+            ),
             // `and` keeps ASC NULLS LAST / DESC NULLS FIRST, `or` keeps ASC
             // NULLS FIRST / DESC NULLS LAST. Both are commutative.
             (
@@ -482,7 +520,7 @@ impl Neg for SortProperties {
 #[derive(Debug, Clone)]
 pub struct ExprProperties {
     /// Properties that describe the sorting behavior of the expression,
-    /// such as whether it is ordered, unordered, or a singleton value.
+    /// such as whether it is ordered, grouped, unordered, or a singleton value.
     pub sort_properties: SortProperties,
     /// A closed interval representing the range of possible values for
     /// the expression. Used to compute reliable bounds.

datafusion/expr/src/udf.rs

@@ -360,8 +360,20 @@ impl ScalarUDF {
 
     /// Calculates the [`SortProperties`] of this function based on its
     /// children's properties.
+    ///
+    /// [`SortProperties::Grouped`] is retained only when the implementation
+    /// also reports that the transformation is strictly order-preserving. A
+    /// many-to-one function can otherwise make an output value occur in
+    /// multiple non-adjacent runs.
     pub fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
-        self.inner.output_ordering(inputs)
+        let sort_properties = self.inner.output_ordering(inputs)?;
+        if sort_properties == SortProperties::Grouped
+            && !self.inner.strictly_order_preserving(inputs)?
+        {
+            Ok(SortProperties::Unordered)
+        } else {
+            Ok(sort_properties)
+        }
     }
 
     pub fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result<bool> {
@@ -936,7 +948,12 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         Ok(Some(vec![]))
     }
 
-    /// Calculates the [`SortProperties`] of this function based on its children's properties.
+    /// Calculates the [`SortProperties`] of this function based on its children's
+    /// properties.
+    ///
+    /// The [`ScalarUDF`] wrapper retains a [`SortProperties::Grouped`] result
+    /// only when [`Self::strictly_order_preserving`] also returns `true` for
+    /// the same inputs.
     fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
         if !self.preserves_lex_ordering(inputs)? {
             return Ok(SortProperties::Unordered);

datafusion/ffi/src/expr/expr_properties.rs

@@ -67,13 +67,15 @@ pub enum FFI_SortProperties {
     Ordered(FFI_SortOptions),
     Unordered,
     Singleton,
+    Grouped,
 }
 
 impl From<&SortProperties> for FFI_SortProperties {
     fn from(value: &SortProperties) -> Self {
         match value {
             SortProperties::Unordered => FFI_SortProperties::Unordered,
             SortProperties::Singleton => FFI_SortProperties::Singleton,
+            SortProperties::Grouped => FFI_SortProperties::Grouped,
             SortProperties::Ordered(o) => FFI_SortProperties::Ordered(o.into()),
         }
     }
@@ -84,6 +86,7 @@ impl From<&FFI_SortProperties> for SortProperties {
         match value {
             FFI_SortProperties::Unordered => SortProperties::Unordered,
             FFI_SortProperties::Singleton => SortProperties::Singleton,
+            FFI_SortProperties::Grouped => SortProperties::Grouped,
             FFI_SortProperties::Ordered(o) => SortProperties::Ordered(o.into()),
         }
     }
@@ -113,3 +116,14 @@ impl From<&FFI_SortOptions> for SortOptions {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn grouped_sort_properties_round_trip() {
+        let ffi = FFI_SortProperties::from(&SortProperties::Grouped);
+        assert_eq!(SortProperties::from(&ffi), SortProperties::Grouped);
+    }
+}

datafusion/ffi/src/session/mod.rs

@@ -986,7 +986,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: [] }, geq_class: GroupingEquivalenceClass { groupings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, geq_cache: GroupingEquivalenceClass { groupings: [] }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None } }"
         );
 
         assert_eq!(

datafusion/ffi/src/tests/udf_udaf_udwf.rs

@@ -21,7 +21,7 @@ use arrow_schema::DataType;
 use datafusion_catalog::TableFunctionImpl;
 use datafusion_common::ScalarValue;
 use datafusion_common::config::ConfigOptions;
-use datafusion_expr::sort_properties::ExprProperties;
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{
     AggregateUDF, ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF,
     ScalarUDFImpl, Signature, Volatility, WindowUDF,
@@ -167,6 +167,11 @@ impl ScalarUDFImpl for PlacementUDF {
         &self,
         inputs: &[ExprProperties],
     ) -> datafusion_common::Result<bool> {
+        // This test-only sentinel verifies that the new `Grouped` variant
+        // survives the cross-library `ExprProperties` conversion.
+        if matches!(inputs, [input] if input.sort_properties == SortProperties::Grouped) {
+            return Ok(true);
+        }
         Ok(inputs.iter().all(|input| input.preserves_lex_ordering))
     }
 }

datafusion/ffi/src/udf/mod.rs

@@ -553,6 +553,7 @@ impl ScalarUDFImpl for ForeignScalarUDF {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use datafusion_expr::sort_properties::SortProperties;
 
     #[derive(Debug, PartialEq, Eq, Hash)]
     struct PlacementUDF {
@@ -594,6 +595,13 @@ mod tests {
                 return internal_err!("preserves_lex_ordering requires an input");
             }
 
+            // This test-only sentinel verifies that the new `Grouped` variant
+            // travels through the foreign path.
+            if matches!(inputs, [input] if input.sort_properties == SortProperties::Grouped)
+            {
+                return Ok(true);
+            }
+
             Ok(inputs.iter().all(|input| input.preserves_lex_ordering))
         }
 
@@ -692,6 +700,8 @@ mod tests {
                 .preserves_lex_ordering(&[preserves, does_not_preserve])
                 .unwrap()
         );
+        let grouped = ExprProperties::new_unknown().with_order(SortProperties::Grouped);
+        assert!(foreign_udf.preserves_lex_ordering(&[grouped]).unwrap());
         assert!(foreign_udf.preserves_lex_ordering(&[]).is_err());
 
         let updated = foreign_udf

datafusion/ffi/tests/ffi_udf.rs

@@ -27,7 +27,7 @@ mod tests {
     use datafusion::prelude::{SessionContext, col};
     use datafusion_execution::config::SessionConfig;
     use datafusion_expr::lit;
-    use datafusion_expr::sort_properties::ExprProperties;
+    use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
     use datafusion_ffi::tests::create_record_batch;
     use datafusion_ffi::tests::utils::get_module;
     use std::sync::Arc;
@@ -119,6 +119,9 @@ mod tests {
         assert!(foreign_func.preserves_lex_ordering(std::slice::from_ref(&preserves))?);
         assert!(!foreign_func.preserves_lex_ordering(&[preserves, does_not_preserve])?);
 
+        let grouped = ExprProperties::new_unknown().with_order(SortProperties::Grouped);
+        assert!(foreign_func.preserves_lex_ordering(&[grouped])?);
+
         Ok(())
     }
 

datafusion/functions/src/datetime/date_bin.rs

@@ -812,12 +812,14 @@ mod tests {
     use crate::datetime::date_bin::{DateBinFunc, date_bin_nanos_interval};
     use arrow::array::types::TimestampNanosecondType;
     use arrow::array::{Array, IntervalDayTimeArray, TimestampNanosecondArray};
+    use arrow::compute::SortOptions;
     use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
     use arrow::datatypes::{DataType, Field, FieldRef, TimeUnit};
 
     use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
     use datafusion_common::{DataFusionError, ScalarValue};
-    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+    use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl};
 
     use chrono::TimeDelta;
     use datafusion_common::config::ConfigOptions;
@@ -867,6 +869,26 @@ mod tests {
         );
     }
 
+    #[test]
+    fn date_bin_does_not_propagate_grouping() -> Result<(), DataFusionError> {
+        let udf = ScalarUDF::from(DateBinFunc::new());
+        let step = ExprProperties::new_unknown().with_order(SortProperties::Singleton);
+        let grouped = ExprProperties::new_unknown().with_order(SortProperties::Grouped);
+
+        assert_eq!(
+            udf.output_ordering(&[step.clone(), grouped])?,
+            SortProperties::Unordered
+        );
+
+        let ordered = ExprProperties::new_unknown()
+            .with_order(SortProperties::Ordered(SortOptions::default()));
+        assert_eq!(
+            udf.output_ordering(&[step, ordered])?,
+            SortProperties::Ordered(SortOptions::default())
+        );
+        Ok(())
+    }
+
     #[test]
     fn test_date_bin() {
         let return_field = &Arc::new(Field::new(