feat(cubesql): Top-down extractor for rewrites

cube-js · Sep 11, 2024 · 1b256a7 · 1b256a7
1 parent db2256d
commit 1b256a7
Show file tree

Hide file tree

Showing 8 changed files with 618 additions and 216 deletions.
diff --git a/packages/cubejs-backend-native/Cargo.lock b/packages/cubejs-backend-native/Cargo.lock
diff --git a/rust/cubesql/Cargo.lock b/rust/cubesql/Cargo.lock
diff --git a/rust/cubesql/cubesql/Cargo.toml b/rust/cubesql/cubesql/Cargo.toml
@@ -58,6 +58,7 @@ minijinja = { version = "1", features = ["json", "loader"] }
 lru = "0.12.1"
 sha2 = "0.10.8"
 bigdecimal = "0.4.2"
+indexmap = "1.9.3"
 
 
 [dev-dependencies]

diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -6962,6 +6962,7 @@ ORDER BY
                 MIN(month_count) min_month_count
             FROM count_by_month
             GROUP BY quarter0
+            ORDER BY quarter0
         "#;
 
         let expected_cube_scan = V1LoadRequestQuery {
@@ -16222,7 +16223,11 @@ ORDER BY "source"."str0" ASC
                 .wrapped_sql
                 .unwrap()
                 .sql;
-            assert!(sql.contains("\"limit\":1000"));
+            if Rewriter::top_down_extractor_enabled() {
+                assert!(sql.contains("LIMIT 1000"));
+            } else {
+                assert!(sql.contains("\"limit\":1000"));
+            }
             assert!(sql.contains("% 7"));
 
             let physical_plan = query_plan.as_physical_plan().await.unwrap();
@@ -18400,7 +18405,7 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
                 measure(count) AS cnt,
                 date_trunc('month', order_date) AS dt
             FROM KibanaSampleDataEcommerce
-            WHERE order_date IN (to_timestamp('2019-01-01 00:00:00.000000', 'YYYY-MM-DD HH24:MI:SS.US'))
+            WHERE date_trunc('month', order_date) IN (to_timestamp('2019-01-01 00:00:00.000000', 'YYYY-MM-DD HH24:MI:SS.US'))
             GROUP BY 2
             ;"#
             .to_string(),
@@ -18418,10 +18423,18 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
                 time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension {
                     dimension: "KibanaSampleDataEcommerce.order_date".to_string(),
                     granularity: Some("month".to_string()),
-                    date_range: Some(json!(vec![
-                        "2019-01-01T00:00:00.000Z".to_string(),
-                        "2019-01-01T00:00:00.000Z".to_string()
-                    ]))
+                    date_range: if Rewriter::top_down_extractor_enabled() {
+                        Some(json!(vec![
+                            "2019-01-01T00:00:00.000Z".to_string(),
+                            "2019-01-31T23:59:59.999Z".to_string()
+                        ]))
+                    } else {
+                        // Non-optimal variant with top down extractor disabled
+                        Some(json!(vec![
+                            "2019-01-01 00:00:00.000".to_string(),
+                            "2019-01-31 23:59:59.999".to_string()
+                        ]))
+                    }
                 }]),
                 order: None,
                 limit: None,

diff --git a/rust/cubesql/cubesql/src/compile/query_engine.rs b/rust/cubesql/cubesql/src/compile/query_engine.rs
@@ -186,7 +186,13 @@ pub trait QueryEngine {
         let mut rewriter = Rewriter::new(finalized_graph, cube_ctx.clone());
 
         let result = rewriter
-            .find_best_plan(root, state.auth_context().unwrap(), qtrace, span_id.clone())
+            .find_best_plan(
+                root,
+                state.auth_context().unwrap(),
+                qtrace,
+                span_id.clone(),
+                self.config_ref().top_down_extractor(),
+            )
             .await
             .map_err(|e| match e.cause {
                 CubeErrorCauseType::Internal(_) => CompilationError::Internal(