test(cubesql): Tests for aggregation over limited CubeScan

Author: mcheshkov

rust/cubesql/cubesql/src/compile/test/test_wrapper.rs

@@ -1,4 +1,6 @@
+use cubeclient::models::V1LoadRequestQuery;
 use datafusion::physical_plan::displayable;
+use pretty_assertions::assert_eq;
 use std::sync::Arc;
 
 use crate::{
@@ -1001,3 +1003,127 @@ async fn test_wrapper_limit_zero() {
 
     let _physical_plan = query_plan.as_physical_plan().await.unwrap();
 }
+
+/// Regular aggregation over CubeScan(limit=n, ungrouped=true) is NOT pushed to CubeScan
+#[tokio::test]
+async fn wrapper_agg_over_limit() {
+    if !Rewriter::sql_push_down_enabled() {
+        return;
+    }
+    init_testing_logger();
+
+    let query_plan = convert_select_to_query_plan(
+        // language=PostgreSQL
+        r#"
+        SELECT
+            customer_gender
+        FROM (
+            SELECT
+                customer_gender
+            FROM
+                KibanaSampleDataEcommerce
+            LIMIT 5
+        ) scan
+        GROUP BY
+            1
+        "#
+        .to_string(),
+        DatabaseProtocol::PostgreSQL,
+    )
+    .await;
+
+    let physical_plan = query_plan.as_physical_plan().await.unwrap();
+    println!(
+        "Physical plan: {}",
+        displayable(physical_plan.as_ref()).indent()
+    );
+
+    let logical_plan = query_plan.as_logical_plan();
+    assert_eq!(
+        logical_plan.find_cube_scan().request,
+        V1LoadRequestQuery {
+            measures: Some(vec![]),
+            dimensions: Some(vec![]),
+            segments: Some(vec![]),
+            order: Some(vec![]),
+            limit: Some(5),
+            ungrouped: Some(true),
+            ..Default::default()
+        }
+    );
+
+    assert!(logical_plan
+        .find_cube_scan_wrapper()
+        .wrapped_sql
+        .unwrap()
+        .sql
+        .contains("\"limit\":5"));
+    assert!(query_plan
+        .as_logical_plan()
+        .find_cube_scan_wrapper()
+        .wrapped_sql
+        .unwrap()
+        .sql
+        .contains("\"ungrouped\":true"));
+}
+
+/// Aggregation(dimension) over CubeScan(limit=n, ungrouped=true) is NOT pushed to CubeScan
+#[tokio::test]
+async fn wrapper_agg_dimension_over_limit() {
+    if !Rewriter::sql_push_down_enabled() {
+        return;
+    }
+    init_testing_logger();
+
+    let query_plan = convert_select_to_query_plan(
+        // language=PostgreSQL
+        r#"
+        SELECT
+            MAX(customer_gender)
+        FROM (
+            SELECT
+                customer_gender
+            FROM
+                KibanaSampleDataEcommerce
+            LIMIT 5
+        ) scan
+        "#
+            .to_string(),
+        DatabaseProtocol::PostgreSQL,
+    )
+        .await;
+
+    let physical_plan = query_plan.as_physical_plan().await.unwrap();
+    println!(
+        "Physical plan: {}",
+        displayable(physical_plan.as_ref()).indent()
+    );
+
+    let logical_plan = query_plan.as_logical_plan();
+    assert_eq!(
+        logical_plan.find_cube_scan().request,
+        V1LoadRequestQuery {
+            measures: Some(vec![]),
+            dimensions: Some(vec![]),
+            segments: Some(vec![]),
+            order: Some(vec![]),
+            limit: Some(5),
+            ungrouped: Some(true),
+            ..Default::default()
+        }
+    );
+
+    assert!(logical_plan
+        .find_cube_scan_wrapper()
+        .wrapped_sql
+        .unwrap()
+        .sql
+        .contains("\"limit\":5"));
+    assert!(query_plan
+        .as_logical_plan()
+        .find_cube_scan_wrapper()
+        .wrapped_sql
+        .unwrap()
+        .sql
+        .contains("\"ungrouped\":true"));
+}