feat(query): DATE_ADD Functions to support week as a unit (#16530)

databendlabs · Sep 27, 2024 · 6f10276 · 6f10276
1 parent 27157a3
commit 6f10276
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/src/query/functions/src/scalars/datetime.rs b/src/query/functions/src/scalars/datetime.rs
@@ -916,6 +916,39 @@ macro_rules! impl_register_arith_functions {
                 ),
             );
 
+            registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
+                concat!($op, "_weeks"),
+
+                |_, _, _| FunctionDomain::MayThrow,
+                vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
+                    let delta = 7 * delta;
+                    match AddDaysImpl::eval_date(date, $signed_wrapper!{delta}) {
+                        Ok(t) => builder.push(t),
+                        Err(e) => {
+                            ctx.set_error(builder.len(), e);
+                            builder.push(0);
+                        },
+                    }
+                }),
+            );
+            registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
+                concat!($op, "_weeks"),
+
+                |_, _, _| FunctionDomain::MayThrow,
+                vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
+                    |ts, delta, builder, ctx| {
+                        let delta = 7 * delta;
+                        match AddDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta}) {
+                            Ok(t) => builder.push(t),
+                            Err(e) => {
+                                ctx.set_error(builder.len(), e);
+                                builder.push(0);
+                            },
+                        }
+                    },
+                ),
+            );
+
             registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, TimestampType, _, _>(
                 concat!($op, "_hours"),
 

diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt
@@ -108,6 +108,10 @@ Functions overloads:
 1 add_seconds(Date NULL, Int64 NULL) :: Timestamp NULL
 2 add_seconds(Timestamp, Int64) :: Timestamp
 3 add_seconds(Timestamp NULL, Int64 NULL) :: Timestamp NULL
+0 add_weeks(Date, Int64) :: Date
+1 add_weeks(Date NULL, Int64 NULL) :: Date NULL
+2 add_weeks(Timestamp, Int64) :: Timestamp
+3 add_weeks(Timestamp NULL, Int64 NULL) :: Timestamp NULL
 0 add_years(Date, Int64) :: Date
 1 add_years(Date NULL, Int64 NULL) :: Date NULL
 2 add_years(Timestamp, Int64) :: Timestamp
@@ -3686,6 +3690,10 @@ Functions overloads:
 1 subtract_seconds(Date NULL, Int64 NULL) :: Timestamp NULL
 2 subtract_seconds(Timestamp, Int64) :: Timestamp
 3 subtract_seconds(Timestamp NULL, Int64 NULL) :: Timestamp NULL
+0 subtract_weeks(Date, Int64) :: Date
+1 subtract_weeks(Date NULL, Int64 NULL) :: Date NULL
+2 subtract_weeks(Timestamp, Int64) :: Timestamp
+3 subtract_weeks(Timestamp NULL, Int64 NULL) :: Timestamp NULL
 0 subtract_years(Date, Int64) :: Date
 1 subtract_years(Date NULL, Int64 NULL) :: Date NULL
 2 subtract_years(Timestamp, Int64) :: Timestamp

diff --git a/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test b/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test
@@ -1132,3 +1132,8 @@ query T
 select to_timestamp('2022-03-27 07:54:31.12');
 ----
 2022-03-27 07:54:31.120000
+
+query T
+SELECT DATE_ADD('week', 1, TODAY())=DATE_ADD('day', 7, TODAY())
+----
+1