From 8e65992f3609b871d61c2df258311b43aeda955c Mon Sep 17 00:00:00 2001 From: Weijun Huang Date: Mon, 6 Mar 2023 22:46:27 +0100 Subject: [PATCH 1/3] feat: interval add timestamp --- datafusion/common/src/scalar.rs | 38 +++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs index c31b37b63b21..c094c10734d9 100644 --- a/datafusion/common/src/scalar.rs +++ b/datafusion/common/src/scalar.rs @@ -516,18 +516,34 @@ macro_rules! impl_op { let value = seconds_add(*ts_s, $RHS, get_sign!($OPERATION))?; Ok(ScalarValue::TimestampSecond(Some(value), zone.clone())) } + (_, ScalarValue::TimestampSecond(Some(ts_s), zone)) => { + let value = seconds_add(*ts_s, $LHS, get_sign!($OPERATION))?; + Ok(ScalarValue::TimestampSecond(Some(value), zone.clone())) + } (ScalarValue::TimestampMillisecond(Some(ts_ms), zone), _) => { let value = milliseconds_add(*ts_ms, $RHS, get_sign!($OPERATION))?; Ok(ScalarValue::TimestampMillisecond(Some(value), zone.clone())) } + (_, ScalarValue::TimestampMillisecond(Some(ts_ms), zone)) => { + let value = milliseconds_add(*ts_ms, $LHS, get_sign!($OPERATION))?; + Ok(ScalarValue::TimestampMillisecond(Some(value), zone.clone())) + } (ScalarValue::TimestampMicrosecond(Some(ts_us), zone), _) => { let value = microseconds_add(*ts_us, $RHS, get_sign!($OPERATION))?; Ok(ScalarValue::TimestampMicrosecond(Some(value), zone.clone())) } + (_, ScalarValue::TimestampMicrosecond(Some(ts_us), zone)) => { + let value = microseconds_add(*ts_us, $LHS, get_sign!($OPERATION))?; + Ok(ScalarValue::TimestampMicrosecond(Some(value), zone.clone())) + } (ScalarValue::TimestampNanosecond(Some(ts_ns), zone), _) => { let value = nanoseconds_add(*ts_ns, $RHS, get_sign!($OPERATION))?; Ok(ScalarValue::TimestampNanosecond(Some(value), zone.clone())) } + (_, ScalarValue::TimestampNanosecond(Some(ts_ns), zone)) => { + let value = nanoseconds_add(*ts_ns, $LHS, get_sign!($OPERATION))?; + Ok(ScalarValue::TimestampNanosecond(Some(value), zone.clone())) + } _ => Err(DataFusionError::Internal(format!( "Operator {} is not implemented for types {:?} and {:?}", stringify!($OPERATION), @@ -2911,6 +2927,28 @@ mod tests { Ok(()) } + #[test] + fn test_interval_add_timestamp() -> Result<()> { + let interval = ScalarValue::IntervalMonthDayNano(Some(123)); + let timestamp = ScalarValue::TimestampNanosecond(Some(123), None); + let result = interval.add(×tamp)?; + let expect = timestamp.add(&interval)?; + assert_eq!(result, expect); + + let interval = ScalarValue::IntervalYearMonth(Some(123)); + let timestamp = ScalarValue::TimestampNanosecond(Some(123), None); + let result = interval.add(×tamp)?; + let expect = timestamp.add(&interval)?; + assert_eq!(result, expect); + + let interval = ScalarValue::IntervalDayTime(Some(123)); + let timestamp = ScalarValue::TimestampNanosecond(Some(123), None); + let result = interval.add(×tamp)?; + let expect = timestamp.add(&interval)?; + assert_eq!(result, expect); + Ok(()) + } + #[test] fn scalar_decimal_test() -> Result<()> { let decimal_value = ScalarValue::Decimal128(Some(123), 10, 1); From 269790a87a35f46502fcaf4e3cf67aa841e53a51 Mon Sep 17 00:00:00 2001 From: Weijun Huang Date: Tue, 7 Mar 2023 14:08:55 +0100 Subject: [PATCH 2/3] add end-to-end tests --- .../sqllogictests/test_files/timestamps.slt | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/datafusion/core/tests/sqllogictests/test_files/timestamps.slt b/datafusion/core/tests/sqllogictests/test_files/timestamps.slt index dce6213f7eb4..e20536e68082 100644 --- a/datafusion/core/tests/sqllogictests/test_files/timestamps.slt +++ b/datafusion/core/tests/sqllogictests/test_files/timestamps.slt @@ -191,3 +191,24 @@ query P SELECT DATE_TRUNC('second', '2022-08-03 14:38:50Z'); ---- 2022-08-03T14:38:50 + +# Test that interval can add a timestamp +query P +SELECT timestamp '2013-07-01 12:00:00' + INTERVAL '8' DAY; +---- +2013-07-09T12:00:00 + +query P +SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' DAY; +---- +2000-01-09T00:00:00 + +query P +SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' YEAR; +---- +2008-01-01T00:00:00 + +query P +SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' MONTH; +---- +2000-09-01T00:00:00 \ No newline at end of file From fbe6c7cc10ef211cf0d6f3894eba56466adb732d Mon Sep 17 00:00:00 2001 From: Weijun Huang Date: Tue, 7 Mar 2023 19:32:51 +0100 Subject: [PATCH 3/3] update interval test and coercion rule --- .../sqllogictests/test_files/timestamps.slt | 20 ++++++++ datafusion/expr/src/type_coercion.rs | 5 ++ datafusion/expr/src/type_coercion/binary.rs | 51 +++++++++++++------ datafusion/physical-expr/src/planner.rs | 10 ++++ 4 files changed, 70 insertions(+), 16 deletions(-) diff --git a/datafusion/core/tests/sqllogictests/test_files/timestamps.slt b/datafusion/core/tests/sqllogictests/test_files/timestamps.slt index e20536e68082..1958adea13ea 100644 --- a/datafusion/core/tests/sqllogictests/test_files/timestamps.slt +++ b/datafusion/core/tests/sqllogictests/test_files/timestamps.slt @@ -211,4 +211,24 @@ SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' YEAR; query P SELECT '2000-01-01T00:00:00'::timestamp + INTERVAL '8' MONTH; ---- +2000-09-01T00:00:00 + +query P +SELECT INTERVAL '8' DAY + timestamp '2013-07-01 12:00:00'; +---- +2013-07-09T12:00:00 + +query P +SELECT INTERVAL '8' DAY + '2000-01-01T00:00:00'::timestamp; +---- +2000-01-09T00:00:00 + +query P +SELECT INTERVAL '8' YEAR + '2000-01-01T00:00:00'::timestamp; +---- +2008-01-01T00:00:00 + +query P +SELECT INTERVAL '8' MONTH + '2000-01-01T00:00:00'::timestamp; +---- 2000-09-01T00:00:00 \ No newline at end of file diff --git a/datafusion/expr/src/type_coercion.rs b/datafusion/expr/src/type_coercion.rs index 373f4a7a2d4a..740b10e28547 100644 --- a/datafusion/expr/src/type_coercion.rs +++ b/datafusion/expr/src/type_coercion.rs @@ -67,6 +67,11 @@ pub fn is_timestamp(dt: &DataType) -> bool { matches!(dt, DataType::Timestamp(_, _)) } +/// Determine whether the given data type 'dt' is a `Interval`. +pub fn is_interval(dt: &DataType) -> bool { + matches!(dt, DataType::Interval(_)) +} + /// Determine whether the given data type `dt` is a `Date`. pub fn is_date(dt: &DataType) -> bool { matches!(dt, DataType::Date32 | DataType::Date64) diff --git a/datafusion/expr/src/type_coercion/binary.rs b/datafusion/expr/src/type_coercion/binary.rs index ba24f4a901d5..70ca44139fde 100644 --- a/datafusion/expr/src/type_coercion/binary.rs +++ b/datafusion/expr/src/type_coercion/binary.rs @@ -17,7 +17,7 @@ //! Coercion rules for matching argument types for binary operators -use crate::type_coercion::{is_date, is_numeric, is_timestamp}; +use crate::type_coercion::{is_date, is_interval, is_numeric, is_timestamp}; use crate::Operator; use arrow::compute::can_cast_types; use arrow::datatypes::{ @@ -114,22 +114,12 @@ pub fn coerce_types( | Operator::GtEq | Operator::LtEq => comparison_coercion(lhs_type, rhs_type), Operator::Plus | Operator::Minus - if is_date(lhs_type) || is_timestamp(lhs_type) => + if is_date(lhs_type) + || is_date(rhs_type) + || is_timestamp(lhs_type) + || is_timestamp(rhs_type) => { - match rhs_type { - // timestamp/date +/- interval returns timestamp/date - DataType::Interval(_) => Some(lhs_type.clone()), - // providing more helpful error message - DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) => { - return Err(DataFusionError::Plan( - format!( - "'{lhs_type:?} {op} {rhs_type:?}' is an unsupported operation. \ - addition/subtraction on dates/timestamps only supported with interval types" - ), - )); - } - _ => None, - } + temporal_add_sub_coercion(lhs_type, rhs_type, op)? } // for math expressions, the final value of the coercion is also the return type // because coercion favours higher information types @@ -199,6 +189,35 @@ pub fn comparison_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option Result> { + // interval + date or timestamp + if is_interval(lhs_type) && (is_date(rhs_type) || is_timestamp(rhs_type)) { + return Ok(Some(rhs_type.clone())); + } + + // date or timestamp + interval + if is_interval(rhs_type) && (is_date(lhs_type) || is_timestamp(lhs_type)) { + return Ok(Some(lhs_type.clone())); + } + + // date or timestamp + date or timestamp + if (is_date(lhs_type) || is_timestamp(lhs_type)) + && (is_date(rhs_type) || is_timestamp(rhs_type)) + { + return Err(DataFusionError::Plan( + format!( + "'{lhs_type:?} {op} {rhs_type:?}' is an unsupported operation. \ + addition/subtraction on dates/timestamps only supported with interval types" + ),)); + } + Ok(None) +} + /// Returns the output type of applying numeric operations such as `=` /// to arguments `lhs_type` and `rhs_type` if one is numeric and one /// is `Utf8`/`LargeUtf8`. diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 105de17c2e74..1fbd73b3ba01 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -197,6 +197,16 @@ pub fn create_physical_expr( rhs, input_schema, )?)), + ( + DataType::Interval(_), + Operator::Plus | Operator::Minus, + DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _), + ) => Ok(Arc::new(DateTimeIntervalExpr::try_new( + rhs, + *op, + lhs, + input_schema, + )?)), _ => { // Note that the logical planner is responsible // for type coercion on the arguments (e.g. if one