Skip to content

Commit 89ca583

Browse files
Weijun-HWeijun-H
authored
feat: interval add timestamp (#5491)
* feat: interval add timestamp * add end-to-end tests * update interval test and coercion rule
1 parent 9e32de3 commit 89ca583

File tree

5 files changed

+129
-16
lines changed

5 files changed

+129
-16
lines changed

‎datafusion/common/src/scalar.rs

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -516,18 +516,34 @@ macro_rules! impl_op {
516516
let value = seconds_add(*ts_s, $RHS, get_sign!($OPERATION))?;
517517
Ok(ScalarValue::TimestampSecond(Some(value), zone.clone()))
518518
}
519+
(_, ScalarValue::TimestampSecond(Some(ts_s), zone)) => {
520+
let value = seconds_add(*ts_s, $LHS, get_sign!($OPERATION))?;
521+
Ok(ScalarValue::TimestampSecond(Some(value), zone.clone()))
522+
}
519523
(ScalarValue::TimestampMillisecond(Some(ts_ms), zone), _) => {
520524
let value = milliseconds_add(*ts_ms, $RHS, get_sign!($OPERATION))?;
521525
Ok(ScalarValue::TimestampMillisecond(Some(value), zone.clone()))
522526
}
527+
(_, ScalarValue::TimestampMillisecond(Some(ts_ms), zone)) => {
528+
let value = milliseconds_add(*ts_ms, $LHS, get_sign!($OPERATION))?;
529+
Ok(ScalarValue::TimestampMillisecond(Some(value), zone.clone()))
530+
}
523531
(ScalarValue::TimestampMicrosecond(Some(ts_us), zone), _) => {
524532
let value = microseconds_add(*ts_us, $RHS, get_sign!($OPERATION))?;
525533
Ok(ScalarValue::TimestampMicrosecond(Some(value), zone.clone()))
526534
}
535+
(_, ScalarValue::TimestampMicrosecond(Some(ts_us), zone)) => {
536+
let value = microseconds_add(*ts_us, $LHS, get_sign!($OPERATION))?;
537+
Ok(ScalarValue::TimestampMicrosecond(Some(value), zone.clone()))
538+
}
527539
(ScalarValue::TimestampNanosecond(Some(ts_ns), zone), _) => {
528540
let value = nanoseconds_add(*ts_ns, $RHS, get_sign!($OPERATION))?;
529541
Ok(ScalarValue::TimestampNanosecond(Some(value), zone.clone()))
530542
}
543+
(_, ScalarValue::TimestampNanosecond(Some(ts_ns), zone)) => {
544+
let value = nanoseconds_add(*ts_ns, $LHS, get_sign!($OPERATION))?;
545+
Ok(ScalarValue::TimestampNanosecond(Some(value), zone.clone()))
546+
}
531547
_ => Err(DataFusionError::Internal(format!(
532548
"Operator {} is not implemented for types {:?} and {:?}",
533549
stringify!($OPERATION),
@@ -2941,6 +2957,28 @@ mod tests {
29412957
Ok(())
29422958
}
29432959

2960+
#[test]
2961+
fn test_interval_add_timestamp() -> Result<()> {
2962+
let interval = ScalarValue::IntervalMonthDayNano(Some(123));
2963+
let timestamp = ScalarValue::TimestampNanosecond(Some(123), None);
2964+
let result = interval.add(&timestamp)?;
2965+
let expect = timestamp.add(&interval)?;
2966+
assert_eq!(result, expect);
2967+
2968+
let interval = ScalarValue::IntervalYearMonth(Some(123));
2969+
let timestamp = ScalarValue::TimestampNanosecond(Some(123), None);
2970+
let result = interval.add(&timestamp)?;
2971+
let expect = timestamp.add(&interval)?;
2972+
assert_eq!(result, expect);
2973+
2974+
let interval = ScalarValue::IntervalDayTime(Some(123));
2975+
let timestamp = ScalarValue::TimestampNanosecond(Some(123), None);
2976+
let result = interval.add(&timestamp)?;
2977+
let expect = timestamp.add(&interval)?;
2978+
assert_eq!(result, expect);
2979+
Ok(())
2980+
}
2981+
29442982
#[test]
29452983
fn scalar_decimal_test() -> Result<()> {
29462984
let decimal_value = ScalarValue::Decimal128(Some(123), 10, 1);

‎datafusion/core/tests/sqllogictests/test_files/timestamps.slt

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,3 +191,44 @@ query P
191191
SELECT DATE_TRUNC('second', '2022-08-03 14:38:50Z');
192192
----
193193
2022-08-03T14:38:50
194+
195+
# Test that interval can add a timestamp
196+
query P
197+
SELECT timestamp '2013-07-01 12:00:00' + INTERVAL '8' DAY;
198+
----
199+
2013-07-09T12:00:00
200+
201+
query P
202+
SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' DAY;
203+
----
204+
2000-01-09T00:00:00
205+
206+
query P
207+
SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' YEAR;
208+
----
209+
2008-01-01T00:00:00
210+
211+
query P
212+
SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' MONTH;
213+
----
214+
2000-09-01T00:00:00
215+
216+
query P
217+
SELECT INTERVAL '8' DAY + timestamp '2013-07-01 12:00:00';
218+
----
219+
2013-07-09T12:00:00
220+
221+
query P
222+
SELECT INTERVAL '8' DAY + '2000-01-01T00:00:00'::timestamp;
223+
----
224+
2000-01-09T00:00:00
225+
226+
query P
227+
SELECT INTERVAL '8' YEAR + '2000-01-01T00:00:00'::timestamp;
228+
----
229+
2008-01-01T00:00:00
230+
231+
query P
232+
SELECT INTERVAL '8' MONTH + '2000-01-01T00:00:00'::timestamp;
233+
----
234+
2000-09-01T00:00:00

‎datafusion/expr/src/type_coercion.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,11 @@ pub fn is_timestamp(dt: &DataType) -> bool {
6767
matches!(dt, DataType::Timestamp(_, _))
6868
}
6969

70+
/// Determine whether the given data type 'dt' is a `Interval`.
71+
pub fn is_interval(dt: &DataType) -> bool {
72+
matches!(dt, DataType::Interval(_))
73+
}
74+
7075
/// Determine whether the given data type `dt` is a `Date`.
7176
pub fn is_date(dt: &DataType) -> bool {
7277
matches!(dt, DataType::Date32 | DataType::Date64)

‎datafusion/expr/src/type_coercion/binary.rs

Lines changed: 35 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
//! Coercion rules for matching argument types for binary operators
1919
20-
use crate::type_coercion::{is_date, is_numeric, is_timestamp};
20+
use crate::type_coercion::{is_date, is_interval, is_numeric, is_timestamp};
2121
use crate::Operator;
2222
use arrow::compute::can_cast_types;
2323
use arrow::datatypes::{
@@ -114,22 +114,12 @@ pub fn coerce_types(
114114
| Operator::GtEq
115115
| Operator::LtEq => comparison_coercion(lhs_type, rhs_type),
116116
Operator::Plus | Operator::Minus
117-
if is_date(lhs_type) || is_timestamp(lhs_type) =>
117+
if is_date(lhs_type)
118+
|| is_date(rhs_type)
119+
|| is_timestamp(lhs_type)
120+
|| is_timestamp(rhs_type) =>
118121
{
119-
match rhs_type {
120-
// timestamp/date +/- interval returns timestamp/date
121-
DataType::Interval(_) => Some(lhs_type.clone()),
122-
// providing more helpful error message
123-
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) => {
124-
return Err(DataFusionError::Plan(
125-
format!(
126-
"'{lhs_type:?} {op} {rhs_type:?}' is an unsupported operation. \
127-
addition/subtraction on dates/timestamps only supported with interval types"
128-
),
129-
));
130-
}
131-
_ => None,
132-
}
122+
temporal_add_sub_coercion(lhs_type, rhs_type, op)?
133123
}
134124
// for math expressions, the final value of the coercion is also the return type
135125
// because coercion favours higher information types
@@ -214,6 +204,35 @@ pub fn comparison_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<D
214204
.or_else(|| string_numeric_coercion(lhs_type, rhs_type))
215205
}
216206

207+
/// Return the output type from performing addition or subtraction operations on temporal data types
208+
pub fn temporal_add_sub_coercion(
209+
lhs_type: &DataType,
210+
rhs_type: &DataType,
211+
op: &Operator,
212+
) -> Result<Option<DataType>> {
213+
// interval + date or timestamp
214+
if is_interval(lhs_type) && (is_date(rhs_type) || is_timestamp(rhs_type)) {
215+
return Ok(Some(rhs_type.clone()));
216+
}
217+
218+
// date or timestamp + interval
219+
if is_interval(rhs_type) && (is_date(lhs_type) || is_timestamp(lhs_type)) {
220+
return Ok(Some(lhs_type.clone()));
221+
}
222+
223+
// date or timestamp + date or timestamp
224+
if (is_date(lhs_type) || is_timestamp(lhs_type))
225+
&& (is_date(rhs_type) || is_timestamp(rhs_type))
226+
{
227+
return Err(DataFusionError::Plan(
228+
format!(
229+
"'{lhs_type:?} {op} {rhs_type:?}' is an unsupported operation. \
230+
addition/subtraction on dates/timestamps only supported with interval types"
231+
),));
232+
}
233+
Ok(None)
234+
}
235+
217236
/// Returns the output type of applying numeric operations such as `=`
218237
/// to arguments `lhs_type` and `rhs_type` if one is numeric and one
219238
/// is `Utf8`/`LargeUtf8`.

‎datafusion/physical-expr/src/planner.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,16 @@ pub fn create_physical_expr(
197197
rhs,
198198
input_schema,
199199
)?)),
200+
(
201+
DataType::Interval(_),
202+
Operator::Plus | Operator::Minus,
203+
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _),
204+
) => Ok(Arc::new(DateTimeIntervalExpr::try_new(
205+
rhs,
206+
*op,
207+
lhs,
208+
input_schema,
209+
)?)),
200210
_ => {
201211
// Note that the logical planner is responsible
202212
// for type coercion on the arguments (e.g. if one

0 commit comments

Comments
 (0)