feat(query): add date_diff function (#16560)

feat: add date_diff function * support date and timestamp
databendlabs · Oct 12, 2024 · 6ad9da2 · 6ad9da2
1 parent e7fcde8
commit 6ad9da2
Show file tree

Hide file tree

Showing 17 changed files with 1,065 additions and 252 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/common/parquet2/src/encoding/delta_length_byte_array/encoder.rs b/src/common/parquet2/src/encoding/delta_length_byte_array/encoder.rs
@@ -15,7 +15,7 @@
 
 use crate::encoding::delta_bitpacked;
 
-/// Encodes a clonable iterator of `&[u8]` into `buffer`. This does not allocated on the heap.
+/// Encodes a cloneable iterator of `&[u8]` into `buffer`. This does not allocated on the heap.
 /// # Implementation
 /// This encoding is equivalent to call [`delta_bitpacked::encode`] on the lengths of the items
 /// of the iterator followed by extending the buffer from each item of the iterator.

diff --git a/src/query/ast/src/ast/expr.rs b/src/query/ast/src/ast/expr.rs
@@ -212,6 +212,12 @@ pub enum Expr {
         interval: Box<Expr>,
         date: Box<Expr>,
     },
+    DateDiff {
+        span: Span,
+        unit: IntervalKind,
+        date_start: Box<Expr>,
+        date_end: Box<Expr>,
+    },
     DateSub {
         span: Span,
         unit: IntervalKind,
@@ -260,6 +266,7 @@ impl Expr {
             | Expr::Map { span, .. }
             | Expr::Interval { span, .. }
             | Expr::DateAdd { span, .. }
+            | Expr::DateDiff { span, .. }
             | Expr::DateSub { span, .. }
             | Expr::DateTrunc { span, .. }
             | Expr::Hole { span, .. } => *span,
@@ -388,6 +395,15 @@ impl Expr {
                 date,
                 ..
             } => merge_span(merge_span(*span, interval.whole_span()), date.whole_span()),
+            Expr::DateDiff {
+                span,
+                date_start,
+                date_end,
+                ..
+            } => merge_span(
+                merge_span(*span, date_start.whole_span()),
+                date_end.whole_span(),
+            ),
             Expr::DateSub {
                 span,
                 interval,
@@ -409,6 +425,7 @@ impl Expr {
             "SUBSTRING",
             "TRIM",
             "DATE_ADD",
+            "DATE_DIFF",
             "DATE_SUB",
             "DATE_TRUNC",
         ]
@@ -698,6 +715,14 @@ impl Display for Expr {
                 } => {
                     write!(f, "DATE_ADD({unit}, {interval}, {date})")?;
                 }
+                Expr::DateDiff {
+                    unit,
+                    date_start,
+                    date_end,
+                    ..
+                } => {
+                    write!(f, "DATE_DIFF({unit}, {date_start}, {date_end})")?;
+                }
                 Expr::DateSub {
                     unit,
                     interval,

diff --git a/src/query/ast/src/ast/format/syntax/expr.rs b/src/query/ast/src/ast/format/syntax/expr.rs
@@ -404,6 +404,20 @@ pub(crate) fn pretty_expr(expr: Expr) -> RcDoc<'static> {
             .append(RcDoc::space())
             .append(pretty_expr(*date))
             .append(RcDoc::text(")")),
+        Expr::DateDiff {
+            unit,
+            date_start,
+            date_end,
+            ..
+        } => RcDoc::text("DATE_DIFF(")
+            .append(RcDoc::text(unit.to_string()))
+            .append(RcDoc::text(","))
+            .append(RcDoc::space())
+            .append(pretty_expr(*date_start))
+            .append(RcDoc::text(","))
+            .append(RcDoc::space())
+            .append(pretty_expr(*date_end))
+            .append(RcDoc::text(")")),
         Expr::DateSub {
             unit,
             interval,

diff --git a/src/query/ast/src/parser/expr.rs b/src/query/ast/src/parser/expr.rs
@@ -296,6 +296,11 @@ pub enum ExprElement {
         interval: Expr,
         date: Expr,
     },
+    DateDiff {
+        unit: IntervalKind,
+        date_start: Expr,
+        date_end: Expr,
+    },
     DateSub {
         unit: IntervalKind,
         interval: Expr,
@@ -408,6 +413,7 @@ impl ExprElement {
             ExprElement::Map { .. } => Affix::Nilfix,
             ExprElement::Interval { .. } => Affix::Nilfix,
             ExprElement::DateAdd { .. } => Affix::Nilfix,
+            ExprElement::DateDiff { .. } => Affix::Nilfix,
             ExprElement::DateSub { .. } => Affix::Nilfix,
             ExprElement::DateTrunc { .. } => Affix::Nilfix,
             ExprElement::Hole { .. } => Affix::Nilfix,
@@ -450,6 +456,7 @@ impl Expr {
             Expr::Map { .. } => Affix::Nilfix,
             Expr::Interval { .. } => Affix::Nilfix,
             Expr::DateAdd { .. } => Affix::Nilfix,
+            Expr::DateDiff { .. } => Affix::Nilfix,
             Expr::DateSub { .. } => Affix::Nilfix,
             Expr::DateTrunc { .. } => Affix::Nilfix,
             Expr::Hole { .. } => Affix::Nilfix,
@@ -625,6 +632,16 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
                 interval: Box::new(interval),
                 date: Box::new(date),
             },
+            ExprElement::DateDiff {
+                unit,
+                date_start,
+                date_end,
+            } => Expr::DateDiff {
+                span: transform_span(elem.span.tokens),
+                unit,
+                date_start: Box::new(date_start),
+                date_end: Box::new(date_end),
+            },
             ExprElement::DateSub {
                 unit,
                 interval,
@@ -1167,6 +1184,18 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             date,
         },
     );
+
+    let date_diff = map(
+        rule! {
+            DATE_DIFF ~ "(" ~ #interval_kind ~ "," ~ #subexpr(0) ~ "," ~ #subexpr(0) ~ ")"
+        },
+        |(_, _, unit, _, date_start, _, date_end, _)| ExprElement::DateDiff {
+            unit,
+            date_start,
+            date_end,
+        },
+    );
+
     let date_sub = map(
         rule! {
             DATE_SUB ~ "(" ~ #interval_kind ~ "," ~ #subexpr(0) ~ "," ~ #subexpr(0) ~ ")"
@@ -1253,6 +1282,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
                 | #unary_op : "<operator>"
                 | #cast : "`CAST(... AS ...)`"
                 | #date_add: "`DATE_ADD(..., ..., (YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND | DOY | DOW))`"
+                | #date_diff: "`DATE_DIFF(..., ..., (YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND | DOY | DOW))`"
                 | #date_sub: "`DATE_SUB(..., ..., (YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND | DOY | DOW))`"
                 | #date_trunc: "`DATE_TRUNC((YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND), ...)`"
                 | #date_expr: "`DATE <str_literal>`"

diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs
@@ -478,6 +478,8 @@ pub enum TokenKind {
     DATE,
     #[token("DATE_ADD", ignore(ascii_case))]
     DATE_ADD,
+    #[token("DATE_DIFF", ignore(ascii_case))]
+    DATE_DIFF,
     #[token("DATE_PART", ignore(ascii_case))]
     DATE_PART,
     #[token("DATE_SUB", ignore(ascii_case))]
@@ -1567,6 +1569,7 @@ impl TokenKind {
             // | TokenKind::WINDOW
             | TokenKind::WITH
             | TokenKind::DATE_ADD
+            | TokenKind::DATE_DIFF
             | TokenKind::DATE_SUB
             | TokenKind::DATE_TRUNC
             | TokenKind::IGNORE_RESULT

diff --git a/src/query/ast/tests/it/testdata/expr-error.txt b/src/query/ast/tests/it/testdata/expr-error.txt
@@ -52,7 +52,7 @@ error:
   --> SQL:1:10
   |
 1 | CAST(col1)
-  | ----     ^ unexpected `)`, expecting `AS`, `,`, `(`, `IS`, `NOT`, `IN`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `LIKE`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `DATE_ADD`, `DATE_SUB`, `DATE_TRUNC`, `DATE`, or 32 more ...
+  | ----     ^ unexpected `)`, expecting `AS`, `,`, `(`, `IS`, `NOT`, `IN`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `LIKE`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `DATE_ADD`, `DATE_DIFF`, `DATE_SUB`, `DATE_TRUNC`, or 33 more ...
   | |         
   | while parsing `CAST(... AS ...)`
   | while parsing expression

diff --git a/src/query/ast/tests/it/testdata/stmt-error.txt b/src/query/ast/tests/it/testdata/stmt-error.txt
@@ -443,7 +443,7 @@ error:
   --> SQL:1:41
   |
 1 | SELECT * FROM t GROUP BY GROUPING SETS ()
-  | ------                                  ^ unexpected `)`, expecting `(`, `IS`, `IN`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `LIKE`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `DATE_ADD`, `DATE_SUB`, `DATE_TRUNC`, `DATE`, `TIMESTAMP`, `INTERVAL`, or 30 more ...
+  | ------                                  ^ unexpected `)`, expecting `(`, `IS`, `IN`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `LIKE`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `DATE_ADD`, `DATE_DIFF`, `DATE_SUB`, `DATE_TRUNC`, `DATE`, `TIMESTAMP`, or 31 more ...
   | |                                        
   | while parsing `SELECT ...`
 
@@ -865,7 +865,7 @@ error:
   --> SQL:1:65
   |
 1 | CREATE FUNCTION IF NOT EXISTS isnotempty AS(p) -> not(is_null(p)
-  | ------                                   --       ----          ^ unexpected end of input, expecting `)`, `IGNORE`, `RESPECT`, `OVER`, `(`, `IS`, `NOT`, `IN`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `LIKE`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `DATE_ADD`, `DATE_SUB`, or 35 more ...
+  | ------                                   --       ----          ^ unexpected end of input, expecting `)`, `IGNORE`, `RESPECT`, `OVER`, `(`, `IS`, `NOT`, `IN`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `LIKE`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `DATE_ADD`, `DATE_DIFF`, or 36 more ...
   | |                                        |        |  |          
   | |                                        |        |  while parsing `(<expr> [, ...])`
   | |                                        |        while parsing expression

diff --git a/src/query/expression/src/types/timestamp.rs b/src/query/expression/src/types/timestamp.rs
@@ -45,8 +45,8 @@ pub const TIMESTAMP_MIN: i64 = -30610224000000000;
 /// Maximum valid timestamp `9999-12-31 23:59:59.999999`, represented by the microsecs offset from 1970-01-01.
 pub const TIMESTAMP_MAX: i64 = 253402300799999999;
 
-pub const MICROS_IN_A_SEC: i64 = 1_000_000;
-pub const MICROS_IN_A_MILLI: i64 = 1_000;
+pub const MICROS_PER_SEC: i64 = 1_000_000;
+pub const MICROS_PER_MILLI: i64 = 1_000;
 
 pub const PRECISION_MICRO: u8 = 6;
 pub const PRECISION_MILLI: u8 = 3;
@@ -262,7 +262,7 @@ impl ArgType for TimestampType {
 }
 
 pub fn microseconds_to_seconds(micros: i64) -> i64 {
-    micros / MICROS_IN_A_SEC
+    micros / MICROS_PER_SEC
 }
 
 pub fn microseconds_to_days(micros: i64) -> i32 {