fix: agg input does not match union output when subquery (#17552)

* fix: agg input does not match union output when subquery

Signed-off-by: Kould <[email protected]>

* test: add fix case on subquery.test

Signed-off-by: Kould <[email protected]>

---------

Signed-off-by: Kould <[email protected]>

Author: KKould

src/query/sql/src/planner/optimizer/decorrelate/decorrelate.rs

@@ -286,7 +286,7 @@ impl SubqueryRewriter {
                 )?;
 
                 let mut join_type = JoinType::LeftSingle;
-                if subquery.contain_agg.unwrap() {
+                if matches!(subquery.contain_agg, Some(true)) {
                     let rel_expr = RelExpr::with_s_expr(&subquery.subquery);
                     let card = rel_expr
                         .derive_cardinality()?
@@ -473,7 +473,9 @@ impl SubqueryRewriter {
                 .table_index(column_entry.table_index())
                 .build(),
             });
-            let derive_column = self.derived_columns.get(correlated_column).unwrap();
+            let Some(derive_column) = self.derived_columns.get(correlated_column) else {
+                continue;
+            };
             let column_entry = metadata.column(*derive_column);
             let left_column = ScalarExpr::BoundColumnRef(BoundColumnRef {
                 span,

src/query/sql/src/planner/optimizer/decorrelate/flatten_plan.rs

@@ -786,6 +786,9 @@ impl SubqueryRewriter {
             flatten_info,
             need_cross_join,
         )?;
+        self.derived_columns
+            .retain(|_, derived_column| op.output_indexes.contains(derived_column));
+
         Ok(SExpr::create_binary(
             Arc::new(op.clone().into()),
             Arc::new(left_flatten_plan),

tests/sqllogictests/suites/query/subquery.test

@@ -912,3 +912,33 @@ group by input_date order by input_date;
 
 statement ok
 drop table test_group;
+
+# https://github.com/databendlabs/databend/issues/17476
+statement ok
+DROP TABLE IF EXISTS t1;
+
+statement ok
+DROP TABLE IF EXISTS t2;
+
+statement ok
+CREATE TABLE t1 (i1 INT);
+
+statement ok
+CREATE TABLE t2 (i1 INT);
+
+statement ok
+insert into t1 values (0);
+
+statement ok
+insert into t2 values (0);
+
+query I
+SELECT 1 FROM t1
+WHERE t1.i1 = (
+  SELECT t1.i1
+    FROM t2 UNION
+  SELECT dt1.i1
+    FROM (t1 AS dt1)
+);
+----
+1