clean & doc

Signed-off-by: Ruihang Xia <[email protected]>

Author: waynexia

datafusion/src/optimizer/common_subexpr_eliminate.rs

@@ -30,6 +30,26 @@ use crate::logical_plan::{
 };
 use crate::optimizer::optimizer::OptimizerRule;
 
+/// A map from expression's identifier to tuple including
+/// - the expression itself (cloned)
+/// - counter
+/// - DataType of this expression.
+type ExprSet = HashMap<Identifier, (Expr, usize, DataType)>;
+
+/// Identifier type. Current implementation use describe of a expression (type String) as
+/// Identifier.
+///
+/// A Identifier should (ideally) be able to "hash", "accumulate", "equal" and "have no
+/// collision (as low as possible)"
+///
+/// Since a identifier is likely to be copied many times, it is better that a identifier
+/// is small or "copy". otherwise some kinds of reference count is needed. String description
+/// here is not such a good choose.
+type Identifier = String;
+/// Perform Common Sub-expression Elimination optimization.
+///
+/// Currently only common sub-expressions within one logical plan will
+/// be eliminated.
 pub struct CommonSubexprEliminate {}
 
 impl OptimizerRule for CommonSubexprEliminate {
@@ -50,26 +70,19 @@ impl CommonSubexprEliminate {}
 
 fn optimize(plan: &LogicalPlan, execution_props: &ExecutionProps) -> Result<LogicalPlan> {
     let mut expr_set = ExprSet::new();
-    let mut addr_map = ExprAddrToId::new();
     let mut affected_id = HashSet::new();
 
     match plan {
         LogicalPlan::Projection {
             expr,
             input,
-            schema,
+            schema: _,
         } => {
             let mut arrays = vec![];
             for e in expr {
                 let data_type = e.get_type(input.schema())?;
                 let mut id_array = vec![];
-                expr_to_identifier(
-                    e,
-                    &mut expr_set,
-                    &mut addr_map,
-                    &mut id_array,
-                    data_type,
-                )?;
+                expr_to_identifier(e, &mut expr_set, &mut id_array, data_type)?;
                 arrays.push(id_array);
             }
 
@@ -78,32 +91,20 @@ fn optimize(plan: &LogicalPlan, execution_props: &ExecutionProps) -> Result<Logi
         LogicalPlan::Filter { predicate, input } => {
             let data_type = predicate.get_type(input.schema())?;
             let mut id_array = vec![];
-            expr_to_identifier(
-                predicate,
-                &mut expr_set,
-                &mut addr_map,
-                &mut id_array,
-                data_type,
-            )?;
+            expr_to_identifier(predicate, &mut expr_set, &mut id_array, data_type)?;
 
             return optimize(input, execution_props);
         }
         LogicalPlan::Window {
             input,
             window_expr,
-            schema,
+            schema: _,
         } => {
             let mut arrays = vec![];
             for e in window_expr {
                 let data_type = e.get_type(input.schema())?;
                 let mut id_array = vec![];
-                expr_to_identifier(
-                    e,
-                    &mut expr_set,
-                    &mut addr_map,
-                    &mut id_array,
-                    data_type,
-                )?;
+                expr_to_identifier(e, &mut expr_set, &mut id_array, data_type)?;
                 arrays.push(id_array);
             }
 
@@ -120,26 +121,14 @@ fn optimize(plan: &LogicalPlan, execution_props: &ExecutionProps) -> Result<Logi
             for e in group_expr {
                 let data_type = e.get_type(input.schema())?;
                 let mut id_array = vec![];
-                expr_to_identifier(
-                    e,
-                    &mut expr_set,
-                    &mut addr_map,
-                    &mut id_array,
-                    data_type,
-                )?;
+                expr_to_identifier(e, &mut expr_set, &mut id_array, data_type)?;
                 group_arrays.push(id_array);
             }
             let mut aggr_arrays = vec![];
             for e in aggr_expr {
                 let data_type = e.get_type(input.schema())?;
                 let mut id_array = vec![];
-                expr_to_identifier(
-                    e,
-                    &mut expr_set,
-                    &mut addr_map,
-                    &mut id_array,
-                    data_type,
-                )?;
+                expr_to_identifier(e, &mut expr_set, &mut id_array, data_type)?;
                 aggr_arrays.push(id_array);
             }
 
@@ -163,8 +152,7 @@ fn optimize(plan: &LogicalPlan, execution_props: &ExecutionProps) -> Result<Logi
 
             let mut new_input = optimize(input, execution_props)?;
             if !affected_id.is_empty() {
-                new_input =
-                    build_project_plan(new_input, affected_id, &expr_set, schema)?;
+                new_input = build_project_plan(new_input, affected_id, &expr_set)?;
             }
 
             return Ok(LogicalPlan::Aggregate {
@@ -174,7 +162,7 @@ fn optimize(plan: &LogicalPlan, execution_props: &ExecutionProps) -> Result<Logi
                 schema: schema.clone(),
             });
         }
-        LogicalPlan::Sort { expr, input } => {}
+        LogicalPlan::Sort { expr: _, input: _ } => {}
         LogicalPlan::Join { .. }
         | LogicalPlan::CrossJoin { .. }
         | LogicalPlan::Repartition { .. }
@@ -193,13 +181,12 @@ fn build_project_plan(
     input: LogicalPlan,
     affected_id: HashSet<Identifier>,
     expr_set: &ExprSet,
-    schema: &DFSchema,
 ) -> Result<LogicalPlan> {
     let mut project_exprs = vec![];
     let mut fields = vec![];
 
     for id in affected_id {
-        let (expr, _, _, _, data_type) = expr_set.get(&id).unwrap();
+        let (expr, _, data_type) = expr_set.get(&id).unwrap();
         // todo: check `nullable`
         fields.push(DFField::new(None, &id, data_type.clone(), true));
         project_exprs.push(expr.clone());
@@ -212,50 +199,44 @@ fn build_project_plan(
     })
 }
 
-// Helper struct & func
-
-/// A map from expression's identifier to tuple including
-/// - the expression itself (cloned)
-/// - a hash set contains all addresses with the same identifier.
-/// - counter
-/// - A alternative plan.
-pub type ExprSet =
-    HashMap<Identifier, (Expr, HashSet<*const Expr>, usize, Option<()>, DataType)>;
-
-pub type ExprAddrToId = HashMap<*const Expr, Identifier>;
-
-/// Identifier type. Current implementation use descript of a expression (type String) as
-/// Identifier.
+/// Go through an expression tree and generate identifier.
 ///
-/// A Identifier should (ideally) be "hashable", "accumulatable", "Eq", "collisionless"
-/// (as low probability as possible).
+/// An identifier contains information of the expression itself and its sub-expression.
+/// This visitor implementation use a stack `visit_stack` to track traversal, which
+/// lets us know when a sub-tree's visiting is finished. When `pre_visit` is called
+/// (traversing to a new node), an `EnterMark` and an `ExprItem` will be pushed into stack.
+/// And try to pop out a `EnterMark` on leaving a node (`post_visit()`). All `ExprItem`
+/// before the first `EnterMark` is considered to be sub-tree of the leaving node.
+///
+/// This visitor also records identifier in `id_array`. Makes the following traverse
+/// pass can get the identifier of a node without recalculate it. We assign each node
+/// in the expr tree a series number, start from 1, maintained by `series_number`.
+/// Series number represents the order we left (`post_visit`) a node. Has the property
+/// that child node's series number always smaller than parent's. While `id_array` is
+/// organized in the order we enter (`pre_visit`) a node. `node_count` helps us to
+/// get the index of `id_array` for each node.
 ///
-/// Since a identifier is likely to be copied many times, it is better that a identifier
-/// is small or "copy". otherwise some kinds of reference count is needed.
-pub type Identifier = String;
-
-/// Go through an expression tree and generate identifier.
 /// `Expr` without sub-expr (column, literal etc.) will not have identifier
 /// because they should not be recognized as common sub-expr.
 struct ExprIdentifierVisitor<'a> {
     // param
     expr_set: &'a mut ExprSet,
-    addr_map: &'a mut ExprAddrToId,
     /// series number (usize) and identifier.
     id_array: &'a mut Vec<(usize, Identifier)>,
     data_type: DataType,
 
     // inner states
-    visit_stack: Vec<Item>,
+    visit_stack: Vec<VisitRecord>,
     /// increased in pre_visit, start from 0.
     node_count: usize,
     /// increased in post_visit, start from 1.
-    post_visit_number: usize,
+    series_number: usize,
 }
 
-enum Item {
+/// Record item that used when traversing a expression tree.
+enum VisitRecord {
     /// `usize` is the monotone increasing series number assigned in pre_visit().
-    /// Starts from 0. Is used to index ithe dentifier array `id_array` in post_visit().
+    /// Starts from 0. Is used to index the identifier array `id_array` in post_visit().
     EnterMark(usize),
     /// Accumulated identifier of sub expression.
     ExprItem(Identifier),
@@ -311,35 +292,38 @@ impl ExprIdentifierVisitor<'_> {
         desc
     }
 
+    /// Find the first `EnterMark` in the stack, and accumulates every `ExprItem`
+    /// before it.
     fn pop_enter_mark(&mut self) -> (usize, Identifier) {
         let mut desc = String::new();
 
         while let Some(item) = self.visit_stack.pop() {
             match item {
-                Item::EnterMark(idx) => {
+                VisitRecord::EnterMark(idx) => {
                     return (idx, desc);
                 }
-                Item::ExprItem(s) => {
+                VisitRecord::ExprItem(s) => {
                     desc.push_str(&s);
                 }
             }
         }
 
-        (0, desc)
+        unreachable!("Enter mark should paired with node number");
     }
 }
 
 impl ExpressionVisitor for ExprIdentifierVisitor<'_> {
     fn pre_visit(mut self, _expr: &Expr) -> Result<Recursion<Self>> {
-        self.visit_stack.push(Item::EnterMark(self.node_count));
+        self.visit_stack
+            .push(VisitRecord::EnterMark(self.node_count));
         self.node_count += 1;
         // put placeholder
         self.id_array.push((0, "".to_string()));
         Ok(Recursion::Continue(self))
     }
 
     fn post_visit(mut self, expr: &Expr) -> Result<Self> {
-        self.post_visit_number += 1;
+        self.series_number += 1;
 
         let (idx, sub_expr_desc) = self.pop_enter_mark();
         // skip exprs should not be recognize.
@@ -350,58 +334,61 @@ impl ExpressionVisitor for ExprIdentifierVisitor<'_> {
                 | Expr::ScalarVariable(..)
                 | Expr::Wildcard
         ) {
-            self.id_array[idx].0 = self.post_visit_number;
+            self.id_array[idx].0 = self.series_number;
             let desc = Self::desc_expr(expr);
-            self.visit_stack.push(Item::ExprItem(desc));
+            self.visit_stack.push(VisitRecord::ExprItem(desc));
             return Ok(self);
         }
         let mut desc = Self::desc_expr(expr);
         desc.push_str(&sub_expr_desc);
 
-        self.id_array[idx] = (self.post_visit_number, desc.clone());
-        self.visit_stack.push(Item::ExprItem(desc.clone()));
+        self.id_array[idx] = (self.series_number, desc.clone());
+        self.visit_stack.push(VisitRecord::ExprItem(desc.clone()));
         let data_type = self.data_type.clone();
         self.expr_set
             .entry(desc.clone())
-            .or_insert_with(|| (expr.clone(), HashSet::new(), 0, None, data_type))
-            .2 += 1;
-        self.addr_map.insert(expr as *const Expr, desc);
+            .or_insert_with(|| (expr.clone(), 0, data_type))
+            .1 += 1;
         Ok(self)
     }
 }
 
-/// Go through an expression tree and generate identity.
+/// Go through an expression tree and generate identifier for every node in this tree.
 fn expr_to_identifier(
     expr: &Expr,
     expr_set: &mut ExprSet,
-    addr_map: &mut ExprAddrToId,
     id_array: &mut Vec<(usize, Identifier)>,
     data_type: DataType,
 ) -> Result<()> {
     expr.accept(ExprIdentifierVisitor {
         expr_set,
-        addr_map,
         id_array,
         data_type,
         visit_stack: vec![],
         node_count: 0,
-        post_visit_number: 0,
+        series_number: 0,
     })?;
 
     Ok(())
 }
 
+/// Rewrite expression by replacing detected common sub-expression with
+/// the corresponding temporary column name. That column contains the
+/// evaluate result of replaced expression.
 struct CommonSubexprRewriter<'a> {
     expr_set: &'a mut ExprSet,
     id_array: &'a Vec<(usize, Identifier)>,
+    /// Which identifier is replaced.
     affected_id: &'a mut HashSet<Identifier>,
 
+    /// the max series number we have rewritten. Other expression nodes
+    /// with smaller series number is already replaced and shouldn't
+    /// do anything with them.
     max_series_number: usize,
+    /// current node's information's index in `id_array`.
     curr_index: usize,
 }
 
-impl CommonSubexprRewriter<'_> {}
-
 impl ExprRewriter for CommonSubexprRewriter<'_> {
     fn pre_visit(&mut self, _: &Expr) -> Result<RewriteRecursion> {
         if self.curr_index >= self.id_array.len()
@@ -416,8 +403,7 @@ impl ExprRewriter for CommonSubexprRewriter<'_> {
             self.curr_index += 1;
             return Ok(RewriteRecursion::Continue);
         }
-        let (stored_expr, somewhat_set, counter, another_thing, _) =
-            self.expr_set.get(curr_id).unwrap();
+        let (_, counter, _) = self.expr_set.get(curr_id).unwrap();
         if *counter > 1 {
             self.affected_id.insert(curr_id.clone());
             Ok(RewriteRecursion::Mutate)
@@ -471,8 +457,16 @@ mod test {
     use crate::logical_plan::{binary_expr, col, lit, sum, LogicalPlanBuilder, Operator};
     use crate::test::*;
 
+    fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
+        let optimizer = CommonSubexprEliminate {};
+        let optimized_plan = optimizer
+            .optimize(plan, &ExecutionProps::new())
+            .expect("failed to optimize plan");
+        let formatted_plan = format!("{:?}", optimized_plan);
+        assert_eq!(formatted_plan, expected);
+    }
+
     #[test]
-    #[ignore]
     fn dev_driver_tpch_q1_simplified() -> Result<()> {
         // SQL:
         //  select
@@ -504,8 +498,11 @@ mod test {
             )?
             .build()?;
 
-        let optimizer = CommonSubexprEliminate {};
-        let new_plan = optimizer.optimize(&plan, &ExecutionProps::new()).unwrap();
+        let expected = "Aggregate: groupBy=[[]], aggr=[[SUM(#BinaryExpr-*BinaryExpr--Column-test.bLiteral1Column-test.a), SUM(#BinaryExpr-*BinaryExpr--Column-test.bLiteral1Column-test.a Multiply Int32(1) Plus #test.c)]]\
+        \n  Projection: #test.a Multiply Int32(1) Minus #test.b\
+        \n    TableScan: test projection=None";
+
+        assert_optimized_plan_eq(&plan, expected);
 
         Ok(())
     }