Skip to content

Commit 364d955

Browse files
TCeasonTCeason
authored
fix(query): add_months need to respect last day of month (#18201)
1 parent 0224108 commit 364d955

File tree

7 files changed

+122
-21
lines changed

7 files changed

+122
-21
lines changed

‎src/query/expression/src/utils/date_helper.rs‎

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,13 @@ pub const FACTOR_MINUTE: i64 = 60;
7474
pub const FACTOR_SECOND: i64 = 1;
7575
const LAST_DAY_LUT: [i8; 13] = [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31];
7676

77-
fn eval_years_base(year: i16, month: i8, day: i8, delta: i64) -> std::result::Result<Date, String> {
77+
fn eval_years_base(
78+
year: i16,
79+
month: i8,
80+
day: i8,
81+
delta: i64,
82+
_add_months: bool,
83+
) -> std::result::Result<Date, String> {
7884
let new_year = year as i64 + delta;
7985
let mut new_day = day;
8086
if std::intrinsics::unlikely(month == 2 && day == 29) {
@@ -91,6 +97,7 @@ fn eval_months_base(
9197
month: i8,
9298
day: i8,
9399
delta: i64,
100+
add_months: bool,
94101
) -> std::result::Result<Date, String> {
95102
let total_months = (month as i64 + delta - 1) as i16;
96103
let mut new_year = year + (total_months / 12);
@@ -101,12 +108,17 @@ fn eval_months_base(
101108
}
102109

103110
// 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".
104-
let new_day = std::cmp::min::<u32>(
105-
day as u32,
106-
last_day_of_year_month(new_year, (new_month0 + 1) as i8) as u32,
107-
);
111+
// 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.
112+
let new_month = (new_month0 + 1) as i8;
113+
// Determine the correct day
114+
let max_day = last_day_of_year_month(new_year, new_month);
115+
let new_day = if add_months && day == last_day_of_year_month(year, month) {
116+
max_day
117+
} else {
118+
day.min(max_day)
119+
};
108120

109-
match Date::new(new_year, (new_month0 + 1) as i8, new_day as i8) {
121+
match Date::new(new_year, (new_month0 + 1) as i8, new_day) {
110122
Ok(d) => Ok(d),
111123
Err(e) => Err(format!("Invalid date: {}", e)),
112124
}
@@ -131,9 +143,16 @@ macro_rules! impl_interval_year_month {
131143
date: i32,
132144
tz: TimeZone,
133145
delta: impl AsPrimitive<i64>,
146+
add_months: bool,
134147
) -> std::result::Result<i32, String> {
135148
let date = date.to_date(tz);
136-
let new_date = $op(date.year(), date.month(), date.day(), delta.as_())?;
149+
let new_date = $op(
150+
date.year(),
151+
date.month(),
152+
date.day(),
153+
delta.as_(),
154+
add_months,
155+
)?;
137156

138157
Ok(clamp_date(
139158
new_date
@@ -147,9 +166,10 @@ macro_rules! impl_interval_year_month {
147166
us: i64,
148167
tz: TimeZone,
149168
delta: impl AsPrimitive<i64>,
169+
add_months: bool,
150170
) -> std::result::Result<i64, String> {
151171
let ts = us.to_timestamp(tz.clone());
152-
let new_date = $op(ts.year(), ts.month(), ts.day(), delta.as_())?;
172+
let new_date = $op(ts.year(), ts.month(), ts.day(), delta.as_(), add_months)?;
153173

154174
let mut ts = new_date
155175
.at(ts.hour(), ts.minute(), ts.second(), ts.subsec_nanosecond())

‎src/query/functions/src/scalars/timestamp/src/datetime.rs‎

Lines changed: 39 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -912,7 +912,7 @@ macro_rules! impl_register_arith_functions {
912912

913913
|_, _, _| FunctionDomain::MayThrow,
914914
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
915-
match EvalYearsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
915+
match EvalYearsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
916916
Ok(t) => builder.push(t),
917917
Err(e) => {
918918
ctx.set_error(builder.len(), e);
@@ -927,7 +927,7 @@ macro_rules! impl_register_arith_functions {
927927
|_, _, _| FunctionDomain::MayThrow,
928928
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
929929
|ts, delta, builder, ctx| {
930-
match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
930+
match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
931931
Ok(t) => builder.push(t),
932932
Err(e) => {
933933
ctx.set_error(builder.len(), e);
@@ -943,7 +943,7 @@ macro_rules! impl_register_arith_functions {
943943

944944
|_, _, _| FunctionDomain::MayThrow,
945945
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
946-
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3) {
946+
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3, false) {
947947
Ok(t) => builder.push(t),
948948
Err(e) => {
949949
ctx.set_error(builder.len(), e);
@@ -958,7 +958,7 @@ macro_rules! impl_register_arith_functions {
958958
|_, _, _| FunctionDomain::MayThrow,
959959
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
960960
|ts, delta, builder, ctx| {
961-
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3) {
961+
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta} * 3, false) {
962962
Ok(t) => builder.push(t),
963963
Err(e) => {
964964
ctx.set_error(builder.len(), e);
@@ -969,12 +969,45 @@ macro_rules! impl_register_arith_functions {
969969
),
970970
);
971971

972+
registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
973+
concat!("date_", $op, "_months"),
974+
975+
|_, _, _| FunctionDomain::MayThrow,
976+
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
977+
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
978+
Ok(t) => builder.push(t),
979+
Err(e) => {
980+
ctx.set_error(builder.len(), e);
981+
builder.push(0);
982+
},
983+
}
984+
}),
985+
);
986+
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
987+
concat!("date_", $op, "_months"),
988+
989+
|_, _, _| FunctionDomain::MayThrow,
990+
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
991+
|ts, delta, builder, ctx| {
992+
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, false) {
993+
Ok(t) => builder.push(t),
994+
Err(e) => {
995+
ctx.set_error(builder.len(), e);
996+
builder.push(0);
997+
},
998+
}
999+
},
1000+
),
1001+
);
1002+
1003+
// 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.
1004+
// 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.
9721005
registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
9731006
concat!($op, "_months"),
9741007

9751008
|_, _, _| FunctionDomain::MayThrow,
9761009
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
977-
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
1010+
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, true) {
9781011
Ok(t) => builder.push(t),
9791012
Err(e) => {
9801013
ctx.set_error(builder.len(), e);
@@ -989,7 +1022,7 @@ macro_rules! impl_register_arith_functions {
9891022
|_, _, _| FunctionDomain::MayThrow,
9901023
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
9911024
|ts, delta, builder, ctx| {
992-
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}) {
1025+
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), $signed_wrapper!{delta}, true) {
9931026
Ok(t) => builder.push(t),
9941027
Err(e) => {
9951028
ctx.set_error(builder.len(), e);

‎src/query/functions/src/scalars/timestamp/src/interval.rs‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,12 @@ fn register_interval_add_sub_mul(registry: &mut FunctionRegistry) {
117117
let ts = a
118118
.wrapping_add(b.microseconds())
119119
.wrapping_add((b.days() as i64).wrapping_mul(86_400_000_000));
120-
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), b.months()) {
120+
match EvalMonthsImpl::eval_timestamp(
121+
ts,
122+
ctx.func_ctx.tz.clone(),
123+
b.months(),
124+
false,
125+
) {
121126
Ok(t) => output.push(t),
122127
Err(e) => {
123128
ctx.set_error(output.len(), e);
@@ -138,7 +143,12 @@ fn register_interval_add_sub_mul(registry: &mut FunctionRegistry) {
138143
let ts = a
139144
.wrapping_add(b.microseconds())
140145
.wrapping_add((b.days() as i64).wrapping_mul(86_400_000_000));
141-
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), b.months()) {
146+
match EvalMonthsImpl::eval_timestamp(
147+
ts,
148+
ctx.func_ctx.tz.clone(),
149+
b.months(),
150+
false,
151+
) {
142152
Ok(t) => output.push(t),
143153
Err(e) => {
144154
ctx.set_error(output.len(), e);
@@ -173,7 +183,12 @@ fn register_interval_add_sub_mul(registry: &mut FunctionRegistry) {
173183
let ts = a
174184
.wrapping_sub(b.microseconds())
175185
.wrapping_sub((b.days() as i64).wrapping_mul(86_400_000_000));
176-
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.tz.clone(), -b.months()) {
186+
match EvalMonthsImpl::eval_timestamp(
187+
ts,
188+
ctx.func_ctx.tz.clone(),
189+
-b.months(),
190+
false,
191+
) {
177192
Ok(t) => output.push(t),
178193
Err(e) => {
179194
ctx.set_error(output.len(), e);

‎src/query/functions/tests/it/scalars/testdata/function_list.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1298,6 +1298,14 @@ Functions overloads:
12981298
1 cot(Float64 NULL) :: Float64 NULL
12991299
0 crc32(String) :: UInt32
13001300
1 crc32(String NULL) :: UInt32 NULL
1301+
0 date_add_months(Date, Int64) :: Date
1302+
1 date_add_months(Date NULL, Int64 NULL) :: Date NULL
1303+
2 date_add_months(Timestamp, Int64) :: Timestamp
1304+
3 date_add_months(Timestamp NULL, Int64 NULL) :: Timestamp NULL
1305+
0 date_subtract_months(Date, Int64) :: Date
1306+
1 date_subtract_months(Date NULL, Int64 NULL) :: Date NULL
1307+
2 date_subtract_months(Timestamp, Int64) :: Timestamp
1308+
3 date_subtract_months(Timestamp NULL, Int64 NULL) :: Timestamp NULL
13011309
0 dayofweek(Date) :: UInt8
13021310
1 dayofweek(Date NULL) :: UInt8 NULL
13031311
2 dayofweek(Timestamp) :: UInt8

‎src/query/sql/src/planner/semantic/type_check.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3342,7 +3342,12 @@ impl<'a> TypeChecker<'a> {
33423342
let func_name = match is_diff {
33433343
Expr::DateDiff { .. } => format!("diff_{}s", interval_kind.to_string().to_lowercase()),
33443344
Expr::DateSub { .. } | Expr::DateAdd { .. } => {
3345-
format!("add_{}s", interval_kind.to_string().to_lowercase())
3345+
let interval_kind = interval_kind.to_string().to_lowercase();
3346+
if interval_kind == "month" {
3347+
format!("date_add_{}s", interval_kind.to_string().to_lowercase())
3348+
} else {
3349+
format!("add_{}s", interval_kind.to_string().to_lowercase())
3350+
}
33463351
}
33473352
Expr::DateBetween { .. } => {
33483353
format!("between_{}s", interval_kind.to_string().to_lowercase())

‎tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -757,13 +757,13 @@ select add_years(to_datetime('9998-12-30 21:59:59'), 1)
757757
query T
758758
select subtract_months(to_date(18321), cast(13, INT16))
759759
----
760-
2019-01-29
760+
2019-01-31
761761

762762
# 2020-2-29T10:00:00 - (12*10 + 2) months
763763
query T
764764
select subtract_months(to_datetime(1582970400000000), cast(122, INT16))
765765
----
766-
2009-12-29 10:00:00.000000
766+
2009-12-31 10:00:00.000000
767767

768768

769769
query T
@@ -1777,3 +1777,23 @@ query T
17771777
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);
17781778
----
17791779
31536000000000 31449600000000 31535999000000 -31536000000000
1780+
1781+
query T
1782+
select add_months('2020-02-29',12), DATE_ADD(month, 12, '2020-02-29');
1783+
----
1784+
2021-02-28 2021-02-28
1785+
1786+
query T
1787+
select add_months('2020-02-28',12), DATE_ADD(month, 12, '2020-02-28');
1788+
----
1789+
2021-02-28 2021-02-28
1790+
1791+
query T
1792+
select add_months('2020-02-27',12), DATE_ADD(month, 12, '2020-02-27');
1793+
----
1794+
2021-02-27 2021-02-27
1795+
1796+
query T
1797+
select add_months('2025-04-30',1), DATE_ADD(month, 1, '2025-04-30');
1798+
----
1799+
2025-05-31 2025-05-30

‎tests/sqllogictests/suites/query/functions/02_0012_function_datetimes_tz.test‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -457,7 +457,7 @@ set timezone = 'UTC'
457457
query T
458458
select add_months(to_timestamp(1619822911999000), 1)
459459
----
460-
2021-05-30 22:48:31.999000
460+
2021-05-31 22:48:31.999000
461461

462462
query T
463463
select to_timestamp(1583013600000000)

0 commit comments

Comments
 (0)