feat(cubesql): Push down limit through projection (#5206)

cube-js · Sep 1, 2022 · 3c6ff7d · 3c6ff7d
1 parent 8a4c403
commit 3c6ff7d
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 6 deletions.
diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -6165,6 +6165,69 @@ ORDER BY \"COUNT(count)\" DESC"
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_limit_push_down() -> Result<(), CubeError> {
+        // 1 level push down
+        let query_plan = convert_select_to_query_plan(
+            "SELECT l1.*, 1 as projection_should_exist_l1 FROM (\
+                    SELECT
+                      \"customer_gender\"
+                    FROM \"KibanaSampleDataEcommerce\"
+                    WHERE TRUE = TRUE
+                ) as l1 LIMIT 1000"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+        .await;
+
+        let logical_plan = query_plan.as_logical_plan();
+        assert_eq!(
+            logical_plan.find_cube_scan().request,
+            V1LoadRequestQuery {
+                measures: Some(vec![]),
+                segments: Some(vec![]),
+                dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
+                time_dimensions: None,
+                order: None,
+                limit: Some(1000),
+                offset: None,
+                filters: None,
+            }
+        );
+
+        // 2 levels push down
+        let query_plan = convert_select_to_query_plan(
+            "SELECT l2.*, 1 as projection_should_exist_l2 FROM (\
+                SELECT l1.*, 1 as projection_should_exist FROM (\
+                    SELECT
+                    \"customer_gender\"
+                    FROM \"KibanaSampleDataEcommerce\"
+                    WHERE TRUE = TRUE
+                ) as l1
+             ) as l2 LIMIT 1000"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+        .await;
+
+        let logical_plan = query_plan.as_logical_plan();
+        assert_eq!(
+            logical_plan.find_cube_scan().request,
+            V1LoadRequestQuery {
+                measures: Some(vec![]),
+                segments: Some(vec![]),
+                dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
+                time_dimensions: None,
+                order: None,
+                limit: Some(1000),
+                offset: None,
+                filters: None,
+            }
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_thought_spot_cte() -> Result<(), CubeError> {
         init_logger();
@@ -6201,8 +6264,7 @@ ORDER BY \"COUNT(count)\" DESC"
                 dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
                 time_dimensions: None,
                 order: None,
-                // TODO: limit: Some(1000),
-                limit: None,
+                limit: Some(1000),
                 offset: None,
                 filters: None,
             }
@@ -6286,8 +6348,7 @@ ORDER BY \"COUNT(count)\" DESC"
                 dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
                 time_dimensions: None,
                 order: None,
-                // TODO: limit: Some(1000),
-                limit: None,
+                limit: Some(1000),
                 offset: None,
                 filters: None,
             }

diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs
@@ -114,6 +114,22 @@ impl RewriteRules for FilterRules {
                 ),
                 self.push_down_limit_filter("?literal", "?new_limit", "?new_limit_n"),
             ),
+            // Transform Filter: Boolean(true)
+            // It's safe to push down filter under projection, next filter-truncate-true will truncate it
+            // TODO: Find a better solution how to drop filter node at all once
+            transforming_rewrite(
+                "push-down-filter-projection",
+                filter(
+                    literal_expr("?literal"),
+                    projection("?expr", "?input", "?alias"),
+                ),
+                projection(
+                    "?expr",
+                    filter(literal_expr("?literal"), "?input"),
+                    "?alias",
+                ),
+                self.match_filter_literal_true("?literal"),
+            ),
             rewrite(
                 "swap-limit-filter",
                 filter(
@@ -149,6 +165,12 @@ impl RewriteRules for FilterRules {
                     ),
                 ),
             ),
+            rewrite(
+                "limit-push-down-projection",
+                limit("?limit", projection("?expr", "?input", "?alias")),
+                projection("?expr", limit("?limit", "?input"), "?alias"),
+            ),
+            // Limit to top node
             rewrite(
                 "swap-limit-projection",
                 projection(
@@ -191,7 +213,7 @@ impl RewriteRules for FilterRules {
                 "filter-truncate-true",
                 filter_replacer(literal_expr("?literal"), "?alias_to_cube", "?members"),
                 cube_scan_filters_empty_tail(),
-                self.truncate_filter_literal_true("?literal"),
+                self.match_filter_literal_true("?literal"),
             ),
             transforming_rewrite(
                 "filter-replacer",
@@ -1005,7 +1027,7 @@ impl FilterRules {
         }
     }
 
-    fn truncate_filter_literal_true(
+    fn match_filter_literal_true(
         &self,
         literal_var: &'static str,
     ) -> impl Fn(&mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>, &mut Subst) -> bool {