refactor(cubesql): Add CubeEGraph and CubeRewrite type aliases (#8848)

cube-js · Oct 22, 2024 · 264f8da · 264f8da
1 parent b7306d2
commit 264f8da
Show file tree

Hide file tree

Showing 61 changed files with 538 additions and 787 deletions.
diff --git a/rust/cubesql/cubesql/src/compile/qtrace.rs b/rust/cubesql/cubesql/src/compile/qtrace.rs
@@ -1,14 +1,10 @@
 use std::{env, fs, sync::Arc};
 
-use super::rewrite::{
-    analysis::{LogicalPlanAnalysis, LogicalPlanData},
-    rewriter::IterInfo,
-    LogicalPlanLanguage,
-};
-use crate::compile::test::find_cube_scans_deep_search;
+use super::rewrite::{analysis::LogicalPlanData, rewriter::IterInfo, LogicalPlanLanguage};
+use crate::compile::{rewrite::rewriter::CubeEGraph, test::find_cube_scans_deep_search};
 use cubeclient::models::V1LoadRequestQuery;
 use datafusion::logical_plan::LogicalPlan;
-use egg::{EClass, EGraph, Iteration, Language};
+use egg::{EClass, Iteration, Language};
 use serde::Serialize;
 use sqlparser::ast::Statement;
 use uuid::Uuid;
@@ -82,10 +78,7 @@ impl Qtrace {
         self.statement(|stmt| stmt.set_optimized_plan(plan));
     }
 
-    pub fn set_original_graph(
-        &mut self,
-        egraph: &EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
-    ) {
+    pub fn set_original_graph(&mut self, egraph: &CubeEGraph) {
         self.statement(|stmt| stmt.set_original_graph(egraph));
     }
 
@@ -189,10 +182,7 @@ impl QtraceStatement {
         self.optimized_plan = Some(format!("{:?}", plan));
     }
 
-    pub fn set_original_graph(
-        &mut self,
-        egraph: &EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
-    ) {
+    pub fn set_original_graph(&mut self, egraph: &CubeEGraph) {
         self.original_graph = egraph
             .classes()
             .map(|eclass| QtraceEclass::make(eclass))

diff --git a/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs b/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs
@@ -1150,7 +1150,7 @@ impl LogicalPlanAnalysis {
     }
 
     fn eval_constant_expr(
-        egraph: &EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
+        egraph: &EGraph<LogicalPlanLanguage, Self>,
         expr: &Expr,
     ) -> Option<ConstantFolding> {
         let schema = DFSchema::empty();

diff --git a/rust/cubesql/cubesql/src/compile/rewrite/converter.rs b/rust/cubesql/cubesql/src/compile/rewrite/converter.rs
@@ -1,11 +1,14 @@
+pub use super::rewriter::CubeRunner;
 use crate::{
     compile::{
         engine::df::{
             scan::{CubeScanNode, CubeScanOptions, MemberField, WrappedSelectNode},
             wrapper::CubeScanWrapperNode,
         },
         rewrite::{
-            analysis::LogicalPlanAnalysis, extract_exprlist_from_groupping_set, rewriter::Rewriter,
+            analysis::LogicalPlanAnalysis,
+            extract_exprlist_from_groupping_set,
+            rewriter::{CubeEGraph, Rewriter},
             AggregateFunctionExprDistinct, AggregateFunctionExprFun, AggregateSplit,
             AggregateUDFExprFun, AliasExprAlias, AnyExprAll, AnyExprOp, BetweenExprNegated,
             BinaryExprOp, CastExprDataType, ChangeUserMemberValue, ColumnExprColumn,
@@ -52,7 +55,7 @@ use datafusion::{
     scalar::ScalarValue,
     sql::planner::ContextProvider,
 };
-use egg::{EGraph, Id, RecExpr};
+use egg::{Id, RecExpr};
 use itertools::Itertools;
 use serde_json::json;
 use std::{
@@ -62,8 +65,6 @@ use std::{
     sync::{Arc, LazyLock},
 };
 
-pub use super::rewriter::CubeRunner;
-
 macro_rules! add_data_node {
     ($converter:expr, $value_expr:expr, $field_variant:ident) => {
         $converter
@@ -125,10 +126,7 @@ macro_rules! add_binary_expr_list_node {
                 $flat_list
             )
         } else {
-            fn to_binary_tree(
-                graph: &mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
-                list: &[Id],
-            ) -> Id {
+            fn to_binary_tree(graph: &mut CubeEGraph, list: &[Id]) -> Id {
                 if list.len() == 0 {
                     graph.add(LogicalPlanLanguage::$field_variant(Vec::new()))
                 } else if list.len() == 1 {
@@ -186,7 +184,7 @@ static EXCLUDED_PARAM_VALUES: LazyLock<HashSet<ScalarValue>> = LazyLock::new(||
 });
 
 pub struct LogicalPlanToLanguageConverter {
-    graph: EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
+    graph: CubeEGraph,
     cube_context: Arc<CubeContext>,
     flat_list: bool,
 }
@@ -199,28 +197,22 @@ pub struct LogicalPlanToLanguageContext {
 impl LogicalPlanToLanguageConverter {
     pub fn new(cube_context: Arc<CubeContext>, flat_list: bool) -> Self {
         Self {
-            graph: EGraph::<LogicalPlanLanguage, LogicalPlanAnalysis>::new(
-                LogicalPlanAnalysis::new(
-                    cube_context.clone(),
-                    Arc::new(DefaultPhysicalPlanner::default()),
-                ),
-            ),
+            graph: CubeEGraph::new(LogicalPlanAnalysis::new(
+                cube_context.clone(),
+                Arc::new(DefaultPhysicalPlanner::default()),
+            )),
             cube_context,
             flat_list,
         }
     }
 
-    pub fn add_expr(
-        graph: &mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
-        expr: &Expr,
-        flat_list: bool,
-    ) -> Result<Id, CubeError> {
+    pub fn add_expr(graph: &mut CubeEGraph, expr: &Expr, flat_list: bool) -> Result<Id, CubeError> {
         // TODO: reference self?
         Self::add_expr_replace_params(graph, expr, &mut None, flat_list)
     }
 
     pub fn add_expr_replace_params(
-        graph: &mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
+        graph: &mut CubeEGraph,
         expr: &Expr,
         query_params: &mut Option<HashMap<usize, ScalarValue>>,
         flat_list: bool,
@@ -833,7 +825,7 @@ impl LogicalPlanToLanguageConverter {
         })
     }
 
-    pub fn take_egraph(self) -> EGraph<LogicalPlanLanguage, LogicalPlanAnalysis> {
+    pub fn take_egraph(self) -> CubeEGraph {
         self.graph
     }