Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(query): fix add/subtract datetime with big integer panic #12940

Merged
merged 4 commits into from
Sep 20, 2023
Merged

fix(query): fix add/subtract datetime with big integer panic #12940

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
07862eb
fix(query): fix add/subtract datetime with big integer panic
b41sh Sep 19, 2023
be826b3
fix
b41sh Sep 20, 2023
300f671
fix domain
b41sh Sep 20, 2023
bfeb83b
Merge branch 'main' into fix-to_quarter
b41sh Sep 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions src/query/expression/src/utils/date_helper.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -417,6 +417,7 @@ pub struct ToYYYYMMDD;
pub struct ToYYYYMMDDHH;
pub struct ToYYYYMMDDHHMMSS;
pub struct ToYear;
pub struct ToQuarter;
pub struct ToMonth;
pub struct ToDayOfYear;
pub struct ToDayOfMonth;
Expand Down Expand Up @@ -464,6 +465,12 @@ impl ToNumber<u16> for ToYear {
}
}

impl ToNumber<u8> for ToQuarter {
fn to_number(dt: &DateTime<Tz>) -> u8 {
(dt.month0() / 3 + 1) as u8
}
}

impl ToNumber<u8> for ToMonth {
fn to_number(dt: &DateTime<Tz>) -> u8 {
dt.month() as u8
Expand Down
102 changes: 74 additions & 28 deletions src/query/functions/src/scalars/datetime.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -909,6 +909,13 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
ToNumberImpl::eval_date::<ToYear, _>(val, ctx.func_ctx.tz)
}),
);
registry.register_passthrough_nullable_1_arg::<DateType, UInt8Type, _, _>(
"to_quarter",
|_, _| FunctionDomain::Full,
vectorize_1_arg::<DateType, UInt8Type>(|val, ctx| {
ToNumberImpl::eval_date::<ToQuarter, _>(val, ctx.func_ctx.tz)
}),
);
registry.register_passthrough_nullable_1_arg::<DateType, UInt8Type, _, _>(
"to_month",
|_, _| FunctionDomain::Full,
Expand Down Expand Up @@ -973,6 +980,13 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
ToNumberImpl::eval_timestamp::<ToYear, _>(val, ctx.func_ctx.tz)
}),
);
registry.register_passthrough_nullable_1_arg::<TimestampType, UInt8Type, _, _>(
"to_quarter",
|_, _| FunctionDomain::Full,
vectorize_1_arg::<TimestampType, UInt8Type>(|val, ctx| {
ToNumberImpl::eval_timestamp::<ToQuarter, _>(val, ctx.func_ctx.tz)
}),
);
registry.register_passthrough_nullable_1_arg::<TimestampType, UInt8Type, _, _>(
"to_month",
|_, _| FunctionDomain::Full,
Expand Down Expand Up @@ -1027,23 +1041,31 @@ fn register_to_number_functions(registry: &mut FunctionRegistry) {
}

fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
registry.register_2_arg::<DateType, Int64Type, DateType, _, _>(
registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
"plus",
|_, lhs, rhs| {
(|| {
let lm = lhs.max;
let ln = lhs.min;
let rm: i32 = num_traits::cast::cast(rhs.max)?;
let rn: i32 = num_traits::cast::cast(rhs.min)?;
let lm: i64 = num_traits::cast::cast(lhs.max)?;
let ln: i64 = num_traits::cast::cast(lhs.min)?;
let rm = rhs.max;
let rn = rhs.min;

Some(FunctionDomain::Domain(SimpleDomain::<i32> {
min: ln.checked_add(rn)?,
max: lm.checked_add(rm)?,
min: check_date(ln + rn).ok()?,
max: check_date(lm + rm).ok()?,
}))
})()
.unwrap_or(FunctionDomain::Full)
.unwrap_or(FunctionDomain::MayThrow)
},
|a, b, _| a + (b as i32),
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|a, b, output, ctx| {
match check_date((a as i64) + b) {
Ok(v) => output.push(v),
Err(err) => {
ctx.set_error(output.len(), err);
output.push(0);
}
}
}),
);

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

registry.register_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
"plus",
|_, lhs, rhs| {
(|| {
Expand All @@ -1074,13 +1096,21 @@ fn register_timestamp_add_sub(registry: &mut FunctionRegistry) {
let rm = rhs.max;
let rn = rhs.min;
Some(FunctionDomain::Domain(SimpleDomain::<i64> {
min: ln.checked_add(rn)?,
max: lm.checked_add(rm)?,
min: check_timestamp(ln + rn).ok()?,
max: check_timestamp(lm + rm).ok()?,
}))
})()
.unwrap_or(FunctionDomain::Full)
.unwrap_or(FunctionDomain::MayThrow)
},
|a, b, _| a + b,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|a, b, output, ctx| match check_timestamp(a + b) {
Ok(v) => output.push(v),
Err(err) => {
ctx.set_error(output.len(), err);
output.push(0);
}
},
),
);

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

registry.register_2_arg::<DateType, Int64Type, DateType, _, _>(
registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
"minus",
|_, lhs, rhs| {
(|| {
let lm = lhs.max;
let ln = lhs.min;
let rm: i32 = num_traits::cast::cast(rhs.max)?;
let rn: i32 = num_traits::cast::cast(rhs.min)?;
let lm: i64 = num_traits::cast::cast(lhs.max)?;
let ln: i64 = num_traits::cast::cast(lhs.min)?;
let rm = rhs.max;
let rn = rhs.min;

Some(FunctionDomain::Domain(SimpleDomain::<i32> {
min: ln.checked_sub(rm)?,
max: lm.checked_sub(rn)?,
min: check_date(ln - rn).ok()?,
max: check_date(lm - rm).ok()?,
}))
})()
.unwrap_or(FunctionDomain::Full)
.unwrap_or(FunctionDomain::MayThrow)
},
|a, b, _| a - b as i32,
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|a, b, output, ctx| {
match check_date((a as i64) - b) {
Ok(v) => output.push(v),
Err(err) => {
ctx.set_error(output.len(), err);
output.push(0);
}
}
}),
);

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

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

Some(FunctionDomain::Domain(SimpleDomain::<i64> {
min: ln.checked_sub(rm)?,
max: lm.checked_sub(rn)?,
min: check_timestamp(ln - rn).ok()?,
max: check_timestamp(lm - rm).ok()?,
}))
})()
.unwrap_or(FunctionDomain::Full)
.unwrap_or(FunctionDomain::MayThrow)
},
|a, b, _| a - b,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|a, b, output, ctx| match check_timestamp(a - b) {
Ok(v) => output.push(v),
Err(err) => {
ctx.set_error(output.len(), err);
output.push(0);
}
},
),
);

registry.register_2_arg::<TimestampType, TimestampType, Int64Type, _, _>(
Expand Down
18 changes: 18 additions & 0 deletions src/query/functions/tests/it/scalars/datetime.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,10 +108,14 @@ fn test_date_add_subtract(file: &mut impl Write) {
run_ast(file, "add_years(to_date(0), 100)", &[]);
run_ast(file, "add_months(to_date(0), 100)", &[]);
run_ast(file, "add_days(to_date(0), 100)", &[]);
run_ast(file, "add(to_date(0), 100)", &[]);
run_ast(file, "add(to_date(0), 10000000)", &[]);
run_ast(file, "subtract_years(to_date(0), 100)", &[]);
run_ast(file, "subtract_quarters(to_date(0), 100)", &[]);
run_ast(file, "subtract_months(to_date(0), 100)", &[]);
run_ast(file, "subtract_days(to_date(0), 100)", &[]);
run_ast(file, "subtract(to_date(0), 100)", &[]);
run_ast(file, "subtract(to_date(0), 10000000)", &[]);
run_ast(file, "add_years(a, b)", &[
("a", DateType::from_data(vec![-100, 0, 100])),
("b", Int32Type::from_data(vec![1, 2, 3])),
Expand Down Expand Up @@ -155,13 +159,17 @@ fn test_timestamp_add_subtract(file: &mut impl Write) {
run_ast(file, "add_hours(to_timestamp(0), 100)", &[]);
run_ast(file, "add_minutes(to_timestamp(0), 100)", &[]);
run_ast(file, "add_seconds(to_timestamp(0), 100)", &[]);
run_ast(file, "add(to_timestamp(0), 100000000000000)", &[]);
run_ast(file, "add(to_timestamp(0), 1000000000000000000)", &[]);
run_ast(file, "subtract_years(to_timestamp(0), 100)", &[]);
run_ast(file, "subtract_quarters(to_timestamp(0), 100)", &[]);
run_ast(file, "subtract_months(to_timestamp(0), 100)", &[]);
run_ast(file, "subtract_days(to_timestamp(0), 100)", &[]);
run_ast(file, "subtract_hours(to_timestamp(0), 100)", &[]);
run_ast(file, "subtract_minutes(to_timestamp(0), 100)", &[]);
run_ast(file, "subtract_seconds(to_timestamp(0), 100)", &[]);
run_ast(file, "subtract(to_timestamp(0), 100000000000000)", &[]);
run_ast(file, "subtract(to_timestamp(0), 1000000000000000000)", &[]);
run_ast(file, "add_years(a, b)", &[
("a", TimestampType::from_data(vec![-100, 0, 100])),
("b", Int32Type::from_data(vec![1, 2, 3])),
Expand Down Expand Up @@ -462,6 +470,7 @@ fn test_to_number(file: &mut impl Write) {
run_ast(file, "to_yyyymmdd(to_date(18875))", &[]);
run_ast(file, "to_yyyymmddhhmmss(to_date(18875))", &[]);
run_ast(file, "to_year(to_date(18875))", &[]);
run_ast(file, "to_quarter(to_date(18875))", &[]);
run_ast(file, "to_month(to_date(18875))", &[]);
run_ast(file, "to_day_of_year(to_date(18875))", &[]);
run_ast(file, "to_day_of_month(to_date(18875))", &[]);
Expand All @@ -482,6 +491,10 @@ fn test_to_number(file: &mut impl Write) {
"a",
DateType::from_data(vec![-100, 0, 100]),
)]);
run_ast(file, "to_quarter(a)", &[(
"a",
DateType::from_data(vec![-100, 0, 100]),
)]);
run_ast(file, "to_month(a)", &[(
"a",
DateType::from_data(vec![-100, 0, 100]),
Expand All @@ -504,6 +517,7 @@ fn test_to_number(file: &mut impl Write) {
run_ast(file, "to_yyyymmdd(to_timestamp(1630812366))", &[]);
run_ast(file, "to_yyyymmddhhmmss(to_timestamp(1630812366))", &[]);
run_ast(file, "to_year(to_timestamp(1630812366))", &[]);
run_ast(file, "to_quarter(to_timestamp(1630812366))", &[]);
run_ast(file, "to_month(to_timestamp(1630812366))", &[]);
run_ast(file, "to_day_of_year(to_timestamp(1630812366))", &[]);
run_ast(file, "to_day_of_month(to_timestamp(1630812366))", &[]);
Expand All @@ -527,6 +541,10 @@ fn test_to_number(file: &mut impl Write) {
"a",
TimestampType::from_data(vec![-100, 0, 100]),
)]);
run_ast(file, "to_quarter(a)", &[(
"a",
TimestampType::from_data(vec![-100, 0, 100]),
)]);
run_ast(file, "to_month(a)", &[(
"a",
TimestampType::from_data(vec![-100, 0, 100]),
Expand Down
Loading