feat: Introduce CASE expression support in SQL parser and planner, enhancing expression handling and type coercion for conditional logic

Author: longshan.lu

qurious/src/common/table_schema.rs

@@ -10,6 +10,14 @@ use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef};
 
 pub type TableSchemaRef = Arc<TableSchema>;
 
+/// Arrow schema metadata key used to preserve per-field qualifiers (table/alias) across planning stages.
+///
+/// This is needed because Arrow `Schema` fields are identified by name only, and we allow duplicate
+/// column names across different relations (e.g. `nation n1`, `nation n2` both have `n_name`).
+/// Physical planning must be able to map a `(relation, column_name)` to the correct field index.
+pub const FIELD_QUALIFIERS_META_KEY: &str = "qurious.field_qualifiers";
+const FIELD_QUALIFIERS_META_SEP: char = '\u{1f}'; // unit separator (unlikely to appear in names)
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct TableSchema {
     pub schema: SchemaRef,
@@ -48,7 +56,24 @@ impl TableSchema {
     }
 
     pub fn arrow_schema(&self) -> SchemaRef {
-        self.schema.clone()
+        // Preserve qualifier information via Schema metadata so physical planning can disambiguate
+        // same-named fields from different relations.
+        let mut metadata = self.schema.metadata().clone();
+        let qualifiers = self
+            .field_qualifiers
+            .iter()
+            .map(|q| q.as_ref().map(|t| t.to_qualified_name()).unwrap_or_default())
+            .collect::<Vec<_>>()
+            .join(&FIELD_QUALIFIERS_META_SEP.to_string());
+        metadata.insert(FIELD_QUALIFIERS_META_KEY.to_string(), qualifiers);
+
+        let fields = self
+            .schema
+            .fields()
+            .iter()
+            .map(|f| f.as_ref().clone())
+            .collect::<Vec<_>>();
+        Arc::new(Schema::new_with_metadata(fields, metadata))
     }
 
     pub fn has_field(&self, qualifier: Option<&TableRelation>, name: &str) -> bool {

qurious/src/datatypes/scalar.rs

@@ -213,6 +213,18 @@ impl ScalarValue {
             DataType::LargeUtf8 => typed_cast!(array, index, LargeStringArray, Utf8),
             DataType::Decimal128(p, s) => typed_cast_decimal!(Decimal128Array, Decimal128, array, index, *p, *s),
             DataType::Decimal256(p, s) => typed_cast_decimal!(Decimal256Array, Decimal256, array, index, *p, *s),
+            DataType::Interval(IntervalUnit::MonthDayNano) => {
+                let arr = array
+                    .as_any()
+                    .downcast_ref::<IntervalMonthDayNanoArray>()
+                    .ok_or_else(|| {
+                        Error::InternalError(format!(
+                            "could not cast value to {}",
+                            type_name::<IntervalMonthDayNanoArray>()
+                        ))
+                    })?;
+                Ok(ScalarValue::IntervalMonthDayNano(Some(arr.value(index))))
+            }
             _ => unimplemented!("data type {} not supported", array.data_type()),
         }
     }
@@ -257,6 +269,9 @@ impl TryFrom<&DataType> for ScalarValue {
             DataType::Float64 => Ok(ScalarValue::Float64(None)),
             DataType::Utf8 => Ok(ScalarValue::Utf8(None)),
             DataType::LargeUtf8 => Ok(ScalarValue::Utf8(None)),
+            DataType::Decimal128(p, s) => Ok(ScalarValue::Decimal128(None, *p, *s)),
+            DataType::Decimal256(p, s) => Ok(ScalarValue::Decimal256(None, *p, *s)),
+            DataType::Interval(IntervalUnit::MonthDayNano) => Ok(ScalarValue::IntervalMonthDayNano(None)),
             _ => unimplemented!("data type {} not supported", value),
         }
     }

qurious/src/execution/session.rs

@@ -381,23 +381,46 @@ mod tests {
         // session.sql("INSERT INTO test VALUES (1, 1), (2, 2), (3, 3), (3, 5), (NULL, NULL);")?;
         // session.sql("select a, b, c, d from x join y on a = c")?;
         println!("++++++++++++++");
-        let batch = session
-            .sql(
-                "
+        // Debug helpers for TPC-H Q8 while this test is ignored:
+        // validate we actually have BRAZIL rows in the derived `all_nations` subquery.
+        let debug = session.sql(
+            "
+
 select
-    sum(l_extendedprice * l_discount) as revenue
+    c_custkey,
+    c_name,
+    sum(l_extendedprice * (1 - l_discount)) as revenue,
+    c_acctbal,
+    n_name,
+    c_address,
+    c_phone,
+    c_comment
 from
-    lineitem
+    customer,
+    orders,
+    lineitem,
+    nation
 where
-        l_shipdate >= date '1994-01-01'
-  and l_shipdate < date '1995-01-01'
-  and l_discount between 0.06 - 0.01 and 0.06 + 0.01
-  and l_quantity < 24;
+        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
+group by
+    c_custkey,
+    c_name,
+    c_acctbal,
+    c_phone,
+    n_name,
+    c_address,
+    c_comment
+order by
+    revenue desc
+limit 10;
 ",
-            )
-            .unwrap();
-
-        print_batches(&batch)?;
+        )?;
+        print_batches(&debug)?;
 
         Ok(())
     }
@@ -422,7 +445,17 @@ where
         session.register_table("t", Arc::new(datasource))?;
 
         let batch = session.sql("SELECT a.* FROM t as a")?;
-        assert_eq!(data, batch);
+        // Compare data ignoring schema-level metadata (we attach qualifiers in schema metadata
+        // to support disambiguation across aliased relations).
+        assert_eq!(data.len(), batch.len());
+        for (expected, actual) in data.iter().zip(batch.iter()) {
+            assert_eq!(expected.schema().fields(), actual.schema().fields());
+            assert_eq!(expected.num_rows(), actual.num_rows());
+            assert_eq!(expected.num_columns(), actual.num_columns());
+            for i in 0..expected.num_columns() {
+                assert_eq!(expected.column(i).as_ref(), actual.column(i).as_ref());
+            }
+        }
 
         Ok(())
     }

qurious/src/functions/datetime/extract.rs

@@ -3,6 +3,7 @@ use std::str::FromStr;
 use arrow::array::Array;
 use arrow::compute::kernels::cast_utils::IntervalUnit;
 use arrow::{
+    array::new_empty_array,
     array::{ArrayRef, AsArray},
     compute::{cast, date_part, DatePart},
     datatypes::DataType,
@@ -31,16 +32,34 @@ impl UserDefinedFunction for DatetimeExtract {
             return Err(Error::InvalidArgumentError("EXTRACT requires 2 arguments".to_string()));
         }
 
+        // EXTRACT is often used in projections that can be evaluated on empty batches (e.g. after filters).
+        // In that case, literal args can be represented as empty arrays, so we must not assume len > 0.
+        if args[1].len() == 0 {
+            return Ok(new_empty_array(&DataType::Int64));
+        }
+
         // get interval_unit value
         let interval_unit = if let Some(val) = args.get(0) {
-            val.as_string::<i32>().value(0)
+            let arr = val.as_string::<i32>();
+            if arr.len() == 0 {
+                // no rows => empty output already handled above; keep defensive
+                return Ok(new_empty_array(&DataType::Int64));
+            }
+            if !arr.is_valid(0) {
+                return Err(Error::InvalidArgumentError(
+                    "First argument of `EXTRACT` must be non-null scalar Utf8".to_string(),
+                ));
+            }
+            arr.value(0)
         } else {
             return Err(Error::InvalidArgumentError(
                 "First argument of `DATE_PART` must be non-null scalar Utf8".to_string(),
             ));
         };
 
-        match IntervalUnit::from_str(interval_unit)? {
+        // Normalize case so we accept YEAR / year / Year, etc.
+        let interval_unit = interval_unit.to_ascii_lowercase();
+        match IntervalUnit::from_str(&interval_unit)? {
             IntervalUnit::Year => date_part_f64(args[1].as_ref(), DatePart::Year),
             IntervalUnit::Month => date_part_f64(args[1].as_ref(), DatePart::Month),
             IntervalUnit::Week => date_part_f64(args[1].as_ref(), DatePart::Week),
@@ -60,3 +79,22 @@ impl UserDefinedFunction for DatetimeExtract {
 fn date_part_f64(array: &dyn Array, part: DatePart) -> Result<ArrayRef> {
     cast(date_part(array, part)?.as_ref(), &DataType::Int64).map_err(|e| arrow_err!(e))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{new_empty_array, StringArray};
+    use arrow::datatypes::DataType;
+    use std::sync::Arc;
+
+    #[test]
+    fn test_extract_empty_input_does_not_panic() {
+        let udf = DatetimeExtract;
+        let unit = Arc::new(StringArray::from(vec!["YEAR"])) as ArrayRef;
+        let empty_date = new_empty_array(&DataType::Date32);
+
+        let out = udf.eval(vec![unit, empty_date]).unwrap();
+        assert_eq!(out.len(), 0);
+        assert_eq!(out.data_type(), &DataType::Int64);
+    }
+}

qurious/src/logical/expr/case.rs

@@ -0,0 +1,68 @@
+use std::fmt::{Display, Formatter};
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, FieldRef, Schema};
+
+use crate::datatypes::scalar::ScalarValue;
+use crate::error::Result;
+use crate::logical::plan::LogicalPlan;
+
+use super::LogicalExpr;
+
+/// SQL CASE expression.
+///
+/// Supports both:
+/// - searched CASE: `CASE WHEN <cond> THEN <value> ... ELSE <value> END`
+/// - simple CASE: `CASE <operand> WHEN <expr> THEN <value> ... ELSE <value> END`
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct CaseExpr {
+    pub operand: Option<Box<LogicalExpr>>,
+    pub when_then: Vec<(LogicalExpr, LogicalExpr)>,
+    pub else_expr: Box<LogicalExpr>,
+}
+
+impl CaseExpr {
+    pub fn field(&self, plan: &LogicalPlan) -> Result<FieldRef> {
+        // Best-effort: use first THEN type (or ELSE) as the output type.
+        // The optimizer's type coercion rule will cast branches to a common type later.
+        let schema = plan.schema();
+        let dt = self.data_type(&schema)?;
+        Ok(Arc::new(Field::new(format!("{}", self), dt, true)))
+    }
+
+    pub fn data_type(&self, schema: &Arc<Schema>) -> Result<DataType> {
+        for (_, then_expr) in &self.when_then {
+            let dt = then_expr.data_type(schema)?;
+            if dt != DataType::Null {
+                return Ok(dt);
+            }
+        }
+        self.else_expr.data_type(schema)
+    }
+}
+
+impl Display for CaseExpr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "CASE")?;
+        if let Some(op) = &self.operand {
+            write!(f, " {op}")?;
+        }
+        for (w, t) in &self.when_then {
+            write!(f, " WHEN {w} THEN {t}")?;
+        }
+        // Always print ELSE for determinism (planner fills missing ELSE with NULL).
+        write!(f, " ELSE {} END", self.else_expr)
+    }
+}
+
+impl From<ScalarValue> for CaseExpr {
+    fn from(value: ScalarValue) -> Self {
+        CaseExpr {
+            operand: None,
+            when_then: vec![],
+            else_expr: Box::new(LogicalExpr::Literal(value)),
+        }
+    }
+}
+
+