fix(cubesql): SQL push down for limit and offset for ungrouped queries

cube-js · Oct 31, 2023 · 67da8c3 · 67da8c3
1 parent 2295fe0
commit 67da8c3
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs
@@ -577,7 +577,15 @@ impl CubeScanWrapperNode {
                                     } else {
                                         None
                                     };
-                                // TODO time dimensions, filters, segments, order, limit, offset
+
+                                if let Some(limit) = limit {
+                                    load_request.limit = Some(limit as i32);
+                                }
+
+                                if let Some(offset) = offset {
+                                    load_request.offset = Some(offset as i32);
+                                }
+                                // TODO time dimensions, filters, segments
 
                                 let sql_response = transport
                                     .sql(

diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -18825,6 +18825,42 @@ ORDER BY \"COUNT(count)\" DESC"
             .contains("ORDER BY \"case_when"));
     }
 
+    #[tokio::test]
+    async fn test_case_wrapper_with_internal_limit() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 LIMIT 1123"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await;
+
+        let logical_plan = query_plan.as_logical_plan();
+        assert!(logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql
+            .contains("CASE WHEN"));
+
+        assert!(logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql
+            .contains("1123"));
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+    }
+
     #[tokio::test]
     async fn test_case_wrapper_with_limit() {
         if !Rewriter::sql_push_down_enabled() {