feat(query): support extract EPOCH from timestamp

databendlabs · Jan 29, 2025 · b23dd38 · b23dd38
1 parent c5e7355
commit b23dd38
Show file tree

Hide file tree

Showing 6 changed files with 112 additions and 49 deletions.
diff --git a/src/query/ast/src/ast/expr.rs b/src/query/ast/src/ast/expr.rs
@@ -828,6 +828,7 @@ pub enum IntervalKind {
     Doy,
     Week,
     Dow,
+    Epoch,
 }
 
 impl Display for IntervalKind {
@@ -843,6 +844,7 @@ impl Display for IntervalKind {
             IntervalKind::Doy => "DOY",
             IntervalKind::Dow => "DOW",
             IntervalKind::Week => "WEEK",
+            IntervalKind::Epoch => "EPOCH",
         })
     }
 }

diff --git a/src/query/ast/src/parser/expr.rs b/src/query/ast/src/parser/expr.rs
@@ -1821,56 +1821,87 @@ pub fn weekday(i: Input) -> IResult<Weekday> {
 }
 
 pub fn interval_kind(i: Input) -> IResult<IntervalKind> {
+    let year = value(IntervalKind::Year, rule! { YEAR });
+    let quarter = value(IntervalKind::Quarter, rule! { QUARTER });
+    let month = value(IntervalKind::Month, rule! { MONTH });
+    let day = value(IntervalKind::Day, rule! { DAY });
+    let hour = value(IntervalKind::Hour, rule! { HOUR });
+    let minute = value(IntervalKind::Minute, rule! { MINUTE });
+    let second = value(IntervalKind::Second, rule! { SECOND });
+    let doy = value(IntervalKind::Doy, rule! { DOY });
+    let dow = value(IntervalKind::Dow, rule! { DOW });
+    let week = value(IntervalKind::Week, rule! { WEEK });
+    let epoch = value(IntervalKind::Epoch, rule! { EPOCH });
+    let year_str = value(
+        IntervalKind::Year,
+        rule! { #literal_string_eq_ignore_case("YEAR")  },
+    );
+    let quarter_str = value(
+        IntervalKind::Quarter,
+        rule! { #literal_string_eq_ignore_case("QUARTER") },
+    );
+    let month_str = value(
+        IntervalKind::Month,
+        rule! { #literal_string_eq_ignore_case("MONTH")  },
+    );
+    let day_str = value(
+        IntervalKind::Day,
+        rule! { #literal_string_eq_ignore_case("DAY")  },
+    );
+    let hour_str = value(
+        IntervalKind::Hour,
+        rule! { #literal_string_eq_ignore_case("HOUR")  },
+    );
+    let minute_str = value(
+        IntervalKind::Minute,
+        rule! { #literal_string_eq_ignore_case("MINUTE")  },
+    );
+    let second_str = value(
+        IntervalKind::Second,
+        rule! { #literal_string_eq_ignore_case("SECOND")  },
+    );
+    let doy_str = value(
+        IntervalKind::Doy,
+        rule! { #literal_string_eq_ignore_case("DOY")  },
+    );
+    let dow_str = value(
+        IntervalKind::Dow,
+        rule! { #literal_string_eq_ignore_case("DOW")  },
+    );
+    let week_str = value(
+        IntervalKind::Week,
+        rule! { #literal_string_eq_ignore_case("WEEK")  },
+    );
+    let epoch_str = value(
+        IntervalKind::Epoch,
+        rule! { #literal_string_eq_ignore_case("EPOCH")  },
+    );
     alt((
-        value(IntervalKind::Year, rule! { YEAR }),
-        value(IntervalKind::Quarter, rule! { QUARTER }),
-        value(IntervalKind::Month, rule! { MONTH }),
-        value(IntervalKind::Day, rule! { DAY }),
-        value(IntervalKind::Hour, rule! { HOUR }),
-        value(IntervalKind::Minute, rule! { MINUTE }),
-        value(IntervalKind::Second, rule! { SECOND }),
-        value(IntervalKind::Doy, rule! { DOY }),
-        value(IntervalKind::Dow, rule! { DOW }),
-        value(IntervalKind::Week, rule! { WEEK }),
-        value(
-            IntervalKind::Year,
-            rule! { #literal_string_eq_ignore_case("YEAR")  },
+        rule!(
+            #year
+            | #quarter
+            | #month
+            | #day
+            | #hour
+            | #minute
+            | #second
+            | #doy
+            | #dow
+            | #week
+            | #epoch
         ),
-        value(
-            IntervalKind::Quarter,
-            rule! { #literal_string_eq_ignore_case("QUARTER") },
-        ),
-        value(
-            IntervalKind::Month,
-            rule! { #literal_string_eq_ignore_case("MONTH")  },
-        ),
-        value(
-            IntervalKind::Day,
-            rule! { #literal_string_eq_ignore_case("DAY")  },
-        ),
-        value(
-            IntervalKind::Hour,
-            rule! { #literal_string_eq_ignore_case("HOUR")  },
-        ),
-        value(
-            IntervalKind::Minute,
-            rule! { #literal_string_eq_ignore_case("MINUTE")  },
-        ),
-        value(
-            IntervalKind::Second,
-            rule! { #literal_string_eq_ignore_case("SECOND")  },
-        ),
-        value(
-            IntervalKind::Doy,
-            rule! { #literal_string_eq_ignore_case("DOY")  },
-        ),
-        value(
-            IntervalKind::Dow,
-            rule! { #literal_string_eq_ignore_case("DOW")  },
-        ),
-        value(
-            IntervalKind::Week,
-            rule! { #literal_string_eq_ignore_case("WEEK")  },
+        rule!(
+            #year_str
+            | #quarter_str
+            | #month_str
+            | #day_str
+            | #hour_str
+            | #minute_str
+            | #second_str
+            | #doy_str
+            | #dow_str
+            | #week_str
+            | #epoch_str
         ),
     ))(i)
 }

diff --git a/src/query/functions/src/scalars/timestamp/src/datetime.rs b/src/query/functions/src/scalars/timestamp/src/datetime.rs
@@ -1516,6 +1516,19 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
         }),
     );
 
+    registry.register_passthrough_nullable_1_arg::<TimestampType, Float64Type, _, _>(
+        "epoch_microsecond",
+        |_, _| FunctionDomain::Full,
+        vectorize_1_arg::<TimestampType, Float64Type>(|val, ctx| {
+            (val.to_timestamp(ctx.func_ctx.tz.clone())
+                .with_time_zone(TimeZone::UTC)
+                .timestamp()
+                .as_microsecond() as f64
+                / 1_000_000f64)
+                .into()
+        }),
+    );
+
     registry.register_passthrough_nullable_1_arg::<TimestampType, UInt8Type, _, _>(
         "to_hour",
         |_, _| FunctionDomain::Full,

diff --git a/src/query/sql/src/planner/semantic/type_check.rs b/src/query/sql/src/planner/semantic/type_check.rs
@@ -2930,6 +2930,9 @@ impl<'a> TypeChecker<'a> {
             ASTIntervalKind::Doy => self.resolve_function(span, "to_day_of_year", vec![], &[arg]),
             ASTIntervalKind::Dow => self.resolve_function(span, "to_day_of_week", vec![], &[arg]),
             ASTIntervalKind::Week => self.resolve_function(span, "to_week_of_year", vec![], &[arg]),
+            ASTIntervalKind::Epoch => {
+                self.resolve_function(span, "epoch_microsecond", vec![], &[arg])
+            }
         }
     }
 

diff --git a/src/tests/sqlsmith/src/sql_gen/expr.rs b/src/tests/sqlsmith/src/sql_gen/expr.rs
@@ -495,7 +495,7 @@ impl<R: Rng> SqlGenerator<'_, R> {
                         DataType::Timestamp
                     };
                     let expr = self.gen_expr(&expr_ty);
-                    let kind = match self.rng.gen_range(0..=9) {
+                    let kind = match self.rng.gen_range(0..=10) {
                         0 => IntervalKind::Year,
                         1 => IntervalKind::Quarter,
                         2 => IntervalKind::Month,
@@ -506,6 +506,7 @@ impl<R: Rng> SqlGenerator<'_, R> {
                         7 => IntervalKind::Doy,
                         8 => IntervalKind::Dow,
                         9 => IntervalKind::Week,
+                        10 => IntervalKind::Epoch,
                         _ => unreachable!(),
                     };
                     Expr::Extract {

diff --git a/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test b/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test
@@ -1124,7 +1124,20 @@ select EXTRACT(YEAR FROM to_datetime('2022-03-04 22:32:09')) = 2022
 ----
 1
 
+query B
+select EXTRACT(EPOCH FROM to_datetime('2022-03-04 22:32:09')) = to_unix_timestamp('2022-03-04 22:32:09');
+----
+1
+
+query B
+select EXTRACT(EPOCH FROM now()) = epoch_microsecond(now())
+----
+1
 
+query B
+SELECT (to_timestamp('2023-11-12 09:38:18.965575') - to_timestamp(0)) / 1000000.0 = epoch_microsecond('2023-11-12 09:38:18.965575')
+----
+1
 
 query B
 select EXTRACT(MONTH FROM to_datetime('2022-03-04 22:32:09')) = 3