Use CubeRewrite type alias in split rules

mcheshkov · mcheshkov · commit 11ff14c198a7 · 2024-10-22T17:33:11.000+03:00
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/aggregate_function.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/aggregate_function.rs
@@ -1,9 +1,9 @@
 use crate::{
     compile::rewrite::{
         agg_fun_expr, alias_expr,
-        analysis::{ConstantFolding, LogicalPlanAnalysis},
+        analysis::ConstantFolding,
         case_expr, cast_expr, column_expr, is_null_expr, literal_expr, literal_int,
-        rewriter::CubeEGraph,
+        rewriter::{CubeEGraph, CubeRewrite},
         rules::{members::MemberRules, split::SplitRules},
         AggregateFunctionExprDistinct, AggregateFunctionExprFun,
         AggregateSplitPushDownReplacerAliasToCube, ColumnExprColumn, LogicalPlanLanguage,
@@ -13,13 +13,9 @@ use crate::{
     var, var_iter,
 };
 use datafusion::{logical_plan::Column, physical_plan::aggregates::AggregateFunction};
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn aggregate_function_rules(
-        &self,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
-    ) {
+    pub fn aggregate_function_rules(&self, rules: &mut Vec<CubeRewrite>) {
         self.single_arg_split_point_rules_aggregate_function(
             "aggregate-function",
             || agg_fun_expr("?fun_name", vec![column_expr("?column")], "?distinct"),
@@ -222,7 +218,7 @@ impl SplitRules {
             + Send
             + Clone
             + 'static,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+        rules: &mut Vec<CubeRewrite>,
     ) {
         self.single_arg_split_point_rules_aggregate(
             name,
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/alias.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/alias.rs
@@ -1,12 +1,11 @@
 use crate::compile::rewrite::{
     aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer, alias_expr,
-    analysis::LogicalPlanAnalysis, projection_split_pullup_replacer,
-    projection_split_pushdown_replacer, rewrite, rules::split::SplitRules, LogicalPlanLanguage,
+    projection_split_pullup_replacer, projection_split_pushdown_replacer, rewrite,
+    rewriter::CubeRewrite, rules::split::SplitRules,
 };
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn alias_rules(&self, rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>) {
+    pub fn alias_rules(&self, rules: &mut Vec<CubeRewrite>) {
         rules.extend(vec![
             rewrite(
                 "split-alias-push-down-aggregate",
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/binary_expr.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/binary_expr.rs
@@ -1,15 +1,10 @@
 use crate::compile::rewrite::{
-    aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer,
-    analysis::LogicalPlanAnalysis, binary_expr, rewrite, rules::split::SplitRules,
-    LogicalPlanLanguage,
+    aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer, binary_expr, rewrite,
+    rewriter::CubeRewrite, rules::split::SplitRules,
 };
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn binary_expr_rules(
-        &self,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
-    ) {
+    pub fn binary_expr_rules(&self, rules: &mut Vec<CubeRewrite>) {
         rules.extend(vec![
             rewrite(
                 "split-binary-push-down-aggregate",
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/case.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/case.rs
@@ -1,12 +1,10 @@
 use crate::compile::rewrite::{
-    aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer,
-    analysis::LogicalPlanAnalysis, case_expr_var_arg, rewrite, rules::split::SplitRules,
-    LogicalPlanLanguage,
+    aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer, case_expr_var_arg, rewrite,
+    rewriter::CubeRewrite, rules::split::SplitRules,
 };
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn case_rules(&self, rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>) {
+    pub fn case_rules(&self, rules: &mut Vec<CubeRewrite>) {
         rules.extend(vec![
             rewrite(
                 "split-case-push-down-aggregate",
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/cast.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/cast.rs
@@ -1,9 +1,9 @@
 use crate::{
     compile::rewrite::{
         aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer,
-        analysis::{LogicalPlanAnalysis, OriginalExpr},
+        analysis::OriginalExpr,
         cast_expr, projection_split_pullup_replacer, projection_split_pushdown_replacer, rewrite,
-        rewriter::CubeEGraph,
+        rewriter::{CubeEGraph, CubeRewrite},
         rules::split::SplitRules,
         transforming_rewrite, AliasExprAlias, CastExprDataType, LiteralExprValue,
         LogicalPlanLanguage, ScalarFunctionExprFun,
@@ -16,10 +16,10 @@ use datafusion::{
     physical_plan::functions::BuiltinScalarFunction,
     scalar::ScalarValue,
 };
-use egg::{Rewrite, Subst};
+use egg::Subst;
 
 impl SplitRules {
-    pub fn cast_rules(&self, rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>) {
+    pub fn cast_rules(&self, rules: &mut Vec<CubeRewrite>) {
         rules.extend([
             transforming_rewrite(
                 "split-cast-push-down-aggregate",
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/column.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/column.rs
@@ -1,12 +1,11 @@
 use crate::compile::rewrite::{
-    aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer,
-    analysis::LogicalPlanAnalysis, column_expr, projection_split_pullup_replacer,
-    projection_split_pushdown_replacer, rewrite, rules::split::SplitRules, LogicalPlanLanguage,
+    aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer, column_expr,
+    projection_split_pullup_replacer, projection_split_pushdown_replacer, rewrite,
+    rewriter::CubeRewrite, rules::split::SplitRules,
 };
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn column_rules(&self, rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>) {
+    pub fn column_rules(&self, rules: &mut Vec<CubeRewrite>) {
         // TODO check for measures?
         rules.push(rewrite(
             "split-column-point-aggregate",
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/dates.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/dates.rs
@@ -1,8 +1,8 @@
 use crate::{
     compile::rewrite::{
-        analysis::{ConstantFolding, LogicalPlanAnalysis},
+        analysis::ConstantFolding,
         cast_expr, literal_expr,
-        rewriter::CubeEGraph,
+        rewriter::{CubeEGraph, CubeRewrite},
         rules::{
             members::min_granularity,
             split::SplitRules,
@@ -13,10 +13,9 @@ use crate::{
     var,
 };
 use datafusion::scalar::ScalarValue;
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn date_rules(&self, rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>) {
+    pub fn date_rules(&self, rules: &mut Vec<CubeRewrite>) {
         // TODO check for time dimension before push down to optimize performance
         // TODO use pass-through instead point rules for epoch
         self.single_arg_split_point_rules(
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/functions.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/functions.rs
@@ -1,14 +1,10 @@
 use crate::compile::rewrite::{
-    analysis::LogicalPlanAnalysis, is_not_null_expr, is_null_expr, literal_expr, negative_expr,
-    rules::split::SplitRules, udf_expr, ListType, LogicalPlanLanguage,
+    is_not_null_expr, is_null_expr, literal_expr, negative_expr, rewriter::CubeRewrite,
+    rules::split::SplitRules, udf_expr, ListType,
 };
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn functions_rules(
-        &self,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
-    ) {
+    pub fn functions_rules(&self, rules: &mut Vec<CubeRewrite>) {
         // Universal rule to traverse any number of function arguments
         // PushDown(ScalarFunctionExprArgs(..., arg, ...)) => ScalarFunctionExprArgs(..., PushDown(arg), ...)
         // ScalarFunctionExprArgs(..., PullUp(arg), ...) => PullUp(ScalarFunctionExprArgs(..., arg, ...))
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/granularity.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/granularity.rs
@@ -1,23 +1,19 @@
 use crate::{
     compile::rewrite::{
         alias_expr,
-        analysis::{ConstantFolding, LogicalPlanAnalysis},
+        analysis::ConstantFolding,
         binary_expr, cast_expr, cast_expr_explicit, column_expr, literal_expr, literal_float,
         literal_int, literal_string,
-        rewriter::CubeEGraph,
+        rewriter::{CubeEGraph, CubeRewrite},
         rules::split::SplitRules,
-        udf_expr, LogicalPlanLanguage,
+        udf_expr,
     },
     var,
 };
 use datafusion::{arrow::datatypes::DataType as ArrowDataType, scalar::ScalarValue};
-use egg::Rewrite;
 
 impl SplitRules {
-    pub fn granularity_rules(
-        &self,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
-    ) {
+    pub fn granularity_rules(&self, rules: &mut Vec<CubeRewrite>) {
         // CAST(CAST(((((EXTRACT(YEAR FROM "ta_1"."LO_COMMITDATE") * 100) + 1) * 100) + 1) AS varchar) AS date)
         self.single_arg_split_point_rules(
             "thoughtspot-year",
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/mod.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/mod.rs
@@ -11,11 +11,10 @@ pub mod top_level;
 
 use crate::{
     compile::rewrite::{
-        aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer, alias_expr,
-        analysis::LogicalPlanAnalysis,
-        fun_expr, fun_expr_var_arg, list_rewrite_with_lists_and_vars, original_expr_name,
+        aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer, alias_expr, fun_expr,
+        fun_expr_var_arg, list_rewrite_with_lists_and_vars, original_expr_name,
         projection_split_pullup_replacer, projection_split_pushdown_replacer, rewrite,
-        rewriter::{CubeEGraph, RewriteRules},
+        rewriter::{CubeEGraph, CubeRewrite, RewriteRules},
         rules::{members::MemberRules, replacer_flat_push_down_node, replacer_push_down_node},
         transforming_chain_rewrite, AliasExprAlias, ListApplierListPattern, ListPattern, ListType,
         LogicalPlanLanguage,
@@ -24,7 +23,6 @@ use crate::{
     transport::MetaContext,
     var,
 };
-use egg::Rewrite;
 use std::{fmt::Display, sync::Arc};
 
 pub struct SplitRules {
@@ -34,7 +32,7 @@ pub struct SplitRules {
 }
 
 impl RewriteRules for SplitRules {
-    fn rewrite_rules(&self) -> Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>> {
+    fn rewrite_rules(&self) -> Vec<CubeRewrite> {
         let mut rules = Vec::new();
 
         self.top_level_rules(&mut rules);
@@ -69,7 +67,7 @@ impl SplitRules {
         name: &str,
         node: impl Fn(String) -> String,
         projection_rules: bool,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+        rules: &mut Vec<CubeRewrite>,
     ) {
         rules.extend(vec![
             rewrite(
@@ -140,7 +138,7 @@ impl SplitRules {
         &self,
         fun_name: &(impl Display + ?Sized),
         projection_rules: bool,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+        rules: &mut Vec<CubeRewrite>,
     ) {
         rules.extend(vec![
             rewrite(
@@ -223,7 +221,7 @@ impl SplitRules {
             + Clone
             + 'static,
         projection_rules: bool,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+        rules: &mut Vec<CubeRewrite>,
     ) {
         if projection_rules {
             self.single_arg_split_point_rules_projection(
@@ -256,7 +254,7 @@ impl SplitRules {
             + Send
             + Clone
             + 'static,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+        rules: &mut Vec<CubeRewrite>,
     ) {
         rules.push(transforming_chain_rewrite(
             &format!("split-{}-point-aggregate", name),
@@ -292,7 +290,7 @@ impl SplitRules {
             + Send
             + Clone
             + 'static,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+        rules: &mut Vec<CubeRewrite>,
     ) {
         rules.push(transforming_chain_rewrite(
             &format!("split-{}-point-projection", name),
@@ -349,11 +347,7 @@ impl SplitRules {
         }
     }
 
-    fn list_pushdown_pullup_rules(
-        name: &str,
-        list_node: &str,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
-    ) {
+    fn list_pushdown_pullup_rules(name: &str, list_node: &str, rules: &mut Vec<CubeRewrite>) {
         let possible_inner_list_nodes = Self::possible_inner_list_nodes();
 
         // Aggregate split replacer
@@ -428,7 +422,7 @@ impl SplitRules {
     fn flat_list_pushdown_pullup_rules(
         name: &str,
         list_type: ListType,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+        rules: &mut Vec<CubeRewrite>,
     ) {
         let possible_inner_list_types = Self::possible_inner_flat_list_types();
 
@@ -525,7 +519,7 @@ impl SplitRules {
         substitute_list_node: &str,
         replacer_node: impl Fn(String, String, String) -> String,
         possible_inner_list_nodes: &Vec<String>,
-    ) -> Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>> {
+    ) -> Vec<CubeRewrite> {
         possible_inner_list_nodes
             .iter()
             .map(|inner_list_node| {
@@ -562,7 +556,7 @@ impl SplitRules {
         replacer_node: impl Fn(String, String, String) -> String,
         possible_inner_list_types: &Vec<ListType>,
         top_level_elem_vars: &[&str],
-    ) -> Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>> {
+    ) -> Vec<CubeRewrite> {
         possible_inner_list_types
             .iter()
             .map(|inner_list_type| {
@@ -608,7 +602,7 @@ impl SplitRules {
         push_down_replacer_node: impl Fn(String, String) -> String,
         pull_up_replacer_node: impl Fn(String, String, String) -> String,
         possible_inner_list_nodes: &Vec<String>,
-    ) -> Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>> {
+    ) -> Vec<CubeRewrite> {
         possible_inner_list_nodes
             .iter()
             .map(|inner_list_node| {
@@ -632,7 +626,7 @@ impl SplitRules {
         push_down_replacer_node: impl Fn(String, String) -> String,
         pull_up_replacer_node: impl Fn(String, String, String) -> String,
         possible_inner_list_types: &Vec<ListType>,
-    ) -> Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>> {
+    ) -> Vec<CubeRewrite> {
         possible_inner_list_types
             .iter()
             .map(|inner_list_type| {
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split/top_level.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split/top_level.rs
@@ -1,22 +1,20 @@
 use crate::{
     compile::rewrite::{
         aggr_aggr_expr_empty_tail, aggr_group_expr_empty_tail, aggregate,
-        aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer,
-        analysis::LogicalPlanAnalysis, cube_scan, projection, projection_expr_empty_tail,
-        projection_split_pullup_replacer, projection_split_pushdown_replacer, rewrite,
-        rewriter::CubeEGraph, rules::split::SplitRules, transforming_rewrite,
-        AggregateSplitPushDownReplacerAliasToCube, CubeScanAliasToCube, ListType,
-        LogicalPlanLanguage, ProjectionAlias, ProjectionSplitPushDownReplacerAliasToCube,
+        aggregate_split_pullup_replacer, aggregate_split_pushdown_replacer, cube_scan, projection,
+        projection_expr_empty_tail, projection_split_pullup_replacer,
+        projection_split_pushdown_replacer, rewrite,
+        rewriter::{CubeEGraph, CubeRewrite},
+        rules::split::SplitRules,
+        transforming_rewrite, AggregateSplitPushDownReplacerAliasToCube, CubeScanAliasToCube,
+        ListType, LogicalPlanLanguage, ProjectionAlias, ProjectionSplitPushDownReplacerAliasToCube,
     },
     var, var_iter,
 };
-use egg::{Rewrite, Subst};
+use egg::Subst;
 
 impl SplitRules {
-    pub fn top_level_rules(
-        &self,
-        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
-    ) {
+    pub fn top_level_rules(&self, rules: &mut Vec<CubeRewrite>) {
         rules.push(transforming_rewrite(
             "split-projection-aggregate",
             aggregate(
