refactor: encapsulate most of the logic within ExprSet, and delineate the expr_identifier from the alias symbol

Author: wiedld

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -17,6 +17,7 @@
 
 //! Eliminate common sub-expression.
 
+use std::collections::hash_map::Entry;
 use std::collections::{BTreeSet, HashMap};
 use std::sync::Arc;
 
@@ -35,16 +36,42 @@ use datafusion_expr::expr::Alias;
 use datafusion_expr::logical_plan::{Aggregate, LogicalPlan, Projection, Window};
 use datafusion_expr::{col, Expr, ExprSchemable};
 
-/// A map from expression's identifier to tuple including
-///
-/// key == Identifier created with only the current node (& subtree)
-///
-/// values:
-/// - the expression itself (cloned)
-/// - counter
-/// - DataType of this expression.
-/// - symbol used as the identifier in the alias
-type ExprSet = HashMap<Identifier, (Expr, usize, DataType, Identifier)>;
+/// Set of expressions generated by the [`ExprIdentifierVisitor`]
+/// and consumed by the [`CommonSubexprRewriter`].
+#[derive(Default)]
+struct ExprSet {
+    /// A map from expression's identifier (stringified expr) to tuple including:
+    /// - the expression itself (cloned)
+    /// - counter
+    /// - DataType of this expression.
+    /// - symbol used as the identifier in the alias.
+    map: HashMap<Identifier, (Expr, usize, DataType, Identifier)>,
+}
+
+impl ExprSet {
+    fn expr_identifier(expr: &Expr) -> Identifier {
+        format!("{expr}")
+    }
+
+    fn get(&self, key: &Identifier) -> Option<&(Expr, usize, DataType, Identifier)> {
+        self.map.get(key)
+    }
+
+    fn entry(
+        &mut self,
+        key: Identifier,
+    ) -> Entry<'_, Identifier, (Expr, usize, DataType, Identifier)> {
+        self.map.entry(key)
+    }
+}
+
+impl From<Vec<(Identifier, (Expr, usize, DataType, Identifier))>> for ExprSet {
+    fn from(entries: Vec<(Identifier, (Expr, usize, DataType, Identifier))>) -> Self {
+        let mut expr_set = Self::default();
+        let _ = entries.into_iter().map(|(k, v)| expr_set.map.insert(k, v));
+        expr_set
+    }
+}
 
 /// Identifier for each subexpression.
 ///
@@ -134,7 +161,7 @@ impl CommonSubexprEliminate {
         config: &dyn OptimizerConfig,
     ) -> Result<LogicalPlan> {
         let mut window_exprs = vec![];
-        let mut expr_set = ExprSet::new();
+        let mut expr_set = ExprSet::default();
 
         // Get all window expressions inside the consecutive window operators.
         // Consecutive window expressions may refer to same complex expression.
@@ -206,7 +233,7 @@ impl CommonSubexprEliminate {
             input,
             ..
         } = aggregate;
-        let mut expr_set = ExprSet::new();
+        let mut expr_set = ExprSet::default();
 
         // build expr_set, with groupby and aggr
         let input_schema = Arc::clone(input.schema());
@@ -226,7 +253,7 @@ impl CommonSubexprEliminate {
         let new_group_expr = pop_expr(&mut new_expr)?;
 
         // create potential projection on top
-        let mut expr_set = ExprSet::new();
+        let mut expr_set = ExprSet::default();
         let new_input_schema = Arc::clone(new_input.schema());
         populate_expr_set(
             &new_aggr_expr,
@@ -277,9 +304,7 @@ impl CommonSubexprEliminate {
                         agg_exprs.push(expr.alias(&name));
                         proj_exprs.push(Expr::Column(Column::from_name(name)));
                     } else {
-                        let id = ExprIdentifierVisitor::<'static>::expr_identifier(
-                            &expr_rewritten,
-                        );
+                        let id = ExprSet::expr_identifier(&expr_rewritten);
                         let out_name =
                             expr_rewritten.to_field(&new_input_schema)?.qualified_name();
                         agg_exprs.push(expr_rewritten.alias(&id));
@@ -313,7 +338,7 @@ impl CommonSubexprEliminate {
         let inputs = plan.inputs();
         let input = inputs[0];
         let input_schema = Arc::clone(input.schema());
-        let mut expr_set = ExprSet::new();
+        let mut expr_set = ExprSet::default();
 
         // Visit expr list and build expr identifier to occuring count map (`expr_set`).
         populate_expr_set(&expr, input_schema, &mut expr_set, ExprMask::Normal)?;
@@ -572,10 +597,6 @@ enum VisitRecord {
 }
 
 impl ExprIdentifierVisitor<'_> {
-    fn expr_identifier(expr: &Expr) -> Identifier {
-        format!("{expr}")
-    }
-
     /// Find the first `EnterMark` in the stack, and accumulates every `ExprItem`
     /// before it.
     fn pop_enter_mark(&mut self) -> (usize, Identifier) {
@@ -619,21 +640,22 @@ impl TreeNodeVisitor for ExprIdentifierVisitor<'_> {
 
         // skip exprs should not be recognize.
         if self.expr_mask.ignores(expr) {
-            let curr_expr_identifier = Self::expr_identifier(expr);
+            let curr_expr_identifier = ExprSet::expr_identifier(expr);
             self.visit_stack
                 .push(VisitRecord::ExprItem(curr_expr_identifier));
             return Ok(TreeNodeRecursion::Continue);
         }
-        let curr_expr_identifier = Self::expr_identifier(expr);
-        let desc = format!("{curr_expr_identifier}{sub_expr_identifier}");
+        let curr_expr_identifier = ExprSet::expr_identifier(expr);
+        let alias_symbol = format!("{curr_expr_identifier}{sub_expr_identifier}");
 
-        self.visit_stack.push(VisitRecord::ExprItem(desc.clone()));
+        self.visit_stack
+            .push(VisitRecord::ExprItem(alias_symbol.clone()));
 
         let data_type = expr.get_type(&self.input_schema)?;
 
         self.expr_set
             .entry(curr_expr_identifier)
-            .or_insert_with(|| (expr.clone(), 0, data_type, desc))
+            .or_insert_with(|| (expr.clone(), 0, data_type, alias_symbol))
             .1 += 1;
         Ok(TreeNodeRecursion::Continue)
     }
@@ -666,12 +688,6 @@ struct CommonSubexprRewriter<'a> {
     affected_id: &'a mut BTreeSet<Identifier>,
 }
 
-impl CommonSubexprRewriter<'_> {
-    fn expr_identifier(expr: &Expr) -> Identifier {
-        format!("{expr}")
-    }
-}
-
 impl TreeNodeRewriter for CommonSubexprRewriter<'_> {
     type Node = Expr;
 
@@ -683,7 +699,7 @@ impl TreeNodeRewriter for CommonSubexprRewriter<'_> {
             return Ok(Transformed::new(expr, false, TreeNodeRecursion::Jump));
         }
 
-        let curr_id = &Self::expr_identifier(&expr);
+        let curr_id = &ExprSet::expr_identifier(&expr);
 
         // lookup previously visited expression
         match self.expr_set.get(curr_id) {
@@ -997,7 +1013,7 @@ mod test {
         let table_scan = test_table_scan().unwrap();
         let affected_id: BTreeSet<Identifier> =
             ["c+a".to_string(), "b+a".to_string()].into_iter().collect();
-        let expr_set_1 = [
+        let expr_set_1 = vec![
             (
                 "c+a".to_string(),
                 (col("c") + col("a"), 1, DataType::UInt32, "c+a".to_string()),
@@ -1007,9 +1023,8 @@ mod test {
                 (col("b") + col("a"), 1, DataType::UInt32, "b+a".to_string()),
             ),
         ]
-        .into_iter()
-        .collect();
-        let expr_set_2 = [
+        .into();
+        let expr_set_2 = vec![
             (
                 "c+a".to_string(),
                 (col("c+a"), 1, DataType::UInt32, "c+a".to_string()),
@@ -1019,8 +1034,7 @@ mod test {
                 (col("b+a"), 1, DataType::UInt32, "b+a".to_string()),
             ),
         ]
-        .into_iter()
-        .collect();
+        .into();
         let project =
             build_common_expr_project_plan(table_scan, affected_id.clone(), &expr_set_1)
                 .unwrap();
@@ -1046,7 +1060,7 @@ mod test {
             ["test1.c+test1.a".to_string(), "test1.b+test1.a".to_string()]
                 .into_iter()
                 .collect();
-        let expr_set_1 = [
+        let expr_set_1 = vec![
             (
                 "test1.c+test1.a".to_string(),
                 (
@@ -1066,9 +1080,8 @@ mod test {
                 ),
             ),
         ]
-        .into_iter()
-        .collect();
-        let expr_set_2 = [
+        .into();
+        let expr_set_2 = vec![
             (
                 "test1.c+test1.a".to_string(),
                 (
@@ -1088,8 +1101,7 @@ mod test {
                 ),
             ),
         ]
-        .into_iter()
-        .collect();
+        .into();
         let project =
             build_common_expr_project_plan(join, affected_id.clone(), &expr_set_1)
                 .unwrap();