fix(cubesql): Replace alias to cube during wrapper pull up

Author: mcheshkov

rust/cubesql/cubesql/src/compile/rewrite/rules/members.rs

@@ -1437,7 +1437,7 @@ impl MemberRules {
         }
     }
 
-    fn replace_alias(
+    pub fn replace_alias(
         alias_to_cube: &Vec<(String, String)>,
         projection_alias: &Option<String>,
     ) -> Vec<(String, String)> {

rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/wrapper_pull_up.rs

@@ -2,14 +2,15 @@ use crate::{
     compile::rewrite::{
         cube_scan_wrapper,
         rewriter::{CubeEGraph, CubeRewrite},
-        rules::wrapper::WrapperRules,
+        rules::{members::MemberRules, wrapper::WrapperRules},
         transforming_rewrite, wrapped_select, wrapped_select_having_expr_empty_tail,
         wrapped_select_joins_empty_tail, wrapper_pullup_replacer, wrapper_replacer_context,
-        WrappedSelectSelectType, WrappedSelectType,
+        LogicalPlanLanguage, WrappedSelectAlias, WrappedSelectSelectType, WrappedSelectType,
+        WrapperReplacerContextAliasToCube,
     },
     var, var_iter, var_list_iter,
 };
-use egg::Subst;
+use egg::{Subst, Var};
 
 impl WrapperRules {
     pub fn wrapper_pull_up_rules(&self, rules: &mut Vec<CubeRewrite>) {
@@ -150,7 +151,7 @@ impl WrapperRules {
                             "?select_ungrouped_scan",
                         ),
                         wrapper_replacer_context(
-                            "?alias_to_cube",
+                            "?alias_to_cube_out",
                             // This is fixed to false for any LHS because we should only allow to push to Cube when from is ungrouped CubeScan
                             // And after pulling replacer over this node it will be WrappedSelect(from=CubeScan), so it should not allow to push for whatever LP is on top of it
                             "WrapperReplacerContextPushToCube:false",
@@ -162,7 +163,12 @@ impl WrapperRules {
                     ),
                     "CubeScanWrapperFinalized:false",
                 ),
-                self.transform_pull_up_wrapper_select("?cube_scan_input"),
+                self.transform_pull_up_wrapper_select(
+                    "?cube_scan_input",
+                    "?alias_to_cube",
+                    "?select_alias",
+                    "?alias_to_cube_out",
+                ),
             ),
             transforming_rewrite(
                 "wrapper-pull-up-to-cube-scan-non-trivial-wrapped-select",
@@ -338,7 +344,7 @@ impl WrapperRules {
                             "?select_ungrouped_scan",
                         ),
                         wrapper_replacer_context(
-                            "?alias_to_cube",
+                            "?alias_to_cube_out",
                             // This is fixed to false for any LHS because we should only allow to push to Cube when from is ungrouped CubeSCan
                             // And after pulling replacer over this node it will be WrappedSelect(from=WrappedSelect), so it should not allow to push for whatever LP is on top of it
                             "WrapperReplacerContextPushToCube:false",
@@ -359,20 +365,66 @@ impl WrapperRules {
                     "?inner_projection_expr",
                     "?inner_group_expr",
                     "?inner_aggr_expr",
+                    "?alias_to_cube",
+                    "?select_alias",
+                    "?alias_to_cube_out",
                 ),
             ),
         ]);
     }
 
+    fn replace_aliases(
+        egraph: &mut CubeEGraph,
+        subst: &mut Subst,
+        alias_to_cube_var: Var,
+        select_alias_var: Var,
+        alias_to_cube_out_var: Var,
+    ) -> bool {
+        for alias_to_cube in var_iter!(
+            egraph[subst[alias_to_cube_var]],
+            WrapperReplacerContextAliasToCube
+        ) {
+            for projection_alias in var_iter!(egraph[subst[select_alias_var]], WrappedSelectAlias) {
+                let replaced_alias_to_cube =
+                    MemberRules::replace_alias(&alias_to_cube, &projection_alias);
+                let new_alias_to_cube =
+                    egraph.add(LogicalPlanLanguage::WrapperReplacerContextAliasToCube(
+                        WrapperReplacerContextAliasToCube(replaced_alias_to_cube),
+                    ));
+                subst.insert(alias_to_cube_out_var, new_alias_to_cube);
+                return true;
+            }
+        }
+
+        false
+    }
+
     fn transform_pull_up_wrapper_select(
         &self,
         cube_scan_input_var: &'static str,
+        alias_to_cube_var: &'static str,
+        select_alias_var: &'static str,
+        alias_to_cube_out_var: &'static str,
     ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
         let cube_scan_input_var = var!(cube_scan_input_var);
+        let alias_to_cube_var = var!(alias_to_cube_var);
+        let select_alias_var = var!(select_alias_var);
+        let alias_to_cube_out_var = var!(alias_to_cube_out_var);
         move |egraph, subst| {
             for _ in var_list_iter!(egraph[subst[cube_scan_input_var]], WrappedSelect).cloned() {
                 return false;
             }
+
+            if !Self::replace_aliases(
+                egraph,
+                subst,
+                alias_to_cube_var,
+                select_alias_var,
+                alias_to_cube_out_var,
+            ) {
+                return false;
+            }
+
             true
         }
     }
@@ -387,12 +439,28 @@ impl WrapperRules {
         inner_projection_expr_var: &'static str,
         _inner_group_expr_var: &'static str,
         _inner_aggr_expr_var: &'static str,
+        alias_to_cube_var: &'static str,
+        select_alias_var: &'static str,
+        alias_to_cube_out_var: &'static str,
     ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
         let select_type_var = var!(select_type_var);
         let projection_expr_var = var!(projection_expr_var);
         let inner_select_type_var = var!(inner_select_type_var);
         let inner_projection_expr_var = var!(inner_projection_expr_var);
+        let alias_to_cube_var = var!(alias_to_cube_var);
+        let select_alias_var = var!(select_alias_var);
+        let alias_to_cube_out_var = var!(alias_to_cube_out_var);
         move |egraph, subst| {
+            if !Self::replace_aliases(
+                egraph,
+                subst,
+                alias_to_cube_var,
+                select_alias_var,
+                alias_to_cube_out_var,
+            ) {
+                return false;
+            }
+
             for select_type in
                 var_iter!(egraph[subst[select_type_var]], WrappedSelectSelectType).cloned()
             {