feat(cubesql): Add separate ungrouped_scan flag to wrapper replacer c…

…ontext (#9120)

It is used to track whether `WrappedSelect` actually wraps ungrouped scan as opposed to old flag, which is used as push to Cube enabler.

Changes in cost are necessary because now ungrouped scans are tracked, we are getting proper values in `wrapped_select_ungrouped_scan` cost component, it would count the `WrappedSelect(ungrouped_scan=true)` nodes in extracted plan. And with old cost it would turns any ungrouped scan under wrapper more expensive than plan with same amount of wrappers (usually it's just 1 anyway), but with more nodes outside wrapper. Consider consuming projection: `Projection(WrappedSelect(ungrouped_scan=true))` vs `WrappedSelect(from=WrappedSelect(ungrouped_scan=true), ungrouped_scan=true)`. Plan with `Projection` would have `ast_size_outside_wrapper=1 wrapped_select_ungrouped_scan=1`, plan with `WrappedSelect` - `ast_size_outside_wrapper=0 wrapped_select_ungrouped_scan=2`, and second one is preferrable.

Also couple of related fixes:
* Mark distinct WrappedSelect as grouped
* Depend on a ungrouped flag for grouped join part: Wrapper can be ungrouped=true, push_to_cube=false, and this is unexpected in ungrouped-grouped join.

rust/cubesql/cubesql/src/compile/rewrite/cost.rs

@@ -255,8 +255,8 @@ pub struct CubePlanCost {
     non_pushed_down_limit_sort: i64,
     joins: usize,
     wrapper_nodes: i64,
-    wrapped_select_ungrouped_scan: usize,
     ast_size_outside_wrapper: usize,
+    wrapped_select_ungrouped_scan: usize,
     filters: i64,
     structure_points: i64,
     filter_members: i64,

rust/cubesql/cubesql/src/compile/rewrite/mod.rs

@@ -469,6 +469,13 @@ crate::plan_to_language! {
             // Known qualifiers of grouped subqueries
             // Used to allow to rewrite columns from them even with push to Cube enabled
             grouped_subqueries: Vec<String>,
+            // When `member` is logical plan this means it is actually ungrouped, even when push_to_cube is disabled.
+            // When `member` is expression it just acts as a pull-through from pushdown.
+            // It will be filled by every wrapper replacer producer rule, essentially same way as
+            // ungrouped_scan flag in wrapped_select is filled:
+            // fixed false for aggregation, copy inner value for projection.
+            // This flag should make roundtrip from top to bottom and back.
+            ungrouped_scan: bool,
         },
         WrapperPushdownReplacer {
             member: Arc<LogicalPlan>,
@@ -1974,9 +1981,10 @@ fn wrapper_replacer_context(
     in_projection: impl Display,
     cube_members: impl Display,
     grouped_subqueries: impl Display,
+    ungrouped_scan: impl Display,
 ) -> String {
     format!(
-        "(WrapperReplacerContext {alias_to_cube} {push_to_cube} {in_projection} {cube_members} {grouped_subqueries})",
+        "(WrapperReplacerContext {alias_to_cube} {push_to_cube} {in_projection} {cube_members} {grouped_subqueries} {ungrouped_scan})",
     )
 }
 

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

@@ -39,6 +39,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     )],
                     "?distinct",
@@ -51,6 +52,7 @@ impl WrapperRules {
                         "?in_projection",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 self.transform_agg_fun_expr("?fun", "?distinct", "?alias_to_cube"),

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

@@ -33,6 +33,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     ),
                     "?op",
@@ -44,6 +45,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     ),
                 ),
@@ -55,6 +57,7 @@ impl WrapperRules {
                         "?in_projection",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 self.transform_binary_expr("?op", "?alias_to_cube"),

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

@@ -33,6 +33,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     ),
                     wrapper_pullup_replacer(
@@ -43,6 +44,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     ),
                     wrapper_pullup_replacer(
@@ -53,6 +55,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     ),
                 ),
@@ -64,6 +67,7 @@ impl WrapperRules {
                         "?in_projection",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 self.transform_case_expr("?alias_to_cube"),

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

@@ -25,6 +25,7 @@ impl WrapperRules {
                         "?in_projection",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 wrapper_pullup_replacer(
@@ -35,6 +36,7 @@ impl WrapperRules {
                         "?in_projection",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
             ),
@@ -50,6 +52,7 @@ impl WrapperRules {
                         "WrapperReplacerContextInProjection:true",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 wrapper_pullup_replacer(
@@ -60,6 +63,7 @@ impl WrapperRules {
                         "WrapperReplacerContextInProjection:true",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 self.pushdown_simple_measure("?name", "?cube_members"),
@@ -75,6 +79,7 @@ impl WrapperRules {
                         "?in_projection",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 wrapper_pullup_replacer(
@@ -85,6 +90,7 @@ impl WrapperRules {
                         "?in_projection",
                         "?cube_members",
                         "?grouped_subqueries",
+                        "?ungrouped_scan",
                     ),
                 ),
                 self.pushdown_dimension(

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

@@ -6,9 +6,9 @@ use crate::{
         transforming_rewrite, wrapper_pullup_replacer, wrapper_replacer_context,
         CubeScanAliasToCube, CubeScanLimit, CubeScanOffset, CubeScanUngrouped, LogicalPlanLanguage,
         WrapperReplacerContextAliasToCube, WrapperReplacerContextGroupedSubqueries,
-        WrapperReplacerContextPushToCube,
+        WrapperReplacerContextPushToCube, WrapperReplacerContextUngroupedScan,
     },
-    var, var_iter,
+    copy_flag, var, var_iter,
 };
 use egg::Subst;
 
@@ -49,6 +49,7 @@ impl WrapperRules {
                             "WrapperReplacerContextInProjection:false",
                             "?members",
                             "?grouped_subqueries_out",
+                            "?ungrouped_scan_out",
                         ),
                     ),
                     "CubeScanWrapperFinalized:false",
@@ -62,6 +63,7 @@ impl WrapperRules {
                     "?alias_to_cube_out",
                     "?push_to_cube_out",
                     "?grouped_subqueries_out",
+                    "?ungrouped_scan_out",
                 ),
             ),
             rewrite(
@@ -85,6 +87,7 @@ impl WrapperRules {
         alias_to_cube_var_out: &'static str,
         push_to_cube_out_var: &'static str,
         grouped_subqueries_out_var: &'static str,
+        ungrouped_scan_out_var: &'static str,
     ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
         let members_var = var!(members_var);
         let alias_to_cube_var = var!(alias_to_cube_var);
@@ -94,6 +97,7 @@ impl WrapperRules {
         let alias_to_cube_var_out = var!(alias_to_cube_var_out);
         let push_to_cube_out_var = var!(push_to_cube_out_var);
         let grouped_subqueries_out_var = var!(grouped_subqueries_out_var);
+        let ungrouped_scan_out_var = var!(ungrouped_scan_out_var);
         move |egraph, subst| {
             let mut has_no_limit_or_offset = true;
             for limit in var_iter!(egraph[subst[limit_var]], CubeScanLimit).cloned() {
@@ -103,6 +107,17 @@ impl WrapperRules {
                 has_no_limit_or_offset &= offset.is_none();
             }
 
+            if !copy_flag!(
+                egraph,
+                subst,
+                ungrouped_cube_var,
+                CubeScanUngrouped,
+                ungrouped_scan_out_var,
+                WrapperReplacerContextUngroupedScan
+            ) {
+                return false;
+            }
+
             if let Some(_) = egraph[subst[members_var]].data.member_name_to_expr {
                 for alias_to_cube in
                     var_iter!(egraph[subst[alias_to_cube_var]], CubeScanAliasToCube).cloned()

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

@@ -51,7 +51,7 @@ impl WrapperRules {
                         "?select_alias",
                         "WrappedSelectDistinct:true",
                         "WrappedSelectPushToCube:false",
-                        "?select_ungrouped_scan",
+                        "WrappedSelectUngroupedScan:false",
                     ),
                     "?context",
                 ),

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

@@ -25,6 +25,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     ),
                     wrapper_pullup_replacer(
@@ -35,6 +36,7 @@ impl WrapperRules {
                             "?in_projection",
                             "?cube_members",
                             "?grouped_subqueries",
+                            "?ungrouped_scan",
                         ),
                     ),
                 ],
@@ -50,6 +52,7 @@ impl WrapperRules {
                     "?in_projection",
                     "?cube_members",
                     "?grouped_subqueries",
+                    "?ungrouped_scan",
                 ),
             ),
             self.transform_date_part_expr("?alias_to_cube"),