WIP

Author: longshan.lu

qurious/src/datatypes/scalar.rs

@@ -2,10 +2,10 @@ use crate::error::{Error, Result};
 use arrow::{
     array::{
         new_null_array, Array, ArrayRef, ArrowPrimitiveType, BooleanArray, Decimal128Array, Decimal256Array,
-        Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeStringArray, PrimitiveArray,
-        StringArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+        Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, IntervalMonthDayNanoArray,
+        LargeStringArray, PrimitiveArray, StringArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
     },
-    datatypes::{i256, DataType, Field},
+    datatypes::{i256, DataType, Field, IntervalMonthDayNano, IntervalUnit},
 };
 use std::any::type_name;
 use std::{fmt::Display, sync::Arc};
@@ -99,6 +99,10 @@ pub enum ScalarValue {
     Decimal128(Option<i128>, u8, i8),
     /// 256bit decimal, using the i256 to represent the decimal, precision scale
     Decimal256(Option<i256>, u8, i8),
+    /// A triple of the number of elapsed months, days, and nanoseconds.
+    /// Months and days are encoded as 32-bit signed integers.
+    /// Nanoseconds is encoded as a 64-bit signed integer (no leap seconds).
+    IntervalMonthDayNano(Option<IntervalMonthDayNano>),
     Utf8(Option<String>),
 }
 
@@ -130,6 +134,11 @@ impl ScalarValue {
             ScalarValue::Utf8(_) => Field::new("utf8", DataType::Utf8, true),
             ScalarValue::Decimal128(_, p, s) => Field::new("decimal128", DataType::Decimal128(*p, *s), true),
             ScalarValue::Decimal256(_, p, s) => Field::new("decimal256", DataType::Decimal256(*p, *s), true),
+            ScalarValue::IntervalMonthDayNano(_) => Field::new(
+                "interval_month_day_nano",
+                DataType::Interval(IntervalUnit::MonthDayNano),
+                true,
+            ),
         }
     }
 
@@ -150,6 +159,7 @@ impl ScalarValue {
             ScalarValue::Utf8(_) => DataType::Utf8,
             ScalarValue::Decimal128(_, p, s) => DataType::Decimal128(*p, *s),
             ScalarValue::Decimal256(_, p, s) => DataType::Decimal256(*p, *s),
+            ScalarValue::IntervalMonthDayNano(_) => DataType::Interval(IntervalUnit::MonthDayNano),
         }
     }
 
@@ -174,6 +184,9 @@ impl ScalarValue {
             ScalarValue::Decimal256(v, p, s) => {
                 Arc::new(build_decimal_array!(*v, Decimal256Array, num_row, *p, *s)) as ArrayRef
             }
+            ScalarValue::IntervalMonthDayNano(v) => {
+                Arc::new(IntervalMonthDayNanoArray::from(vec![*v; num_row])) as ArrayRef
+            }
         })
     }
 
@@ -275,6 +288,9 @@ impl Display for ScalarValue {
             ScalarValue::Decimal128(v, p, s) => format_decimal!(f, v, "Decimal128", p, s),
             ScalarValue::Decimal256(v, p, s) => format_decimal!(f, v, "Decimal256", p, s),
             ScalarValue::Utf8(v) => format_string!(f, v, "Utf8"),
+            ScalarValue::IntervalMonthDayNano(v) => {
+                format_string!(f, v.map(|v| format!("{:?}", v)), "IntervalMonthDayNano")
+            }
         }
     }
 }

qurious/src/execution/session.rs

@@ -98,6 +98,7 @@ impl ExecuteSession {
             LogicalPlan::Dml(stmt) => self.execute_dml(stmt),
             plan => {
                 let plan = self.optimizer.optimize(plan)?;
+                println!("plan: {}", crate::utils::format(&plan, 0));
                 self.planner.create_physical_plan(&plan)?.execute()
             }
         }
@@ -385,49 +386,26 @@ mod tests {
             .sql(
                 "
 select
-    s_acctbal,
-    s_name,
-    n_name,
-    p_partkey,
-    p_mfgr,
-    s_address,
-    s_phone,
-    s_comment
+    o_orderpriority,
+    count(*) as order_count
 from
-    part,
-    supplier,
-    partsupp,
-    nation,
-    region
+    orders
 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'
-)
+        o_orderdate >= '1993-07-01'
+  and o_orderdate < date '1993-07-01' + interval '3' month
+  and exists (
+        select
+            *
+        from
+            lineitem
+        where
+                l_orderkey = o_orderkey
+          and l_commitdate < l_receiptdate
+    )
+group by
+    o_orderpriority
 order by
-    s_acctbal desc,
-    n_name,
-    s_name,
-    p_partkey
-limit 10;
+    o_orderpriority;
 ",
             )
             .unwrap();

qurious/src/logical/expr/mod.rs

@@ -46,6 +46,7 @@ pub enum LogicalExpr {
     Like(Like),
     Negative(Box<LogicalExpr>),
     SubQuery(SubQuery),
+    Exists(Exists),
 }
 
 macro_rules! impl_logical_expr_methods {
@@ -83,6 +84,7 @@ impl_logical_expr_methods! {
 impl Display for LogicalExpr {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
+            LogicalExpr::Exists(Exists { subquery, negated }) => write!(f, "{} EXISTS ({})", if *negated { "NOT" } else { "" }, subquery),
             LogicalExpr::Negative(e) => write!(f, "- {}", e),
             LogicalExpr::Literal(v) => write!(f, "{}", v),
             LogicalExpr::Wildcard => write!(f, "*"),
@@ -213,6 +215,7 @@ impl LogicalExpr {
             LogicalExpr::SortExpr(SortExpr { expr, .. }) | LogicalExpr::Negative(expr) => expr.data_type(schema),
             LogicalExpr::Like(_) | LogicalExpr::IsNull(_) | LogicalExpr::IsNotNull(_) => Ok(DataType::Boolean),
             LogicalExpr::SubQuery(subquery) => Ok(subquery.subquery.schema().fields[0].data_type().clone()),
+            LogicalExpr::Exists(_) => Ok(DataType::Boolean),
             _ => internal_err!("[{}] has no data type", self),
         }
     }
@@ -233,6 +236,12 @@ impl LogicalExpr {
 impl TransformNode for LogicalExpr {
     fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>(self, mut f: F) -> Result<Transformed<Self>> {
         Ok(match self {
+            LogicalExpr::Exists(Exists { subquery, negated }) => subquery.map_exprs(f)?.update(|subquery| {
+                LogicalExpr::Exists(Exists {
+                    subquery: Box::new(subquery),
+                    negated,
+                })
+            }),
             LogicalExpr::Alias(Alias { expr, name }) => f(*expr)?.update(|expr| {
                 LogicalExpr::Alias(Alias {
                     expr: Box::new(expr),
@@ -312,7 +321,11 @@ impl TransformNode for LogicalExpr {
             | LogicalExpr::IsNull(expr)
             | LogicalExpr::IsNotNull(expr)
             | LogicalExpr::Alias(Alias { expr, .. }) => vec![expr.as_ref()],
-            LogicalExpr::SubQuery(_) | LogicalExpr::Wildcard | LogicalExpr::Column(_) | LogicalExpr::Literal(_) => {
+            LogicalExpr::Exists(_)
+            | LogicalExpr::SubQuery(_)
+            | LogicalExpr::Wildcard
+            | LogicalExpr::Column(_)
+            | LogicalExpr::Literal(_) => {
                 vec![]
             }
             LogicalExpr::Like(like) => vec![like.expr.as_ref(), like.pattern.as_ref()],
@@ -336,6 +349,12 @@ pub(crate) fn get_expr_value(expr: LogicalExpr) -> Result<i64> {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct Exists {
+    pub negated: bool,
+    pub subquery: Box<LogicalPlan>,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Like {
     pub negated: bool,

qurious/src/planner/sql.rs

@@ -1,9 +1,15 @@
 use std::{collections::HashMap, sync::Arc};
 
-use arrow::datatypes::{Field, Schema, TimeUnit};
-use sqlparser::ast::{
-    Assignment, BinaryOperator, CopyOption, CopySource, CopyTarget, Cte, Expression, From, FunctionArgument, Ident,
-    Literal, Order, Select, SelectItem, Statement,
+use arrow::{
+    compute::kernels::cast_utils::{parse_interval_month_day_nano_config, IntervalParseConfig, IntervalUnit},
+    datatypes::{Field, Schema, TimeUnit},
+};
+use sqlparser::{
+    ast::{
+        Assignment, BinaryOperator, CopyOption, CopySource, CopyTarget, Cte, Expression, From, FunctionArgument, Ident,
+        Literal, Order, Select, SelectItem, Statement,
+    },
+    datatype::IntervalFields,
 };
 
 use crate::{
@@ -1015,10 +1021,28 @@ impl<'a> SqlQueryPlanner<'a> {
                 expr: Box::new(self.sql_to_expr(*left)?),
                 pattern: Box::new(self.sql_to_expr(*right)?),
             })),
+            Expression::Exists { subquery, negated } => Ok(LogicalExpr::Exists(Exists {
+                negated,
+                subquery: Box::new(self.new_context_scope(|planner| planner.select_to_plan(*subquery))?),
+            })),
+            Expression::Interval { expr, field } => self.interval_to_expr(*expr, field),
             _ => todo!("sql_to_expr: {:?}", expr),
         }
     }
 
+    fn interval_to_expr(&mut self, expr: Expression, field: IntervalFields) -> Result<LogicalExpr> {
+        match expr {
+            Expression::BinaryOperator(binary_op) => self.parse_binary_op(binary_op),
+            val => {
+                let interval_val = format!("{} {}", val, field);
+                let config = IntervalParseConfig::new(IntervalUnit::Second);
+                let val = parse_interval_month_day_nano_config(&interval_val, config)?;
+
+                Ok(LogicalExpr::Literal(ScalarValue::IntervalMonthDayNano(Some(val))))
+            }
+        }
+    }
+
     fn handle_function(&self, name: &str, mut args: Vec<LogicalExpr>) -> Result<LogicalExpr> {
         if let Some(udf) = self.udfs.get(name.to_uppercase().as_str()) {
             return Ok(LogicalExpr::Function(Function {
@@ -1285,6 +1309,34 @@ mod tests {
 
     use super::SqlQueryPlanner;
 
+    #[test]
+    fn test_interval() {
+        quick_test(
+            "SELECT INTERVAL '1' YEAR;",
+            "Projection: (IntervalMonthDayNano('IntervalMonthDayNano { months: 12, days: 0, nanoseconds: 0 }'))\n  Empty Relation\n",
+        );
+
+        quick_test(
+            "SELECT INTERVAL '1' MONTH;",
+            "Projection: (IntervalMonthDayNano('IntervalMonthDayNano { months: 1, days: 0, nanoseconds: 0 }'))\n  Empty Relation\n",
+        );
+
+        quick_test(
+            "SELECT INTERVAL '1' DAY;",
+            "Projection: (IntervalMonthDayNano('IntervalMonthDayNano { months: 0, days: 1, nanoseconds: 0 }'))\n  Empty Relation\n",
+        );
+
+        quick_test("SELECT DATE '1993-07-01' + INTERVAL '3' month", "Projection: (CAST(Utf8('1993-07-01') AS Date32) + IntervalMonthDayNano('IntervalMonthDayNano { months: 3, days: 0, nanoseconds: 0 }'))\n  Empty Relation\n");
+    }
+
+    #[test]
+    fn test_exists() {
+        quick_test(
+            "SELECT * FROM person WHERE EXISTS (SELECT * FROM other_tbl WHERE name = first_name)",
+            "",
+        );
+    }
+
     #[test]
     fn test_outer_field_reference() {
         quick_test(

qurious/tests/tpch/q4.slt

@@ -0,0 +1,28 @@
+query TI
+select
+    o_orderpriority,
+    count(*) as order_count
+from
+    orders
+where
+        o_orderdate >= '1993-07-01'
+  and o_orderdate < date '1993-07-01' + interval '3' month
+  and exists (
+        select
+            *
+        from
+            lineitem
+        where
+                l_orderkey = o_orderkey
+          and l_commitdate < l_receiptdate
+    )
+group by
+    o_orderpriority
+order by
+    o_orderpriority;
+----
+1-URGENT 999
+2-HIGH 997
+3-MEDIUM 1031
+4-NOT SPECIFIED 989
+5-LOW 1077

sqlparser/src/ast.rs

@@ -1,4 +1,7 @@
-use crate::{datatype::DataType, error::Error};
+use crate::{
+    datatype::{DataType, IntervalFields},
+    error::Error,
+};
 use std::fmt::{Display, Formatter};
 
 #[derive(Clone, PartialEq, Debug)]
@@ -670,6 +673,17 @@ pub enum Expression {
         left: Box<Expression>,
         right: Box<Expression>,
     },
+    /// `INTERVAL <expr> <field>`
+    Interval {
+        expr: Box<Expression>,
+        field: IntervalFields,
+    },
+    /// An exists expression `[ NOT ] EXISTS(SELECT ...)`, used in expressions like
+    /// `WHERE [ NOT ] EXISTS (SELECT ...)`.
+    Exists {
+        subquery: Box<Select>,
+        negated: bool,
+    },
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -797,6 +811,10 @@ impl Display for Expression {
             Expression::Like { negated, left, right } => {
                 write!(f, "{} {} LIKE {}", left, if *negated { "NOT" } else { "" }, right)
             }
+            Expression::Interval { expr, field } => write!(f, "INTERVAL {} {}", expr, field),
+            Expression::Exists { subquery, negated } => {
+                write!(f, "{} EXISTS ({})", if *negated { "NOT" } else { "" }, subquery)
+            }
         }
     }
 }