feat(query): support extract EPOCH from timestamp (#17385)

* feat(query): support extract EPOCH from timestamp * use val directly * add domain * function rename to epoch
databendlabs · Feb 1, 2025 · e76aabc · e76aabc
1 parent c2e462e
commit e76aabc
Show file tree

Hide file tree

Showing 7 changed files with 132 additions and 51 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
@@ -1524,6 +1524,17 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
         }),
     );
 
+    registry.register_passthrough_nullable_1_arg::<TimestampType, Float64Type, _, _>(
+        "epoch",
+        |_, domain| {
+            FunctionDomain::Domain(SimpleDomain::<F64> {
+                min: (domain.min as f64 / 1_000_000f64).into(),
+                max: (domain.max as f64 / 1_000_000f64).into(),
+            })
+        },
+        vectorize_1_arg::<TimestampType, Float64Type>(|val, _| (val as f64 / 1_000_000f64).into()),
+    );
+
     registry.register_passthrough_nullable_1_arg::<TimestampType, UInt8Type, _, _>(
         "to_hour",
         |_, _| FunctionDomain::Full,

diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt
@@ -1671,8 +1671,10 @@ Functions overloads:
 199 divide(Float64, Float64) :: Float64
 200 divide(Float64 NULL, Float64 NULL) :: Float64 NULL
 0 divnull(Float64 NULL, Float64 NULL) :: Float64 NULL
-0 epoch(Int64) :: Interval
-1 epoch(Int64 NULL) :: Interval NULL
+0 epoch(Timestamp) :: Float64
+1 epoch(Timestamp NULL) :: Float64 NULL
+2 epoch(Int64) :: Interval
+3 epoch(Int64 NULL) :: Interval NULL
 0 eq(Variant, Variant) :: Boolean
 1 eq(Variant NULL, Variant NULL) :: Boolean NULL
 2 eq(String, String) :: Boolean

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,7 @@ 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", 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,40 @@ select EXTRACT(YEAR FROM to_datetime('2022-03-04 22:32:09')) = 2022
 ----
 1
 
+query T
+select EXTRACT(EPOCH FROM to_datetime('1969-12-31 23:59:59'))
+----
+-1.0
+
+statement ok
+create or replace table t_epoch(c timestamp);
+
+statement ok
+insert into t_epoch values('1969-12-31 23:59:59'), ('1970-01-01 00:00:01.123456789');
+
+query T
+select EXTRACT(EPOCH FROM c) from t_epoch order by c;
+----
+-1.0
+1.123456
+
+statement ok
+drop table if exists t_epoch;
 
+query T
+select EXTRACT(EPOCH FROM to_datetime('1969-12-31 23:59:59.999999'))
+----
+-1.0e-6
+
+query B
+select EXTRACT(EPOCH FROM now()) = epoch(now())
+----
+1
+
+query B
+SELECT (to_timestamp('2023-11-12 09:38:18.965575') - to_timestamp(0)) / 1000000.0 = epoch('2023-11-12 09:38:18.965575')
+----
+1
 
 query B
 select EXTRACT(MONTH FROM to_datetime('2022-03-04 22:32:09')) = 3