fix(query): add_months need to respect last day of month

TCeason · TCeason · commit 995f1aafb3e8 · 2025-06-20T13:05:20.000+08:00
ADD_MONTHS returns slightly different results than DATEADD used with a MONTH component:

For both ADD_MONTHS and DATEADD, if the result month has fewer days than the original day, the result day of the month is the last day of the result month.

For ADD_MONTHS only, if the original day is the last day of the month, the result day of month will be the last day of the result month.

In main:

```sql
select add_month('2025-04-30',1)
----
2025-05-30
```

In PR:

```sql
select add_month('2025-04-30',1)
----
2025-05-31
```
diff --git a/src/query/expression/src/utils/date_helper.rs b/src/query/expression/src/utils/date_helper.rs
@@ -74,7 +74,13 @@ pub const FACTOR_MINUTE: i64 = 60;
 pub const FACTOR_SECOND: i64 = 1;
 const LAST_DAY_LUT: [i8; 13] = [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31];
 
-fn eval_years_base(year: i16, month: i8, day: i8, delta: i64) -> std::result::Result<Date, String> {
+fn eval_years_base(
+    year: i16,
+    month: i8,
+    day: i8,
+    delta: i64,
+    _add_months: bool,
+) -> std::result::Result<Date, String> {
     let new_year = year as i64 + delta;
     let mut new_day = day;
     if std::intrinsics::unlikely(month == 2 && day == 29) {
@@ -91,6 +97,7 @@ fn eval_months_base(
     month: i8,
     day: i8,
     delta: i64,
+    add_months: bool,
 ) -> std::result::Result<Date, String> {
     let total_months = (month as i64 + delta - 1) as i16;
     let mut new_year = year + (total_months / 12);
@@ -101,12 +108,17 @@ fn eval_months_base(
     }
 
     // Handle month last day overflow, "2020-2-29" + "1 year" should be "2021-2-28", or "1990-1-31" + "3 month" should be "1990-4-30".
-    let new_day = std::cmp::min::<u32>(
-        day as u32,
-        last_day_of_year_month(new_year, (new_month0 + 1) as i8) as u32,
-    );
+    // For ADD_MONTHS only, if the original day is the last day of the month, the result day of month will be the last day of the result month.
+    let new_month = (new_month0 + 1) as i8;
+    // Determine the correct day
+    let max_day = last_day_of_year_month(new_year, new_month);
+    let new_day = if add_months && day == last_day_of_year_month(year, month) {
+        max_day
+    } else {
+        day.min(max_day)
+    };
 
-    match Date::new(new_year, (new_month0 + 1) as i8, new_day as i8) {
+    match Date::new(new_year, (new_month0 + 1) as i8, new_day) {
         Ok(d) => Ok(d),
         Err(e) => Err(format!("Invalid date: {}", e)),
     }
@@ -131,9 +143,16 @@ macro_rules! impl_interval_year_month {
                 date: i32,
                 tz: TimeZone,
                 delta: impl AsPrimitive<i64>,
+                add_months: bool,
             ) -> std::result::Result<i32, String> {
                 let date = date.to_date(tz);
-                let new_date = $op(date.year(), date.month(), date.day(), delta.as_())?;
+                let new_date = $op(
+                    date.year(),
+                    date.month(),
+                    date.day(),
+                    delta.as_(),
+                    add_months,
+                )?;
 
                 Ok(clamp_date(
                     new_date
@@ -147,9 +166,10 @@ macro_rules! impl_interval_year_month {
                 us: i64,
                 tz: TimeZone,
                 delta: impl AsPrimitive<i64>,
+                add_months: bool,
             ) -> std::result::Result<i64, String> {
                 let ts = us.to_timestamp(tz.clone());
-                let new_date = $op(ts.year(), ts.month(), ts.day(), delta.as_())?;
+                let new_date = $op(ts.year(), ts.month(), ts.day(), delta.as_(), add_months)?;
 
                 let mut ts = new_date
                     .at(ts.hour(), ts.minute(), ts.second(), ts.subsec_nanosecond())
diff --git a/src/query/functions/src/scalars/timestamp/src/datetime.rs b/src/query/functions/src/scalars/timestamp/src/datetime.rs
@@ -912,7 +912,7 @@ macro_rules! impl_register_arith_functions {
 
                 |_, _, _| FunctionDomain::MayThrow,
                 vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
-                    match EvalYearsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
+                    match EvalYearsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
                         Ok(t) => builder.push(t),
                         Err(e) => {
                             ctx.set_error(builder.len(), e);
@@ -927,7 +927,7 @@ macro_rules! impl_register_arith_functions {
                 |_, _, _| FunctionDomain::MayThrow,
                 vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
                     |ts, delta, builder, ctx| {
-                        match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
+                        match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
                             Ok(t) => builder.push(t),
                             Err(e) => {
                                 ctx.set_error(builder.len(), e);
@@ -943,7 +943,7 @@ macro_rules! impl_register_arith_functions {
 
                 |_, _, _| FunctionDomain::MayThrow,
                 vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
-                    match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3) {
+                    match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3, false) {
                         Ok(t) => builder.push(t),
                         Err(e) => {
                             ctx.set_error(builder.len(), e);
@@ -958,7 +958,7 @@ macro_rules! impl_register_arith_functions {
                 |_, _, _| FunctionDomain::MayThrow,
                 vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
                     |ts, delta, builder, ctx| {
-                        match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3) {
+                        match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3, false) {
                             Ok(t) => builder.push(t),
                             Err(e) => {
                                 ctx.set_error(builder.len(), e);
@@ -969,12 +969,45 @@ macro_rules! impl_register_arith_functions {
                 ),
             );
 
+            registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
+                concat!("date_", $op, "_months"),
+
+                |_, _, _| FunctionDomain::MayThrow,
+                vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
+                    match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
+                        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!("date_", $op, "_months"),
+
+                |_, _, _| FunctionDomain::MayThrow,
+                vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
+                    |ts, delta, builder, ctx| {
+                        match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
+                            Ok(t) => builder.push(t),
+                            Err(e) => {
+                                ctx.set_error(builder.len(), e);
+                                builder.push(0);
+                            },
+                        }
+                    },
+                ),
+            );
+
+            // For both ADD_MONTHS and DATEADD, if the result month has fewer days than the original day, the result day of the month is the last day of the result month.
+            // For ADD_MONTHS only, if the original day is the last day of the month, the result day of month will be the last day of the result month.
             registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
                 concat!($op, "_months"),
 
                 |_, _, _| FunctionDomain::MayThrow,
                 vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
-                    match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
+                    match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, true) {
                         Ok(t) => builder.push(t),
                         Err(e) => {
                             ctx.set_error(builder.len(), e);
@@ -989,7 +1022,7 @@ macro_rules! impl_register_arith_functions {
                 |_, _, _| FunctionDomain::MayThrow,
                 vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
                     |ts, delta, builder, ctx| {
-                        match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
+                        match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, true) {
                             Ok(t) => builder.push(t),
                             Err(e) => {
                                 ctx.set_error(builder.len(), e);
diff --git a/src/query/functions/src/scalars/timestamp/src/interval.rs b/src/query/functions/src/scalars/timestamp/src/interval.rs
@@ -117,7 +117,12 @@ fn register_interval_add_sub_mul(registry: &mut FunctionRegistry) {
                     let ts = a
                         .wrapping_add(b.microseconds())
                         .wrapping_add((b.days() as i64).wrapping_mul(86_400_000_000));
-                    match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), b.months()) {
+                    match EvalMonthsImpl::eval_timestamp(
+                        ts,
+                        ctx.func_ctx.tz.clone(),
+                        b.months(),
+                        false,
+                    ) {
                         Ok(t) => output.push(t),
                         Err(e) => {
                             ctx.set_error(output.len(), e);
@@ -138,7 +143,12 @@ fn register_interval_add_sub_mul(registry: &mut FunctionRegistry) {
                     let ts = a
                         .wrapping_add(b.microseconds())
                         .wrapping_add((b.days() as i64).wrapping_mul(86_400_000_000));
-                    match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), b.months()) {
+                    match EvalMonthsImpl::eval_timestamp(
+                        ts,
+                        ctx.func_ctx.tz.clone(),
+                        b.months(),
+                        false,
+                    ) {
                         Ok(t) => output.push(t),
                         Err(e) => {
                             ctx.set_error(output.len(), e);
@@ -173,7 +183,12 @@ fn register_interval_add_sub_mul(registry: &mut FunctionRegistry) {
                     let ts = a
                         .wrapping_sub(b.microseconds())
                         .wrapping_sub((b.days() as i64).wrapping_mul(86_400_000_000));
-                    match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), -b.months()) {
+                    match EvalMonthsImpl::eval_timestamp(
+                        ts,
+                        ctx.func_ctx.tz.clone(),
+                        -b.months(),
+                        false,
+                    ) {
                         Ok(t) => output.push(t),
                         Err(e) => {
                             ctx.set_error(output.len(), e);
diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt
@@ -1270,6 +1270,14 @@ Functions overloads:
 1 cot(Float64 NULL) :: Float64 NULL
 0 crc32(String) :: UInt32
 1 crc32(String NULL) :: UInt32 NULL
+0 date_add_months(Date, Int64) :: Date
+1 date_add_months(Date NULL, Int64 NULL) :: Date NULL
+2 date_add_months(Timestamp, Int64) :: Timestamp
+3 date_add_months(Timestamp NULL, Int64 NULL) :: Timestamp NULL
+0 date_subtract_months(Date, Int64) :: Date
+1 date_subtract_months(Date NULL, Int64 NULL) :: Date NULL
+2 date_subtract_months(Timestamp, Int64) :: Timestamp
+3 date_subtract_months(Timestamp NULL, Int64 NULL) :: Timestamp NULL
 0 dayofweek(Date) :: UInt8
 1 dayofweek(Date NULL) :: UInt8 NULL
 2 dayofweek(Timestamp) :: UInt8
diff --git a/src/query/sql/src/planner/semantic/type_check.rs b/src/query/sql/src/planner/semantic/type_check.rs
@@ -3283,7 +3283,12 @@ impl<'a> TypeChecker<'a> {
         let func_name = match is_diff {
             Expr::DateDiff { .. } => format!("diff_{}s", interval_kind.to_string().to_lowercase()),
             Expr::DateSub { .. } | Expr::DateAdd { .. } => {
-                format!("add_{}s", interval_kind.to_string().to_lowercase())
+                let interval_kind = interval_kind.to_string().to_lowercase();
+                if interval_kind == "month" {
+                    format!("date_add_{}s", interval_kind.to_string().to_lowercase())
+                } else {
+                    format!("add_{}s", interval_kind.to_string().to_lowercase())
+                }
             }
             Expr::DateBetween { .. } => {
                 format!("between_{}s", interval_kind.to_string().to_lowercase())
diff --git a/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test b/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test
@@ -757,13 +757,13 @@ select add_years(to_datetime('9998-12-30 21:59:59'), 1)
 query T
 select subtract_months(to_date(18321), cast(13, INT16))
 ----
-2019-01-29
+2019-01-31
 
 # 2020-2-29T10:00:00 - (12*10 + 2) months
 query T
 select subtract_months(to_datetime(1582970400000000), cast(122, INT16))
 ----
-2009-12-29 10:00:00.000000
+2009-12-31 10:00:00.000000
 
 
 query T
@@ -1777,3 +1777,23 @@ query T
 select datebetween('microsecond','2019-02-28'::Date, '2020-02-28'::Date), date_between('microsecond','2019-03-01'::Date, '2020-02-28'::Date), date_between('microsecond','2019-02-28 22:00:01'::Timestamp, '2020-02-28 22:00:00'::Timestamp), date_between('microsecond','2020-02-28 22:00:01'::Timestamp, '2019-02-28 22:00:01'::Timestamp);
 ----
 31536000000000 31449600000000 31535999000000 -31536000000000
+
+query T
+select add_months('2020-02-29',12), DATE_ADD(month, 12, '2020-02-29');
+----
+2021-02-28 2021-02-28
+
+query T
+select add_months('2020-02-28',12), DATE_ADD(month, 12, '2020-02-28');
+----
+2021-02-28 2021-02-28
+
+query T
+select add_months('2020-02-27',12), DATE_ADD(month, 12, '2020-02-27');
+----
+2021-02-27 2021-02-27
+
+query T
+select add_months('2025-04-30',1), DATE_ADD(month, 1, '2025-04-30');
+----
+2025-05-31 2025-05-30
