perf(cubesql): Make incremental rule scheduler distinguish between ge…

…neration 0 and 1
cube-js · Jun 17, 2024 · 7f4cfba · 7f4cfba
1 parent 23330b6
commit 7f4cfba
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs b/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs
@@ -35,7 +35,7 @@ pub type MemberNameToExpr = (Option<String>, Member, Expr);
 
 #[derive(Clone, Debug)]
 pub struct LogicalPlanData {
-    pub time: usize,
+    pub iteration_timestamp: usize,
     pub original_expr: Option<OriginalExpr>,
     pub member_name_to_expr: Option<Vec<MemberNameToExpr>>,
     pub trivial_push_down: Option<usize>,
@@ -222,7 +222,9 @@ impl Member {
 
 #[derive(Clone)]
 pub struct LogicalPlanAnalysis {
-    pub time: usize,
+    /* This is 0, when creating the EGraph.  It's set to 1 before iteration 0,
+    2 before the iteration 1, etc. */ 
+    pub iteration_timestamp: usize,
     cube_context: Arc<CubeContext>,
     planner: Arc<DefaultPhysicalPlanner>,
 }
@@ -254,7 +256,7 @@ impl<'a> Index<Id> for SingleNodeIndex<'a> {
 impl LogicalPlanAnalysis {
     pub fn new(cube_context: Arc<CubeContext>, planner: Arc<DefaultPhysicalPlanner>) -> Self {
         Self {
-            time: 0,
+            iteration_timestamp: 0,
             cube_context,
             planner,
         }
@@ -1257,7 +1259,7 @@ impl Analysis<LogicalPlanLanguage> for LogicalPlanAnalysis {
         enode: &LogicalPlanLanguage,
     ) -> Self::Data {
         LogicalPlanData {
-            time: egraph.analysis.time,
+            iteration_timestamp: egraph.analysis.iteration_timestamp,
             original_expr: Self::make_original_expr(egraph, enode),
             member_name_to_expr: Self::make_member_name_to_expr(egraph, enode),
             trivial_push_down: Self::make_trivial_push_down(egraph, enode),
@@ -1297,7 +1299,7 @@ impl Analysis<LogicalPlanLanguage> for LogicalPlanAnalysis {
             | column_name
             | filter_operators
             | is_empty_list
-            | self.merge_max_field(&mut a.time, b.time)
+            | self.merge_max_field(&mut a.iteration_timestamp, b.iteration_timestamp)
     }
 
     fn modify(egraph: &mut EGraph<LogicalPlanLanguage, Self>, id: Id) {
@@ -1337,4 +1339,18 @@ impl Analysis<LogicalPlanLanguage> for LogicalPlanAnalysis {
     fn allow_ematching_cycles(&self) -> bool {
         false
     }
+
+    #[cfg(any())]
+    fn pre_union(
+        egraph: &EGraph<LogicalPlanLanguage, Self>,
+        id1: Id,
+        id2: Id,
+        _justification: &Option<egg::Justification>
+    ) {
+        let eclass1: &egg::EClass<LogicalPlanLanguage, LogicalPlanData> = &egraph[id1];
+        let data1 = &eclass1.data;
+        let eclass2 = &egraph[id2];
+        let data2 = &eclass2.data;
+        log::warn!("id1 = {}, time = {}, id2 = {}, id2 time = {}", id1, data1.time, id2, data2.time);
+    }
 }
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rewriter.rs b/rust/cubesql/cubesql/src/compile/rewrite/rewriter.rs
@@ -221,7 +221,7 @@ impl Rewriter {
         ))
         .with_scheduler(IncrementalScheduler::default())
         .with_hook(|runner| {
-            runner.egraph.analysis.time = runner.iterations.len();
+            runner.egraph.analysis.iteration_timestamp = runner.iterations.len() + 1;
             Ok(())
         })
         .with_egraph(egraph)
@@ -587,7 +587,7 @@ impl egg::RewriteScheduler<LogicalPlanLanguage, LogicalPlanAnalysis> for Increme
             self.current_eclasses.extend(
                 egraph
                     .classes()
-                    .filter_map(|class| (class.data.time + 1 >= iteration).then(|| class.id)),
+                    .filter_map(|class| (class.data.iteration_timestamp >= iteration).then(|| class.id)),
             );
         };
         assert_eq!(iteration, self.current_iter);