feat: Implement IN and NOT IN list expression rewriting in SQL planner, enhancing filter handling for SQL queries

Author: longshan.lu

qurious/src/execution/session.rs

@@ -385,39 +385,36 @@ mod tests {
         // validate we actually have BRAZIL rows in the derived `all_nations` subquery.
         let debug = session.sql(
             "
-
 select
-    c_custkey,
-    c_name,
-    sum(l_extendedprice * (1 - l_discount)) as revenue,
-    c_acctbal,
-    n_name,
-    c_address,
-    c_phone,
-    c_comment
+    l_shipmode,
+    sum(case
+            when o_orderpriority = '1-URGENT'
+                or o_orderpriority = '2-HIGH'
+                then 1
+            else 0
+        end) as high_line_count,
+    sum(case
+            when o_orderpriority != '1-URGENT'
+                and o_orderpriority != '2-HIGH'
+                then 1
+            else 0
+        end) as low_line_count
 from
-    customer,
-    orders,
-    lineitem,
-    nation
+    lineitem
+        join
+    orders
+    on
+            l_orderkey = o_orderkey
 where
-        c_custkey = o_custkey
-  and l_orderkey = o_orderkey
-  and o_orderdate >= date '1993-10-01'
-  and o_orderdate < date '1994-01-01'
-  and l_returnflag = 'R'
-  and c_nationkey = n_nationkey
+        l_shipmode in ('MAIL', 'SHIP')
+  and l_commitdate < l_receiptdate
+  and l_shipdate < l_commitdate
+  and l_receiptdate >= date '1994-01-01'
+  and l_receiptdate < date '1995-01-01'
 group by
-    c_custkey,
-    c_name,
-    c_acctbal,
-    c_phone,
-    n_name,
-    c_address,
-    c_comment
+    l_shipmode
 order by
-    revenue desc
-limit 10;
+    l_shipmode;
 ",
         )?;
         print_batches(&debug)?;

qurious/src/logical/expr/aggregate.rs

@@ -12,6 +12,46 @@ use std::fmt::Display;
 use std::sync::Arc;
 
 use super::Column;
+use crate::datatypes::scalar::ScalarValue;
+
+/// Format an expression for *naming* purposes, stripping out CAST(...) (and nested alias) wrappers.
+///
+/// Type coercion may insert casts without changing the logical meaning of an expression; we don't
+/// want those casts to affect output field names, otherwise downstream column lookups can break
+/// (e.g. SUM(a*b) vs SUM(a*CAST(b AS ...))).
+fn fmt_expr_for_name(expr: &LogicalExpr) -> String {
+    match expr {
+        LogicalExpr::Cast(c) => fmt_expr_for_name(&c.expr),
+        LogicalExpr::Alias(a) => fmt_expr_for_name(&a.expr),
+        LogicalExpr::Column(c) => c.to_string(),
+        LogicalExpr::Literal(v) => v.to_string(),
+        LogicalExpr::Negative(e) => format!("- {}", fmt_expr_for_name(e)),
+        LogicalExpr::BinaryExpr(b) => format!(
+            "{} {} {}",
+            fmt_expr_for_name(&b.left),
+            b.op,
+            fmt_expr_for_name(&b.right)
+        ),
+        LogicalExpr::Case(case) => {
+            let mut s = String::from("CASE");
+            if let Some(op) = &case.operand {
+                s.push(' ');
+                s.push_str(&fmt_expr_for_name(op));
+            }
+            for (w, t) in &case.when_then {
+                s.push_str(" WHEN ");
+                s.push_str(&fmt_expr_for_name(w));
+                s.push_str(" THEN ");
+                s.push_str(&fmt_expr_for_name(t));
+            }
+            s.push_str(" ELSE ");
+            s.push_str(&fmt_expr_for_name(&case.else_expr));
+            s.push_str(" END");
+            s
+        }
+        other => other.to_string(),
+    }
+}
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum AggregateOperator {
@@ -96,10 +136,16 @@ pub struct AggregateExpr {
 impl AggregateExpr {
     pub fn field(&self, plan: &LogicalPlan) -> Result<FieldRef> {
         self.expr.field(plan).and_then(|field| {
-            let col_name = if let LogicalExpr::Column(inner) = self.expr.as_ref() {
-                &inner.qualified_name()
-            } else {
-                field.name()
+            // Use the *expression string* for non-column arguments, otherwise we may generate
+            // names like COUNT(i32) from Arrow field names which won't match expression display.
+            //
+            // Special case: COUNT(*) is rewritten to COUNT(1) by `CountWildcardRule`, but the
+            // output column name must remain COUNT(*) for SQL compatibility / tests.
+            let col_name = match (self.op.clone(), self.expr.as_ref()) {
+                (AggregateOperator::Count, LogicalExpr::Literal(ScalarValue::Int32(Some(1))))
+                | (AggregateOperator::Count, LogicalExpr::Literal(ScalarValue::Int64(Some(1)))) => "*".to_string(),
+                (_, LogicalExpr::Column(inner)) => inner.qualified_name(),
+                (_, other) => fmt_expr_for_name(other),
             };
 
             Ok(Arc::new(Field::new(
@@ -111,13 +157,30 @@ impl AggregateExpr {
     }
 
     pub(crate) fn as_column(&self) -> Result<LogicalExpr> {
-        self.expr.as_column().map(|inner_col| {
-            LogicalExpr::Column(Column {
-                name: format!("{}({})", self.op, inner_col),
-                relation: None,
-                is_outer_ref: false,
-            })
-        })
+        // Keep COUNT(*) naming stable even if it was rewritten to COUNT(1) internally.
+        if self.op == AggregateOperator::Count {
+            if matches!(
+                self.expr.as_ref(),
+                LogicalExpr::Literal(ScalarValue::Int32(Some(1))) | LogicalExpr::Literal(ScalarValue::Int64(Some(1)))
+            ) {
+                return Ok(LogicalExpr::Column(Column {
+                    name: "COUNT(*)".to_string(),
+                    relation: None,
+                    is_outer_ref: false,
+                }));
+            }
+        }
+
+        let arg_name = match self.expr.as_ref() {
+            LogicalExpr::Column(c) => c.to_string(),
+            other => fmt_expr_for_name(other),
+        };
+
+        Ok(LogicalExpr::Column(Column {
+            name: format!("{}({})", self.op, arg_name),
+            relation: None,
+            is_outer_ref: false,
+        }))
     }
 }
 
@@ -126,3 +189,117 @@ impl Display for AggregateExpr {
         write!(f, "{}({})", self.op, self.expr)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::datatypes::operator::Operator;
+    use crate::logical::expr::{BinaryExpr, CaseExpr, CastExpr, Column, LogicalExpr};
+    use crate::logical::plan::{EmptyRelation, LogicalPlan};
+    use arrow::datatypes::{DataType, Field, Schema};
+    use std::sync::Arc;
+
+    fn empty_plan_with_schema(fields: Vec<Field>) -> LogicalPlan {
+        LogicalPlan::EmptyRelation(EmptyRelation {
+            produce_one_row: true,
+            schema: Arc::new(Schema::new(fields)),
+        })
+    }
+
+    #[test]
+    fn count_star_keeps_output_name_after_rewrite_to_count_1() {
+        // Optimizer rule rewrites COUNT(*) -> COUNT(1) for execution.
+        // However, the output column name must remain COUNT(*) to match SQL surface semantics
+        // and sqllogictest expectations.
+        let plan = empty_plan_with_schema(vec![]);
+        let agg = AggregateExpr {
+            op: AggregateOperator::Count,
+            expr: Box::new(LogicalExpr::Literal(ScalarValue::Int32(Some(1)))),
+        };
+
+        let field = agg.field(&plan).unwrap();
+        assert_eq!(field.name(), "COUNT(*)");
+
+        let col = agg.as_column().unwrap();
+        assert_eq!(col.to_string(), "COUNT(*)");
+    }
+
+    #[test]
+    fn aggregate_naming_ignores_casts_in_argument_expression() {
+        // TypeCoercion may insert CASTs (e.g. to make DECIMAL * INT valid),
+        // but we don't want those CASTs to affect the aggregate output column name.
+        let plan = empty_plan_with_schema(vec![
+            Field::new("a", DataType::Decimal128(15, 2), false),
+            Field::new("b", DataType::Int64, false),
+        ]);
+
+        let expr = LogicalExpr::BinaryExpr(BinaryExpr::new(
+            LogicalExpr::Column(Column::new(
+                "a",
+                None::<crate::common::table_relation::TableRelation>,
+                false,
+            )),
+            Operator::Mul,
+            LogicalExpr::Cast(CastExpr::new(
+                LogicalExpr::Column(Column::new(
+                    "b",
+                    None::<crate::common::table_relation::TableRelation>,
+                    false,
+                )),
+                DataType::Decimal128(20, 0),
+            )),
+        ));
+
+        let agg = AggregateExpr {
+            op: AggregateOperator::Sum,
+            expr: Box::new(expr),
+        };
+
+        let field = agg.field(&plan).unwrap();
+        // cast is ignored for naming: b (not CAST(b AS ...))
+        assert_eq!(field.name(), "SUM(a * b)");
+        assert_eq!(agg.as_column().unwrap().to_string(), "SUM(a * b)");
+    }
+
+    #[test]
+    fn aggregate_naming_ignores_casts_inside_case_expression() {
+        // Similar to TPCH Q8: CASE branch literals may get casted by type coercion,
+        // but the aggregate output name should stay stable.
+        let plan = empty_plan_with_schema(vec![
+            Field::new("cond", DataType::Boolean, false),
+            Field::new("v", DataType::Decimal128(38, 4), false),
+        ]);
+
+        let case = CaseExpr {
+            operand: None,
+            when_then: vec![(
+                LogicalExpr::Column(Column::new(
+                    "cond",
+                    None::<crate::common::table_relation::TableRelation>,
+                    false,
+                )),
+                LogicalExpr::Column(Column::new(
+                    "v",
+                    None::<crate::common::table_relation::TableRelation>,
+                    false,
+                )),
+            )],
+            else_expr: Box::new(LogicalExpr::Cast(CastExpr::new(
+                LogicalExpr::Literal(ScalarValue::Int64(Some(0))),
+                DataType::Decimal128(38, 4),
+            ))),
+        };
+
+        let agg = AggregateExpr {
+            op: AggregateOperator::Sum,
+            expr: Box::new(LogicalExpr::Case(case)),
+        };
+
+        let field = agg.field(&plan).unwrap();
+        assert_eq!(field.name(), "SUM(CASE WHEN cond THEN v ELSE Int64(0) END)");
+        assert_eq!(
+            agg.as_column().unwrap().to_string(),
+            "SUM(CASE WHEN cond THEN v ELSE Int64(0) END)"
+        );
+    }
+}

qurious/src/optimizer/rule/rule_optimizer.rs

@@ -28,13 +28,15 @@ impl RuleBaseOptimizer {
         Self {
             rules: vec![
                 Box::new(CountWildcardRule),
-                Box::new(TypeCoercion),
                 Box::new(SimplifyExprs),
                 Box::new(ScalarSubqueryToJoin::default()),
                 Box::new(DecorrelatePredicateSubquery::default()),
                 Box::new(EliminateCrossJoin),
                 Box::new(ExtractEquijoinPredicate),
                 Box::new(PushdownFilter),
+                // Run type coercion late so correlated subqueries have been rewritten/decorrelated
+                // (avoids trying to type-check outer-ref columns inside subquery schemas).
+                Box::new(TypeCoercion),
             ],
         }
     }
@@ -54,3 +56,75 @@ impl Optimizer for RuleBaseOptimizer {
         Ok(current_plan)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{test_utils::sql_to_plan, utils};
+
+    #[test]
+    fn tpch_q2_should_keep_part_table_scan() {
+        // Regression test: optimizer must not drop the `part` relation for TPC-H Q2.
+        let sql = r#"
+select
+  s_acctbal,
+  s_name,
+  n_name,
+  p_partkey,
+  p_mfgr,
+  s_address,
+  s_phone,
+  s_comment
+from
+  part,
+  supplier,
+  partsupp,
+  nation,
+  region
+where
+  p_partkey = ps_partkey
+  and s_suppkey = ps_suppkey
+  and p_size = 15
+  and p_type like '%BRASS'
+  and s_nationkey = n_nationkey
+  and n_regionkey = r_regionkey
+  and r_name = 'EUROPE'
+  and ps_supplycost = (
+  select
+  min(ps_supplycost)
+  from
+  partsupp,
+  supplier,
+  nation,
+  region
+  where
+  p_partkey = ps_partkey
+  and s_suppkey = ps_suppkey
+  and s_nationkey = n_nationkey
+  and n_regionkey = r_regionkey
+  and r_name = 'EUROPE'
+)
+order by
+  s_acctbal desc,
+  n_name,
+  s_name,
+  p_partkey
+limit 10;
+"#;
+
+        let plan = sql_to_plan(sql);
+        let original = utils::format(&plan, 0);
+        assert!(
+            original.contains("TableScan: part"),
+            "original plan lost part scan:\n{original}"
+        );
+
+        let optimizer = RuleBaseOptimizer::new();
+        let optimized = optimizer.optimize(&plan).unwrap();
+        let formatted = utils::format(&optimized, 0);
+        assert!(
+            formatted.contains("TableScan: part"),
+            "optimized plan lost part scan:\n{formatted}"
+        );
+    }
+}