fix(cubesql): Disallow pulling up ungrouped wrapper over CubeScan wit…

…h limit
cube-js · Nov 8, 2024 · b5fc0e4 · b5fc0e4
1 parent 4b9d94b
commit b5fc0e4
Showing 1 changed file with 23 additions and 3 deletions.
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cube_scan_wrapper.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cube_scan_wrapper.rs
@@ -3,8 +3,9 @@ use crate::{
         cube_scan, cube_scan_wrapper, rewrite,
         rewriter::{CubeEGraph, CubeRewrite},
         rules::wrapper::WrapperRules,
-        transforming_rewrite, wrapper_pullup_replacer, CubeScanAliasToCube, CubeScanUngrouped,
-        LogicalPlanLanguage, WrapperPullupReplacerAliasToCube, WrapperPullupReplacerUngrouped,
+        transforming_rewrite, wrapper_pullup_replacer, CubeScanAliasToCube, CubeScanLimit,
+        CubeScanOffset, CubeScanUngrouped, LogicalPlanLanguage, WrapperPullupReplacerAliasToCube,
+        WrapperPullupReplacerUngrouped,
     },
     var, var_iter,
 };
@@ -51,6 +52,8 @@ impl WrapperRules {
                 self.transform_wrap_cube_scan(
                     "?members",
                     "?alias_to_cube",
+                    "?limit",
+                    "?offset",
                     "?ungrouped",
                     "?alias_to_cube_out",
                     "?ungrouped_out",
@@ -77,27 +80,44 @@ impl WrapperRules {
         &self,
         members_var: &'static str,
         alias_to_cube_var: &'static str,
+        limit_var: &'static str,
+        offset_var: &'static str,
         ungrouped_cube_var: &'static str,
         alias_to_cube_var_out: &'static str,
         ungrouped_cube_var_out: &'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);
+        let limit_var = var!(limit_var);
+        let offset_var = var!(offset_var);
         let ungrouped_cube_var = var!(ungrouped_cube_var);
         let alias_to_cube_var_out = var!(alias_to_cube_var_out);
         let ungrouped_cube_var_out = var!(ungrouped_cube_var_out);
         move |egraph, subst| {
+            let mut has_no_limit_or_offset = true;
+            for limit in var_iter!(egraph[subst[limit_var]], CubeScanLimit).cloned() {
+                has_no_limit_or_offset &= limit.is_none();
+            }
+            for offset in var_iter!(egraph[subst[offset_var]], CubeScanOffset).cloned() {
+                has_no_limit_or_offset &= offset.is_none();
+            }
+
             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()
                 {
                     for ungrouped in
                         var_iter!(egraph[subst[ungrouped_cube_var]], CubeScanUngrouped).cloned()
                     {
+                        // When CubeScan already has limit or offset, it's unsafe to allow to push
+                        // anything on top to Cube.
+                        // Especially aggregation: aggregate does not commute with limit,
+                        // so it would be incorrect to join them to single CubeScan
+                        let ungrouped_out = ungrouped && has_no_limit_or_offset;
                         subst.insert(
                             ungrouped_cube_var_out,
                             egraph.add(LogicalPlanLanguage::WrapperPullupReplacerUngrouped(
-                                WrapperPullupReplacerUngrouped(ungrouped),
+                                WrapperPullupReplacerUngrouped(ungrouped_out),
                             )),
                         );
                         subst.insert(