feat: support normalized expr in CSE

zhuliquan · zhuliquan · commit 4e60434c6a13 · 2024-11-09T02:26:50.000+08:00
diff --git a/datafusion/common/src/cse.rs b/datafusion/common/src/cse.rs
@@ -50,33 +50,53 @@ impl<T: HashNode + ?Sized> HashNode for Arc<T> {
     }
 }
 
+/// A trait that defines how to normalize a node.
+///
+/// This trait is used to normalize nodes before comparing them for CSE. Normalization
+/// can be used to ensure that two nodes that are semantically equivalent are considered
+/// equal for CSE.
+/// For example：`a + b` and `b + a` are semantically equivalent.
+pub trait NormalizeNode: Eq {
+    fn normalize(&self) -> Self;
+    fn enable_normalized(&self) -> bool;
+}
+
 /// Identifier that represents a [`TreeNode`] tree.
 ///
 /// This identifier is designed to be efficient and  "hash", "accumulate", "equal" and
 /// "have no collision (as low as possible)"
-#[derive(Debug, Eq, PartialEq)]
-struct Identifier<'n, N> {
+#[derive(Debug, Eq)]
+struct Identifier<'n, N: NormalizeNode> {
     // Hash of `node` built up incrementally during the first, visiting traversal.
     // Its value is not necessarily equal to default hash of the node. E.g. it is not
     // equal to `expr.hash()` if the node is `Expr`.
     hash: u64,
     node: &'n N,
 }
 
-impl<N> Clone for Identifier<'_, N> {
+impl<N: NormalizeNode> Clone for Identifier<'_, N> {
     fn clone(&self) -> Self {
         *self
     }
 }
-impl<N> Copy for Identifier<'_, N> {}
+impl<N: NormalizeNode> Copy for Identifier<'_, N> {}
 
-impl<N> Hash for Identifier<'_, N> {
+impl<N: NormalizeNode> Hash for Identifier<'_, N> {
     fn hash<H: Hasher>(&self, state: &mut H) {
         state.write_u64(self.hash);
     }
 }
 
-impl<'n, N: HashNode> Identifier<'n, N> {
+impl<N: NormalizeNode> PartialEq for Identifier<'_, N> {
+    fn eq(&self, other: &Self) -> bool {
+        self.node.normalize() == other.node.normalize()
+    }
+}
+
+impl<'n, N> Identifier<'n, N>
+where
+    N: HashNode + NormalizeNode,
+{
     fn new(node: &'n N, random_state: &RandomState) -> Self {
         let mut hasher = random_state.build_hasher();
         node.hash_node(&mut hasher);
@@ -213,7 +233,11 @@ pub enum FoundCommonNodes<N> {
 ///
 /// A [`TreeNode`] without any children (column, literal etc.) will not have identifier
 /// because they should not be recognized as common subtree.
-struct CSEVisitor<'a, 'n, N, C: CSEController<Node = N>> {
+struct CSEVisitor<'a, 'n, N, C>
+where
+    N: NormalizeNode,
+    C: CSEController<Node = N>,
+{
     /// statistics of [`TreeNode`]s
     node_stats: &'a mut NodeStats<'n, N>,
 
@@ -244,7 +268,10 @@ struct CSEVisitor<'a, 'n, N, C: CSEController<Node = N>> {
 }
 
 /// Record item that used when traversing a [`TreeNode`] tree.
-enum VisitRecord<'n, N> {
+enum VisitRecord<'n, N>
+where
+    N: NormalizeNode,
+{
     /// Marks the beginning of [`TreeNode`]. It contains:
     /// - The post-order index assigned during the first, visiting traversal.
     EnterMark(usize),
@@ -258,7 +285,11 @@ enum VisitRecord<'n, N> {
     NodeItem(Identifier<'n, N>, bool),
 }
 
-impl<'n, N: TreeNode + HashNode, C: CSEController<Node = N>> CSEVisitor<'_, 'n, N, C> {
+impl<'n, N, C> CSEVisitor<'_, 'n, N, C>
+where
+    N: TreeNode + HashNode + NormalizeNode,
+    C: CSEController<Node = N>,
+{
     /// Find the first `EnterMark` in the stack, and accumulates every `NodeItem` before
     /// it. Returns a tuple that contains:
     /// - The pre-order index of the [`TreeNode`] we marked.
@@ -271,17 +302,26 @@ impl<'n, N: TreeNode + HashNode, C: CSEController<Node = N>> CSEVisitor<'_, 'n,
     ///   information up from children to parents via `visit_stack` during the first,
     ///   visiting traversal and no need to test the expression's validity beforehand with
     ///   an extra traversal).
-    fn pop_enter_mark(&mut self) -> (usize, Option<Identifier<'n, N>>, bool) {
-        let mut node_id = None;
+    fn pop_enter_mark(
+        &mut self,
+        enable_normalize: bool,
+    ) -> (usize, Option<Identifier<'n, N>>, bool) {
+        let mut node_ids: Vec<Identifier<'n, N>> = vec![];
         let mut is_valid = true;
 
         while let Some(item) = self.visit_stack.pop() {
             match item {
                 VisitRecord::EnterMark(down_index) => {
+                    if enable_normalize {
+                        node_ids.sort_by_key(|i| i.hash);
+                    }
+                    let node_id = node_ids
+                        .into_iter()
+                        .fold(None, |accum, item| Some(item.combine(accum)));
                     return (down_index, node_id, is_valid);
                 }
                 VisitRecord::NodeItem(sub_node_id, sub_node_is_valid) => {
-                    node_id = Some(sub_node_id.combine(node_id));
+                    node_ids.push(sub_node_id);
                     is_valid &= sub_node_is_valid;
                 }
             }
@@ -290,8 +330,10 @@ impl<'n, N: TreeNode + HashNode, C: CSEController<Node = N>> CSEVisitor<'_, 'n,
     }
 }
 
-impl<'n, N: TreeNode + HashNode + Eq, C: CSEController<Node = N>> TreeNodeVisitor<'n>
-    for CSEVisitor<'_, 'n, N, C>
+impl<'n, N, C> TreeNodeVisitor<'n> for CSEVisitor<'_, 'n, N, C>
+where
+    N: TreeNode + HashNode + NormalizeNode,
+    C: CSEController<Node = N>,
 {
     type Node = N;
 
@@ -331,7 +373,8 @@ impl<'n, N: TreeNode + HashNode + Eq, C: CSEController<Node = N>> TreeNodeVisito
     }
 
     fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
-        let (down_index, sub_node_id, sub_node_is_valid) = self.pop_enter_mark();
+        let (down_index, sub_node_id, sub_node_is_valid) =
+            self.pop_enter_mark(node.enable_normalized());
 
         let node_id = Identifier::new(node, self.random_state).combine(sub_node_id);
         let is_valid = C::is_valid(node) && sub_node_is_valid;
@@ -369,7 +412,11 @@ impl<'n, N: TreeNode + HashNode + Eq, C: CSEController<Node = N>> TreeNodeVisito
 /// Rewrite a [`TreeNode`] tree by replacing detected common subtrees with the
 /// corresponding temporary [`TreeNode`], that column contains the evaluate result of
 /// replaced [`TreeNode`] tree.
-struct CSERewriter<'a, 'n, N, C: CSEController<Node = N>> {
+struct CSERewriter<'a, 'n, N, C>
+where
+    N: NormalizeNode,
+    C: CSEController<Node = N>,
+{
     /// statistics of [`TreeNode`]s
     node_stats: &'a NodeStats<'n, N>,
 
@@ -386,8 +433,10 @@ struct CSERewriter<'a, 'n, N, C: CSEController<Node = N>> {
     controller: &'a mut C,
 }
 
-impl<N: TreeNode + Eq, C: CSEController<Node = N>> TreeNodeRewriter
-    for CSERewriter<'_, '_, N, C>
+impl<N, C> TreeNodeRewriter for CSERewriter<'_, '_, N, C>
+where
+    N: TreeNode + NormalizeNode,
+    C: CSEController<Node = N>,
 {
     type Node = N;
 
@@ -408,13 +457,30 @@ impl<N: TreeNode + Eq, C: CSEController<Node = N>> TreeNodeRewriter
                     self.down_index += 1;
                 }
 
-                let (node, alias) =
-                    self.common_nodes.entry(node_id).or_insert_with(|| {
-                        let node_alias = self.controller.generate_alias();
-                        (node, node_alias)
-                    });
-
-                let rewritten = self.controller.rewrite(node, alias);
+                // We *must* replace all original nodes with same `node_id`, not just the first
+                // node which is inserted into the common_nodes. This is because nodes with the same
+                // `node_id` are semantically equivalent, but not exactly the same.
+                //
+                // For example, `a + 1` and `1 + a` are semantically equivalent but not identical.
+                // In this case, we should replace the common expression `1 + a` with a new variable
+                // (e.g., `__common_cse_1`). So, `a + 1` and `1 + a` would both be replaced by
+                // `__common_cse_1`.
+                //
+                // The final result would be:
+                // - `__common_cse_1 as a + 1`
+                // - `__common_cse_1 as 1 + a`
+                //
+                // This way, we can efficiently handle semantically equivalent expressions without
+                // incorrectly treating them as identical.
+                let rewritten = if let Some((_, alias)) = self.common_nodes.get(&node_id)
+                {
+                    self.controller.rewrite(&node, alias)
+                } else {
+                    let node_alias = self.controller.generate_alias();
+                    let rewritten = self.controller.rewrite(&node, &node_alias);
+                    self.common_nodes.insert(node_id, (node, node_alias));
+                    rewritten
+                };
 
                 return Ok(Transformed::new(rewritten, true, TreeNodeRecursion::Jump));
             }
@@ -441,7 +507,11 @@ pub struct CSE<N, C: CSEController<Node = N>> {
     controller: C,
 }
 
-impl<N: TreeNode + HashNode + Clone + Eq, C: CSEController<Node = N>> CSE<N, C> {
+impl<N, C> CSE<N, C>
+where
+    N: TreeNode + HashNode + Clone + NormalizeNode,
+    C: CSEController<Node = N>,
+{
     pub fn new(controller: C) -> Self {
         Self {
             random_state: RandomState::new(),
@@ -557,6 +627,7 @@ impl<N: TreeNode + HashNode + Clone + Eq, C: CSEController<Node = N>> CSE<N, C>
     ) -> Result<FoundCommonNodes<N>> {
         let mut found_common = false;
         let mut node_stats = NodeStats::new();
+
         let id_arrays_list = nodes_list
             .iter()
             .map(|nodes| {
@@ -596,7 +667,9 @@ impl<N: TreeNode + HashNode + Clone + Eq, C: CSEController<Node = N>> CSE<N, C>
 #[cfg(test)]
 mod test {
     use crate::alias::AliasGenerator;
-    use crate::cse::{CSEController, HashNode, IdArray, Identifier, NodeStats, CSE};
+    use crate::cse::{
+        CSEController, HashNode, IdArray, Identifier, NodeStats, NormalizeNode, CSE,
+    };
     use crate::tree_node::tests::TestTreeNode;
     use crate::Result;
     use std::collections::HashSet;
@@ -662,6 +735,15 @@ mod test {
         }
     }
 
+    impl NormalizeNode for TestTreeNode<String> {
+        fn normalize(&self) -> Self {
+            self.clone()
+        }
+        fn enable_normalized(&self) -> bool {
+            false
+        }
+    }
+
     #[test]
     fn id_array_visitor() -> Result<()> {
         let alias_generator = AliasGenerator::new();
diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs
@@ -34,7 +34,7 @@ use crate::{
 };
 
 use arrow::datatypes::{DataType, FieldRef};
-use datafusion_common::cse::HashNode;
+use datafusion_common::cse::{HashNode, NormalizeNode};
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
 };
@@ -1666,6 +1666,69 @@ impl Expr {
     }
 }
 
+impl NormalizeNode for Expr {
+    fn enable_normalized(&self) -> bool {
+        #[allow(clippy::match_like_matches_macro)]
+        match self {
+            Expr::BinaryExpr(BinaryExpr {
+                op:
+                    _op @ (Operator::Plus
+                    | Operator::Multiply
+                    | Operator::BitwiseAnd
+                    | Operator::BitwiseOr
+                    | Operator::BitwiseXor
+                    | Operator::Eq
+                    | Operator::NotEq),
+                ..
+            }) => true,
+            _ => false,
+        }
+    }
+
+    fn normalize(&self) -> Expr {
+        match self {
+            Expr::BinaryExpr(BinaryExpr {
+                ref left,
+                ref op,
+                ref right,
+            }) => {
+                let normalized_left = left.normalize();
+                let normalized_right = right.normalize();
+                let new_binary = if matches!(
+                    op,
+                    Operator::Plus
+                        | Operator::Multiply
+                        | Operator::BitwiseAnd
+                        | Operator::BitwiseOr
+                        | Operator::BitwiseXor
+                        | Operator::Eq
+                        | Operator::NotEq
+                ) {
+                    let (l_expr, r_expr) =
+                        if format!("{normalized_left}") < format!("{normalized_right}") {
+                            (normalized_left, normalized_right)
+                        } else {
+                            (normalized_right, normalized_left)
+                        };
+                    BinaryExpr {
+                        left: Box::new(l_expr),
+                        op: *op,
+                        right: Box::new(r_expr),
+                    }
+                } else {
+                    BinaryExpr {
+                        left: Box::new(normalized_left),
+                        op: *op,
+                        right: Box::new(normalized_right),
+                    }
+                };
+                Expr::BinaryExpr(new_binary)
+            }
+            other => other.clone(),
+        }
+    }
+}
+
 impl HashNode for Expr {
     /// As it is pretty easy to forget changing this method when `Expr` changes the
     /// implementation doesn't use wildcard patterns (`..`, `_`) to catch changes
diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs
