perf: reuse the projected equivalence group when only child orderings change

datafusion/physical-expr/src/equivalence/class.rs

@@ -310,6 +310,33 @@ pub struct EquivalenceGroup {
 }
 
 impl EquivalenceGroup {
+    /// A cheap, deliberately conservative check that two groups hold the same
+    /// equivalence classes.
+    ///
+    /// This is not `PartialEq`, and the distinction is the point. `classes` is a
+    /// `Vec` whose order carries no meaning -- `remove_class_at_idx` uses
+    /// `swap_remove` -- so two groups describing exactly the same equalities can
+    /// hold their classes in different orders and this returns `false` for them.
+    /// Naming it `PartialEq` would invite callers to read it as semantic equality,
+    /// which it is not.
+    ///
+    /// The comparison is positional because it runs on a hot path:
+    /// `ProjectionExec` consults it every time a rule replaces its child. Set
+    /// semantics would mean scanning the other group once per class, and that
+    /// quadratic term costs more than the recomputation the caller is trying to
+    /// skip, by a margin that widens with the number of classes.
+    ///
+    /// Only the false direction is reachable: a group can be reported different
+    /// when it is not, never the same when it is not. Callers using this to skip
+    /// work must be built so that a `false` merely costs them that work, which is
+    /// exactly how the projection fast path uses it.
+    ///
+    /// `map` is an index into `classes` and carries no information the classes do
+    /// not already have, so it takes no part in the comparison.
+    pub fn has_same_classes(&self, other: &Self) -> bool {
+        self.classes == other.classes
+    }
+
     /// Creates an equivalence group from the given equivalence classes.
     pub fn new(classes: impl IntoIterator<Item = EquivalenceClass>) -> Self {
         classes.into_iter().collect::<Vec<_>>().into()
@@ -924,7 +951,7 @@ mod tests {
     use super::*;
     use crate::equivalence::tests::create_test_params;
     use crate::expressions::{BinaryExpr, Column, binary, col, lit};
-    use arrow::datatypes::{DataType, Field, Schema};
+    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 
     use datafusion_expr::Operator;
 
@@ -1245,4 +1272,92 @@ mod tests {
 
         Ok(())
     }
+
+    /// Builds a group from a list of equated column pairs.
+    fn group_of(schema: &SchemaRef, pairs: &[(&str, &str)]) -> Result<EquivalenceGroup> {
+        let mut group = EquivalenceGroup::default();
+        for (lhs, rhs) in pairs {
+            group.add_equal_conditions(col(lhs, schema)?, col(rhs, schema)?);
+        }
+        Ok(group)
+    }
+
+    fn abc_schema() -> SchemaRef {
+        Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Int32, false),
+            Field::new("c", DataType::Int32, false),
+            Field::new("d", DataType::Int32, false),
+        ]))
+    }
+
+    #[test]
+    fn test_has_same_classes_is_conservative_about_class_order() -> Result<()> {
+        // `classes` is a `Vec` and `remove_class_at_idx` uses `swap_remove`, so
+        // the order two groups hold their classes in depends on how they were
+        // built. This check is positional and reports such a pair as different.
+        //
+        // That is deliberate, and it is why this is not `PartialEq`: set
+        // semantics would mean scanning the other group per class, and the
+        // quadratic cost dwarfs the recomputation the caller is trying to skip.
+        // The error can only go this way -- different when they match, never the
+        // reverse -- so a caller only forfeits an optimization.
+        let schema = abc_schema();
+        let ab_then_cd = group_of(&schema, &[("a", "b"), ("c", "d")])?;
+        let cd_then_ab = group_of(&schema, &[("c", "d"), ("a", "b")])?;
+
+        assert_eq!(ab_then_cd.len(), 2, "expected two disjoint classes");
+        assert!(!ab_then_cd.has_same_classes(&cd_then_ab));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_has_same_classes_compares_classes() -> Result<()> {
+        let schema = abc_schema();
+
+        // Two empty groups agree.
+        assert!(group_of(&schema, &[])?.has_same_classes(&group_of(&schema, &[])?));
+        // The same class, built the same way.
+        assert!(
+            group_of(&schema, &[("a", "b")])?
+                .has_same_classes(&group_of(&schema, &[("a", "b")])?)
+        );
+        // A class is a set, so the order within a pair is immaterial.
+        assert!(
+            group_of(&schema, &[("a", "b")])?
+                .has_same_classes(&group_of(&schema, &[("b", "a")])?)
+        );
+        // A populated group is not an empty one.
+        assert!(
+            !group_of(&schema, &[("a", "b")])?.has_same_classes(&group_of(&schema, &[])?)
+        );
+        // Equating a different pair yields a different group.
+        assert!(
+            !group_of(&schema, &[("a", "b")])?
+                .has_same_classes(&group_of(&schema, &[("a", "c")])?)
+        );
+        // Widening a class yields a different group.
+        assert!(
+            !group_of(&schema, &[("a", "b")])?
+                .has_same_classes(&group_of(&schema, &[("a", "b"), ("b", "c")])?)
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_has_same_classes_ignores_the_map() -> Result<()> {
+        // `map` indexes into `classes`, so equal classes must imply equal
+        // groups no matter how the classes were arrived at. Bridging `a = b`
+        // and `b = c` into one class must match stating `a = c` and `a = b`.
+        let schema = abc_schema();
+        let bridged = group_of(&schema, &[("a", "b"), ("b", "c")])?;
+        let direct = group_of(&schema, &[("a", "c"), ("a", "b")])?;
+
+        assert_eq!(bridged.len(), 1, "expected a single bridged class");
+        assert!(bridged.has_same_classes(&direct));
+
+        Ok(())
+    }
 }

datafusion/physical-expr/src/equivalence/properties/mod.rs

@@ -194,6 +194,31 @@ impl OrderingEquivalenceCache {
     }
 }
 
+/// Counts how often [`EquivalenceProperties::project_reusing`] reused a cached
+/// equivalence group instead of reprojecting it.
+///
+/// Thread local rather than a global counter: the test binary runs tests in
+/// parallel, and a shared count would let one test observe another's hits.
+#[cfg(test)]
+mod eq_group_reuse_probe {
+    use std::cell::Cell;
+
+    thread_local! {
+        static HITS: Cell<usize> = const { Cell::new(0) };
+    }
+
+    pub(super) fn record_hit() {
+        HITS.with(|h| h.set(h.get() + 1));
+    }
+
+    /// Runs `f` and reports how many reuses happened while it did.
+    pub(super) fn count<T>(f: impl FnOnce() -> T) -> (T, usize) {
+        let before = HITS.with(Cell::get);
+        let value = f();
+        (value, HITS.with(Cell::get) - before)
+    }
+}
+
 impl EquivalenceProperties {
     /// Helper used by the ordering equivalence rule when considering whether
     /// an expression can replace an existing sort key without invalidating
@@ -1188,6 +1213,48 @@ impl EquivalenceProperties {
     /// dropped.
     pub fn project(&self, mapping: &ProjectionMapping, output_schema: SchemaRef) -> Self {
         let eq_group = self.eq_group.project(mapping);
+        // Built here, so it satisfies the precondition by construction; going
+        // through the checked entry point would reproject it under
+        // `debug_assertions` for nothing.
+        self.project_with_eq_group_unchecked(mapping, output_schema, eq_group)
+    }
+
+    /// Projects `self`, reusing `cached`'s already-projected equivalence group
+    /// when `self`'s group is unchanged from `previous`'s.
+    ///
+    /// [`EquivalenceGroup::project`] is a pure function of the group and the
+    /// mapping, so when the group is unchanged the previous result can be handed
+    /// back rather than recomputed. Orderings are still derived here: they are
+    /// precisely what changes when a sort is introduced below this node.
+    ///
+    /// Falls back to a full projection when the groups differ, so there is no
+    /// precondition to violate. The caller does still have to pass a `cached`
+    /// that came from projecting `previous` through this same `mapping`; that
+    /// part cannot be checked here, and in practice it holds because the caller
+    /// carries its projection over untouched.
+    pub fn project_reusing(
+        &self,
+        mapping: &ProjectionMapping,
+        output_schema: SchemaRef,
+        previous: &EquivalenceProperties,
+        cached: &EquivalenceProperties,
+    ) -> Self {
+        let eq_group = if self.eq_group.has_same_classes(&previous.eq_group) {
+            #[cfg(test)]
+            eq_group_reuse_probe::record_hit();
+            cached.eq_group.clone()
+        } else {
+            self.eq_group.project(mapping)
+        };
+        self.project_with_eq_group_unchecked(mapping, output_schema, eq_group)
+    }
+
+    fn project_with_eq_group_unchecked(
+        &self,
+        mapping: &ProjectionMapping,
+        output_schema: SchemaRef,
+        eq_group: EquivalenceGroup,
+    ) -> Self {
         let orderings =
             self.projected_orderings(mapping, self.oeq_cache.normal_cls.clone());
         let normal_orderings = orderings
@@ -1559,3 +1626,127 @@ fn get_expr_properties(
         expr.get_properties(&child_states)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use crate::equivalence::tests::create_test_params;
+    use crate::expressions::col;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    /// Renames `a`, `b`, `c` and `d` so the projection is non-trivial while
+    /// still carrying every ordering and the `a = c` class across.
+    fn renaming_mapping(
+        schema: &SchemaRef,
+        output_schema: &SchemaRef,
+    ) -> Result<ProjectionMapping> {
+        [
+            ("a", "a1", 0),
+            ("b", "b1", 1),
+            ("c", "c1", 2),
+            ("d", "d1", 3),
+        ]
+        .into_iter()
+        .map(|(source, target, index)| {
+            Ok((
+                col(source, schema)?,
+                vec![(col(target, output_schema)?, index)].into(),
+            ))
+        })
+        .collect::<Result<Vec<_>>>()
+        .map(|entries| entries.into_iter().collect())
+    }
+
+    fn renamed_schema() -> SchemaRef {
+        Arc::new(Schema::new(vec![
+            Field::new("a1", DataType::Int32, true),
+            Field::new("b1", DataType::Int32, true),
+            Field::new("c1", DataType::Int32, true),
+            Field::new("d1", DataType::Int32, true),
+        ]))
+    }
+
+    #[test]
+    fn test_project_reusing_matches_project() -> Result<()> {
+        let (schema, eq_properties) = create_test_params()?;
+        let output_schema = renamed_schema();
+        let mapping = renaming_mapping(&schema, &output_schema)?;
+
+        let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema));
+
+        // Reusing the projection of an identical group must reproduce it exactly.
+        let (reused, hits) = eq_group_reuse_probe::count(|| {
+            eq_properties.project_reusing(
+                &mapping,
+                Arc::clone(&output_schema),
+                &eq_properties,
+                &baseline,
+            )
+        });
+        assert_eq!(
+            hits, 1,
+            "the reuse path was not taken, so this compares nothing"
+        );
+
+        // Guard against a vacuous comparison.
+        assert!(
+            !baseline.eq_group().is_empty(),
+            "the projection dropped the equivalence class, nothing is being tested"
+        );
+        assert!(
+            !baseline.oeq_class().is_empty(),
+            "the projection dropped every ordering, nothing is being tested"
+        );
+
+        assert!(
+            baseline.eq_group().has_same_classes(reused.eq_group()),
+            "equivalence group"
+        );
+        assert_eq!(baseline.oeq_class(), reused.oeq_class(), "orderings");
+        assert_eq!(baseline.constraints(), reused.constraints(), "constraints");
+        assert_eq!(baseline.schema(), reused.schema(), "schema");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_project_reusing_falls_back_when_the_group_moved() -> Result<()> {
+        // A `previous` whose group differs must not have its projection carried
+        // over. There is no precondition to violate here: the fallback is what
+        // keeps a mismatched pair from producing properties that disagree with
+        // the projection.
+        let (schema, eq_properties) = create_test_params()?;
+        let output_schema = renamed_schema();
+        let mapping = renaming_mapping(&schema, &output_schema)?;
+
+        let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema));
+        let unrelated = EquivalenceProperties::new(Arc::clone(&schema));
+        assert!(
+            !eq_properties
+                .eq_group()
+                .has_same_classes(unrelated.eq_group()),
+            "the two groups were meant to differ"
+        );
+
+        let (recomputed, hits) = eq_group_reuse_probe::count(|| {
+            eq_properties.project_reusing(
+                &mapping,
+                Arc::clone(&output_schema),
+                &unrelated,
+                &baseline,
+            )
+        });
+        assert_eq!(hits, 0, "a mismatched group was carried over anyway");
+
+        // Falling back must land on exactly what `project` would have produced.
+        assert!(
+            baseline.eq_group().has_same_classes(recomputed.eq_group()),
+            "equivalence group"
+        );
+        assert_eq!(baseline.oeq_class(), recomputed.oeq_class(), "orderings");
+
+        Ok(())
+    }
+}