feat(cubesql): Support split and SUM(1)

cube-js · Mar 1, 2024 · 16e2ee0 · 16e2ee0
1 parent 76038a2
commit 16e2ee0
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 2 deletions.
diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -3305,6 +3305,32 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn sum_to_count_push_down() {
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            "SELECT SUM(1) AS \"count\" FROM \"public\".\"KibanaSampleDataEcommerce\" \"KibanaSampleDataEcommerce\"".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!["KibanaSampleDataEcommerce.count".to_string()]),
+                segments: Some(vec![]),
+                dimensions: Some(vec![]),
+                time_dimensions: None,
+                order: None,
+                limit: None,
+                offset: None,
+                filters: None,
+                ungrouped: None,
+            }
+        );
+    }
+
     #[tokio::test]
     async fn tableau_having_count_on_cube_without_count() {
         if !Rewriter::sql_push_down_enabled() {

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
@@ -76,6 +76,28 @@ impl SplitRules {
             true,
             rules,
         );
+        self.single_arg_split_point_rules(
+            "aggregate-function-sum-count-constant",
+            || agg_fun_expr("?fun_name", vec![literal_int(1)], "?distinct"),
+            || {
+                agg_fun_expr(
+                    "Count",
+                    vec![literal_int(1)],
+                    "AggregateFunctionExprDistinct:false",
+                )
+            },
+            |alias_column| agg_fun_expr("?output_fun_name", vec![alias_column], "?distinct"),
+            self.transform_aggregate_function(
+                Some("?fun_name"),
+                None,
+                Some("?distinct"),
+                "?output_fun_name",
+                "?alias_to_cube",
+                true,
+            ),
+            true,
+            rules,
+        );
         self.single_arg_split_point_rules(
             "aggregate-function-invariant-constant",
             || agg_fun_expr("?fun_name", vec!["?constant"], "?distinct"),

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
@@ -109,7 +109,7 @@ impl SplitRules {
                         "?orders",
                         "?limit",
                         "?offset",
-                        "CubeScanSplit:false",
+                        "CubeScanSplit:true",
                         "?can_pushdown_join",
                         "CubeScanWrapped:false",
                         "?ungrouped",
@@ -214,7 +214,7 @@ impl SplitRules {
                         "?orders",
                         "?limit",
                         "?offset",
-                        "CubeScanSplit:false",
+                        "CubeScanSplit:true",
                         "?can_pushdown_join",
                         "CubeScanWrapped:false",
                         "?ungrouped",