Skip to content

Commit

Permalink
feat(expr, frontend): Impl timestamptz [+-] variable interval (#7751)
Browse files Browse the repository at this point in the history 
Handle the date part as naive datetime, and time part as unix milliseconds. This replicates postgres behaviour. We make use of const eval assumption for performance concern.

Approved-By: xiangjinwu
  • Loading branch information
@jon-chuang
jon-chuang authored Feb 8, 2023
1 parent c66e496 commit b60c0ca
Show file tree
Hide file tree
Showing 3 changed files with 118 additions and 10 deletions.
30 changes: 30 additions & 0 deletions e2e_test/batch/functions/session_timezone.slt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,5 +55,35 @@ select '2022-01-01 00:00:00-08:00'::timestamp with time zone::varchar;
----
2022-01-01 00:00:00-08:00

statement ok
set timezone = 'europe/london'

# Add/Subtract timestamptz with interval across a daylight savings boundary
# Daylight savings falls on 2016-10-30 in London timezone

# This should first add the 24 hours crossing the daylight saving boundary from UTC+1->UTC, then the day
query T
select '2016-10-29 12:00:00'::timestamptz + interval '24 hours' + interval '1 day';
----
2016-10-31 11:00:00+00:00

# This should first add the days at UTC+1->UTC boundary (no change to time), then the hours
query T
select (interval '24 hours' + interval '1 day') + '2016-10-29 12:00:00'::timestamptz;
----
2016-10-31 12:00:00+00:00

# Test inverse for subtract, only -1 day is applied at the UTC->UTC+1 boundary (no change to time)
query T
select '2016-10-31 11:00:00+00:00'::timestamptz - interval '24 hours' - interval '1 day';
----
2016-10-29 11:00:00+01:00

# Test inverse for subtract, this time we apply diff 1 day first, then -24 hours at the UTC->UTC+1 boundary
query T
select '2016-10-31 11:00:00+00:00'::timestamptz - (interval '24 hours' + interval '1 day');
----
2016-10-29 12:00:00+01:00

statement ok
set timezone = 'UTC';
11 changes: 3 additions & 8 deletions src/expr/src/vector_op/arithmetic_op.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -239,19 +239,14 @@ fn timestamptz_interval_inner(
f: fn(i64, i64) -> Option<i64>,
) -> Result<i64> {
// Without session TimeZone, we cannot add month/day in local time. See #5826.
// However, we only reject months but accept days, assuming them are always 24-hour and ignoring
// Daylight Saving.
// This is to keep consistent with `tumble_start` of RisingWave / `date_bin` of PostgreSQL.
if r.get_months() != 0 {
if r.get_months() != 0 || r.get_days() != 0 {
return Err(ExprError::UnsupportedFunction(
"timestamp with time zone +/- interval of months".into(),
"timestamp with time zone +/- interval of days".into(),
));
}

let result: Option<i64> = try {
let d = (r.get_days() as i64).checked_mul(24 * 60 * 60 * 1_000_000)?;
let ms = r.get_ms().checked_mul(1000)?;
let delta_usecs = d.checked_add(ms)?;
let delta_usecs = r.get_ms().checked_mul(1000)?;
f(l, delta_usecs)?
};

Expand Down
87 changes: 85 additions & 2 deletions src/frontend/src/expr/session_timezone.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ pub use risingwave_pb::expr::expr_node::Type as ExprType;

pub use crate::expr::expr_rewriter::ExprRewriter;
pub use crate::expr::function_call::FunctionCall;
use crate::expr::{Expr, ExprImpl};
use crate::expr::{Expr, ExprImpl, Literal};

/// `SessionTimezone` will be used to resolve session
/// timezone-dependent casts, comparisons or arithmetic.
Expand Down Expand Up @@ -135,7 +135,90 @@ impl SessionTimezone {
None
}
// TODO: handle tstz-related arithmetic with timezone
ExprType::Add | ExprType::Subtract => None,
// We first translate to timestamp to handle years, months and days,
// then we translate back to timestamptz handle hours and milliseconds
//
// For performance concern, we assume that most the intervals are const-evaled.
//
// We impl the following expression tree:
//
// [+/-]
// / \
// timestamptz [-]
// / / \
// [+/-] interval date_trunc
// / \ / \
// timestamp date_trunc 'day' interval
// / / \
// timestamptz 'day' interval
//
//
// Const-evaled expr tree:
//
// [+/-]
// / \
// timestamptz interval_non_date_part
// /
// [+/-]
// / \
// timestamp interval_date_part
// /
// timestamptz
ExprType::Subtract | ExprType::Add => {
assert_eq!(inputs.len(), 2);
let canonical_match = matches!(inputs[0].return_type(), DataType::Timestamptz)
&& matches!(inputs[1].return_type(), DataType::Interval);
let inverse_match = matches!(inputs[1].return_type(), DataType::Timestamptz)
&& matches!(inputs[0].return_type(), DataType::Interval);
assert!(!(inverse_match && func_type == ExprType::Subtract)); // This should never have been parsed.
if canonical_match || inverse_match {
let (orig_timestamptz, interval) =
if func_type == ExprType::Add && inverse_match {
(inputs[1].clone(), inputs[0].clone())
} else {
(inputs[0].clone(), inputs[1].clone())
};
let interval_date_part: ExprImpl = FunctionCall::new_unchecked(
ExprType::DateTrunc,
vec![
Literal::new(Some("day".into()), DataType::Varchar).into(),
interval.clone(),
],
DataType::Interval,
)
.into();
let interval_non_date_part = FunctionCall::new_unchecked(
ExprType::Subtract,
vec![interval, interval_date_part.clone()],
DataType::Interval,
)
.into();
let timestamp = self
.with_timezone(ExprType::Cast, &vec![orig_timestamptz], DataType::Timestamp)
.unwrap();
let timestamp_op_date_part = FunctionCall::new_unchecked(
func_type,
vec![timestamp, interval_date_part],
DataType::Timestamp,
)
.into();
let timestamptz = self
.with_timezone(
ExprType::Cast,
&vec![timestamp_op_date_part],
DataType::Timestamptz,
)
.unwrap();
let timestamptz_op_non_date_part = FunctionCall::new_unchecked(
func_type,
vec![timestamptz, interval_non_date_part],
DataType::Timestamptz,
)
.into();
return Some(timestamptz_op_non_date_part);
}
None
}
_ => None,
}
}
Expand Down

0 comments on commit b60c0ca

Please sign in to comment.