update interval test and coercion rule

apache · Mar 7, 2023 · 4b3b94a · 4b3b94a
1 parent 269790a
commit 4b3b94a
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 16 deletions.
diff --git 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
diff --git 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
@@ -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,9 @@ 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_interval(lhs_type) || is_interval(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 +186,28 @@ pub fn comparison_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<D
         .or_else(|| string_numeric_coercion(lhs_type, rhs_type))
 }
 
+/// Return the output type from performing addition or subtraction operations on temporal data types
+pub fn temporal_add_sub_coercion(
+    lhs_type: &DataType,
+    rhs_type: &DataType,
+    op: &Operator,
+) -> Result<Option<DataType>> {
+    if is_interval(lhs_type) && (is_date(rhs_type) || is_timestamp(rhs_type)) {
+        return Ok(Some(rhs_type.clone()));
+    } else if is_interval(rhs_type) && (is_date(lhs_type) || is_timestamp(lhs_type)) {
+        return Ok(Some(lhs_type.clone()));
+    }
+
+    if lhs_type == rhs_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
@@ -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