From 51f0dbedfd3f80cc09b3018cd9268f9e32d143c1 Mon Sep 17 00:00:00 2001
From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
Date: Mon, 25 Aug 2025 16:02:43 +0400
Subject: [PATCH] fix(cubesql): Merge subqueries with SQL push down

Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
---
 rust/cubesql/cubesql/src/compile/mod.rs       | 45 +++++++++++++++++++
 .../rewrite/rules/wrapper/aggregate.rs        |  4 +-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
index 455a60c7eed3c..4cc5f752d7496 100644
--- a/rust/cubesql/cubesql/src/compile/mod.rs
+++ b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -17504,4 +17504,49 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
             displayable(physical_plan.as_ref()).indent()
         );
     }
+
+    #[tokio::test]
+    async fn test_where_subquery_sql_push_down_measure_fn() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_testing_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            r#"
+            WITH top_customers AS (
+                SELECT
+                    KibanaSampleDataEcommerce.id,
+                    MEASURE(KibanaSampleDataEcommerce.sumPrice) AS sum_value
+                FROM KibanaSampleDataEcommerce
+                GROUP BY 1
+                ORDER BY 2 DESC
+                LIMIT 3
+            )
+            SELECT
+                KibanaSampleDataEcommerce.id,
+                MEASURE(KibanaSampleDataEcommerce.sumPrice) AS sum_value
+            FROM KibanaSampleDataEcommerce
+            WHERE KibanaSampleDataEcommerce.id IN (
+                SELECT id FROM top_customers
+            )
+            GROUP BY 1
+            ORDER BY 1
+            "#
+            .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+        .await;
+
+        let logical_plan = query_plan.as_logical_plan();
+        let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
+        assert!(sql.contains("IN (SELECT"));
+        assert!(sql.contains(r#"\\\"limit\\\": 3\\n"#));
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+    }
 }
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs
index 203cd1b2b578c..9f3945b44f57d 100644
--- a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs
+++ b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs
@@ -609,7 +609,7 @@ impl WrapperRules {
                             ),
                         ),
                         wrapper_pullup_replacer(
-                            wrapped_select_subqueries_empty_tail(),
+                            "?inner_subqueries",
                             wrapper_replacer_context(
                                 "?alias_to_cube",
                                 "WrapperReplacerContextPushToCube:true",
@@ -735,7 +735,7 @@ impl WrapperRules {
                         ),
                     ),
                     wrapper_pullup_replacer(
-                        wrapped_select_subqueries_empty_tail(),
+                        "?inner_subqueries",
                         wrapper_replacer_context(
                             "?alias_to_cube",
                             "WrapperReplacerContextPushToCube:true",