Skip to content

Commit 321b99b

Browse files
b41shb41sh
authored
fix(query): fix add/subtract datetime with big integer panic (#12940)
* fix(query): fix add/subtract datetime with big integer panic * fix * fix domain
1 parent 2add53d commit 321b99b

File tree

6 files changed

+268
-28
lines changed

6 files changed

+268
-28
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -417,6 +417,7 @@ pub struct ToYYYYMMDD;
417417
pub struct ToYYYYMMDDHH;
418418
pub struct ToYYYYMMDDHHMMSS;
419419
pub struct ToYear;
420+
pub struct ToQuarter;
420421
pub struct ToMonth;
421422
pub struct ToDayOfYear;
422423
pub struct ToDayOfMonth;
@@ -464,6 +465,12 @@ impl ToNumber<u16> for ToYear {
464465
}
465466
}
466467

468+
impl ToNumber<u8> for ToQuarter {
469+
fn to_number(dt: &DateTime<Tz>) -> u8 {
470+
(dt.month0() / 3 + 1) as u8
471+
}
472+
}
473+
467474
impl ToNumber<u8> for ToMonth {
468475
fn to_number(dt: &DateTime<Tz>) -> u8 {
469476
dt.month() as u8

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

Lines changed: 74 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -909,6 +909,13 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
909909
ToNumberImpl::eval_date::<ToYear, _>(val, ctx.func_ctx.tz)
910910
}),
911911
);
912+
registry.register_passthrough_nullable_1_arg::<DateType, UInt8Type, _, _>(
913+
"to_quarter",
914+
|_, _| FunctionDomain::Full,
915+
vectorize_1_arg::<DateType, UInt8Type>(|val, ctx| {
916+
ToNumberImpl::eval_date::<ToQuarter, _>(val, ctx.func_ctx.tz)
917+
}),
918+
);
912919
registry.register_passthrough_nullable_1_arg::<DateType, UInt8Type, _, _>(
913920
"to_month",
914921
|_, _| FunctionDomain::Full,
@@ -973,6 +980,13 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
973980
ToNumberImpl::eval_timestamp::<ToYear, _>(val, ctx.func_ctx.tz)
974981
}),
975982
);
983+
registry.register_passthrough_nullable_1_arg::<TimestampType, UInt8Type, _, _>(
984+
"to_quarter",
985+
|_, _| FunctionDomain::Full,
986+
vectorize_1_arg::<TimestampType, UInt8Type>(|val, ctx| {
987+
ToNumberImpl::eval_timestamp::<ToQuarter, _>(val, ctx.func_ctx.tz)
988+
}),
989+
);
976990
registry.register_passthrough_nullable_1_arg::<TimestampType, UInt8Type, _, _>(
977991
"to_month",
978992
|_, _| FunctionDomain::Full,
@@ -1027,23 +1041,31 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
10271041
}
10281042

10291043
fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
1030-
registry.register_2_arg::<DateType, Int64Type, DateType, _, _>(
1044+
registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
10311045
"plus",
10321046
|_, lhs, rhs| {
10331047
(|| {
1034-
let lm = lhs.max;
1035-
let ln = lhs.min;
1036-
let rm: i32 = num_traits::cast::cast(rhs.max)?;
1037-
let rn: i32 = num_traits::cast::cast(rhs.min)?;
1048+
let lm: i64 = num_traits::cast::cast(lhs.max)?;
1049+
let ln: i64 = num_traits::cast::cast(lhs.min)?;
1050+
let rm = rhs.max;
1051+
let rn = rhs.min;
10381052

10391053
Some(FunctionDomain::Domain(SimpleDomain::<i32> {
1040-
min: ln.checked_add(rn)?,
1041-
max: lm.checked_add(rm)?,
1054+
min: check_date(ln + rn).ok()?,
1055+
max: check_date(lm + rm).ok()?,
10421056
}))
10431057
})()
1044-
.unwrap_or(FunctionDomain::Full)
1058+
.unwrap_or(FunctionDomain::MayThrow)
10451059
},
1046-
|a, b, _| a + (b as i32),
1060+
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|a, b, output, ctx| {
1061+
match check_date((a as i64) + b) {
1062+
Ok(v) => output.push(v),
1063+
Err(err) => {
1064+
ctx.set_error(output.len(), err);
1065+
output.push(0);
1066+
}
1067+
}
1068+
}),
10471069
);
10481070

10491071
registry.register_2_arg::<DateType, DateType, Int32Type, _, _>(
@@ -1065,7 +1087,7 @@ fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
10651087
|a, b, _| a + b,
10661088
);
10671089

1068-
registry.register_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
1090+
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
10691091
"plus",
10701092
|_, lhs, rhs| {
10711093
(|| {
@@ -1074,13 +1096,21 @@ fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
10741096
let rm = rhs.max;
10751097
let rn = rhs.min;
10761098
Some(FunctionDomain::Domain(SimpleDomain::<i64> {
1077-
min: ln.checked_add(rn)?,
1078-
max: lm.checked_add(rm)?,
1099+
min: check_timestamp(ln + rn).ok()?,
1100+
max: check_timestamp(lm + rm).ok()?,
10791101
}))
10801102
})()
1081-
.unwrap_or(FunctionDomain::Full)
1103+
.unwrap_or(FunctionDomain::MayThrow)
10821104
},
1083-
|a, b, _| a + b,
1105+
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
1106+
|a, b, output, ctx| match check_timestamp(a + b) {
1107+
Ok(v) => output.push(v),
1108+
Err(err) => {
1109+
ctx.set_error(output.len(), err);
1110+
output.push(0);
1111+
}
1112+
},
1113+
),
10841114
);
10851115

10861116
registry.register_2_arg::<TimestampType, TimestampType, Int64Type, _, _>(
@@ -1101,23 +1131,31 @@ fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
11011131
|a, b, _| a + b,
11021132
);
11031133

1104-
registry.register_2_arg::<DateType, Int64Type, DateType, _, _>(
1134+
registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
11051135
"minus",
11061136
|_, lhs, rhs| {
11071137
(|| {
1108-
let lm = lhs.max;
1109-
let ln = lhs.min;
1110-
let rm: i32 = num_traits::cast::cast(rhs.max)?;
1111-
let rn: i32 = num_traits::cast::cast(rhs.min)?;
1138+
let lm: i64 = num_traits::cast::cast(lhs.max)?;
1139+
let ln: i64 = num_traits::cast::cast(lhs.min)?;
1140+
let rm = rhs.max;
1141+
let rn = rhs.min;
11121142

11131143
Some(FunctionDomain::Domain(SimpleDomain::<i32> {
1114-
min: ln.checked_sub(rm)?,
1115-
max: lm.checked_sub(rn)?,
1144+
min: check_date(ln - rn).ok()?,
1145+
max: check_date(lm - rm).ok()?,
11161146
}))
11171147
})()
1118-
.unwrap_or(FunctionDomain::Full)
1148+
.unwrap_or(FunctionDomain::MayThrow)
11191149
},
1120-
|a, b, _| a - b as i32,
1150+
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|a, b, output, ctx| {
1151+
match check_date((a as i64) - b) {
1152+
Ok(v) => output.push(v),
1153+
Err(err) => {
1154+
ctx.set_error(output.len(), err);
1155+
output.push(0);
1156+
}
1157+
}
1158+
}),
11211159
);
11221160

11231161
registry.register_2_arg::<DateType, DateType, Int32Type, _, _>(
@@ -1139,7 +1177,7 @@ fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
11391177
|a, b, _| a - b,
11401178
);
11411179

1142-
registry.register_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
1180+
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
11431181
"minus",
11441182
|_, lhs, rhs| {
11451183
(|| {
@@ -1149,13 +1187,21 @@ fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
11491187
let rn = rhs.min;
11501188

11511189
Some(FunctionDomain::Domain(SimpleDomain::<i64> {
1152-
min: ln.checked_sub(rm)?,
1153-
max: lm.checked_sub(rn)?,
1190+
min: check_timestamp(ln - rn).ok()?,
1191+
max: check_timestamp(lm - rm).ok()?,
11541192
}))
11551193
})()
1156-
.unwrap_or(FunctionDomain::Full)
1194+
.unwrap_or(FunctionDomain::MayThrow)
11571195
},
1158-
|a, b, _| a - b,
1196+
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
1197+
|a, b, output, ctx| match check_timestamp(a - b) {
1198+
Ok(v) => output.push(v),
1199+
Err(err) => {
1200+
ctx.set_error(output.len(), err);
1201+
output.push(0);
1202+
}
1203+
},
1204+
),
11591205
);
11601206

11611207
registry.register_2_arg::<TimestampType, TimestampType, Int64Type, _, _>(

‎src/query/functions/tests/it/scalars/datetime.rs

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,10 +108,14 @@ fn test_date_add_subtract(file: &mut impl Write) {
108108
run_ast(file, "add_years(to_date(0), 100)", &[]);
109109
run_ast(file, "add_months(to_date(0), 100)", &[]);
110110
run_ast(file, "add_days(to_date(0), 100)", &[]);
111+
run_ast(file, "add(to_date(0), 100)", &[]);
112+
run_ast(file, "add(to_date(0), 10000000)", &[]);
111113
run_ast(file, "subtract_years(to_date(0), 100)", &[]);
112114
run_ast(file, "subtract_quarters(to_date(0), 100)", &[]);
113115
run_ast(file, "subtract_months(to_date(0), 100)", &[]);
114116
run_ast(file, "subtract_days(to_date(0), 100)", &[]);
117+
run_ast(file, "subtract(to_date(0), 100)", &[]);
118+
run_ast(file, "subtract(to_date(0), 10000000)", &[]);
115119
run_ast(file, "add_years(a, b)", &[
116120
("a", DateType::from_data(vec![-100, 0, 100])),
117121
("b", Int32Type::from_data(vec![1, 2, 3])),
@@ -155,13 +159,17 @@ fn test_timestamp_add_subtract(file: &mut impl Write) {
155159
run_ast(file, "add_hours(to_timestamp(0), 100)", &[]);
156160
run_ast(file, "add_minutes(to_timestamp(0), 100)", &[]);
157161
run_ast(file, "add_seconds(to_timestamp(0), 100)", &[]);
162+
run_ast(file, "add(to_timestamp(0), 100000000000000)", &[]);
163+
run_ast(file, "add(to_timestamp(0), 1000000000000000000)", &[]);
158164
run_ast(file, "subtract_years(to_timestamp(0), 100)", &[]);
159165
run_ast(file, "subtract_quarters(to_timestamp(0), 100)", &[]);
160166
run_ast(file, "subtract_months(to_timestamp(0), 100)", &[]);
161167
run_ast(file, "subtract_days(to_timestamp(0), 100)", &[]);
162168
run_ast(file, "subtract_hours(to_timestamp(0), 100)", &[]);
163169
run_ast(file, "subtract_minutes(to_timestamp(0), 100)", &[]);
164170
run_ast(file, "subtract_seconds(to_timestamp(0), 100)", &[]);
171+
run_ast(file, "subtract(to_timestamp(0), 100000000000000)", &[]);
172+
run_ast(file, "subtract(to_timestamp(0), 1000000000000000000)", &[]);
165173
run_ast(file, "add_years(a, b)", &[
166174
("a", TimestampType::from_data(vec![-100, 0, 100])),
167175
("b", Int32Type::from_data(vec![1, 2, 3])),
@@ -462,6 +470,7 @@ fn test_to_number(file: &mut impl Write) {
462470
run_ast(file, "to_yyyymmdd(to_date(18875))", &[]);
463471
run_ast(file, "to_yyyymmddhhmmss(to_date(18875))", &[]);
464472
run_ast(file, "to_year(to_date(18875))", &[]);
473+
run_ast(file, "to_quarter(to_date(18875))", &[]);
465474
run_ast(file, "to_month(to_date(18875))", &[]);
466475
run_ast(file, "to_day_of_year(to_date(18875))", &[]);
467476
run_ast(file, "to_day_of_month(to_date(18875))", &[]);
@@ -482,6 +491,10 @@ fn test_to_number(file: &mut impl Write) {
482491
"a",
483492
DateType::from_data(vec![-100, 0, 100]),
484493
)]);
494+
run_ast(file, "to_quarter(a)", &[(
495+
"a",
496+
DateType::from_data(vec![-100, 0, 100]),
497+
)]);
485498
run_ast(file, "to_month(a)", &[(
486499
"a",
487500
DateType::from_data(vec![-100, 0, 100]),
@@ -504,6 +517,7 @@ fn test_to_number(file: &mut impl Write) {
504517
run_ast(file, "to_yyyymmdd(to_timestamp(1630812366))", &[]);
505518
run_ast(file, "to_yyyymmddhhmmss(to_timestamp(1630812366))", &[]);
506519
run_ast(file, "to_year(to_timestamp(1630812366))", &[]);
520+
run_ast(file, "to_quarter(to_timestamp(1630812366))", &[]);
507521
run_ast(file, "to_month(to_timestamp(1630812366))", &[]);
508522
run_ast(file, "to_day_of_year(to_timestamp(1630812366))", &[]);
509523
run_ast(file, "to_day_of_month(to_timestamp(1630812366))", &[]);
@@ -527,6 +541,10 @@ fn test_to_number(file: &mut impl Write) {
527541
"a",
528542
TimestampType::from_data(vec![-100, 0, 100]),
529543
)]);
544+
run_ast(file, "to_quarter(a)", &[(
545+
"a",
546+
TimestampType::from_data(vec![-100, 0, 100]),
547+
)]);
530548
run_ast(file, "to_month(a)", &[(
531549
"a",
532550
TimestampType::from_data(vec![-100, 0, 100]),

0 commit comments

Comments
 (0)