feat: type coercion support date - date (#6578)

apache · Jun 8, 2023 · 9c13a77 · 9c13a77
1 parent 1d3860d
commit 9c13a77
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 21 deletions.
diff --git a/datafusion/core/tests/sqllogictests/test_files/dates.slt b/datafusion/core/tests/sqllogictests/test_files/dates.slt
@@ -90,7 +90,8 @@ select i_item_desc from test
 where d3_date > now() + '5 days';
 
 # DATE minus DATE
-query error DataFusion error: Error during planning: Unsupported argument types\. Can not evaluate Date32 \- Date32
+# https://github.com/apache/arrow-rs/issues/4383
+query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nArrow error: Cast error: Cannot perform arithmetic operation between array of type Date32 and array of type Date32
 SELECT DATE '2023-04-09' - DATE '2023-04-02';
 
 # DATE minus Timestamp

diff --git a/datafusion/expr/src/type_coercion/binary.rs b/datafusion/expr/src/type_coercion/binary.rs
@@ -55,8 +55,7 @@ fn mathematics_temporal_result_type(
         (Interval(_), Date64) => Some(rhs_type.clone()),
         (Date64, Interval(_)) => Some(lhs_type.clone()),
         // interval +/-
-        (Interval(YearMonth), Interval(YearMonth)) => Some(Interval(YearMonth)),
-        (Interval(DayTime), Interval(DayTime)) => Some(Interval(DayTime)),
+        (Interval(l), Interval(h)) if l == h => Some(lhs_type.clone()),
         (Interval(_), Interval(_)) => Some(Interval(MonthDayNano)),
         // timestamp - timestamp
         (Timestamp(Second, _), Timestamp(Second, _))
@@ -68,7 +67,10 @@ fn mathematics_temporal_result_type(
             Some(Interval(MonthDayNano))
         }
         (Timestamp(_, _), Timestamp(_, _)) => None,
-        // TODO: date minus date
+        // date - date
+        (Date32, Date32) => Some(Interval(DayTime)),
+        (Date64, Date64) => Some(Interval(MonthDayNano)),
+        (Date32, Date64) | (Date64, Date32) => Some(Interval(MonthDayNano)),
         // date - timestamp, timestamp - date
         (Date32, Timestamp(_, _))
         | (Timestamp(_, _), Date32)
@@ -733,10 +735,13 @@ fn is_time_with_valid_unit(datatype: DataType) -> bool {
 fn temporal_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataType> {
     use arrow::datatypes::DataType::*;
     use arrow::datatypes::IntervalUnit::*;
+    use arrow::datatypes::TimeUnit::*;
+
+    if lhs_type == rhs_type {
+        return Some(lhs_type.clone());
+    }
     match (lhs_type, rhs_type) {
         // interval +/-
-        (Interval(YearMonth), Interval(YearMonth)) => Some(Interval(YearMonth)),
-        (Interval(DayTime), Interval(DayTime)) => Some(Interval(DayTime)),
         (Interval(_), Interval(_)) => Some(Interval(MonthDayNano)),
         (Date64, Date32) | (Date32, Date64) => Some(Date64),
         (Utf8, Date32) | (Date32, Utf8) => Some(Date32),
@@ -754,13 +759,13 @@ fn temporal_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataTyp
             }
         }
         (Timestamp(_, tz), Utf8) | (Utf8, Timestamp(_, tz)) => {
-            Some(Timestamp(TimeUnit::Nanosecond, tz.clone()))
+            Some(Timestamp(Nanosecond, tz.clone()))
         }
         (Timestamp(_, None), Date32) | (Date32, Timestamp(_, None)) => {
-            Some(Timestamp(TimeUnit::Nanosecond, None))
+            Some(Timestamp(Nanosecond, None))
         }
         (Timestamp(_, _tz), Date32) | (Date32, Timestamp(_, _tz)) => {
-            Some(Timestamp(TimeUnit::Nanosecond, None))
+            Some(Timestamp(Nanosecond, None))
         }
         (Timestamp(lhs_unit, lhs_tz), Timestamp(rhs_unit, rhs_tz)) => {
             let tz = match (lhs_tz, rhs_tz) {
@@ -778,18 +783,18 @@ fn temporal_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataTyp
             };
 
             let unit = match (lhs_unit, rhs_unit) {
-                (TimeUnit::Second, TimeUnit::Millisecond) => TimeUnit::Second,
-                (TimeUnit::Second, TimeUnit::Microsecond) => TimeUnit::Second,
-                (TimeUnit::Second, TimeUnit::Nanosecond) => TimeUnit::Second,
-                (TimeUnit::Millisecond, TimeUnit::Second) => TimeUnit::Second,
-                (TimeUnit::Millisecond, TimeUnit::Microsecond) => TimeUnit::Millisecond,
-                (TimeUnit::Millisecond, TimeUnit::Nanosecond) => TimeUnit::Millisecond,
-                (TimeUnit::Microsecond, TimeUnit::Second) => TimeUnit::Second,
-                (TimeUnit::Microsecond, TimeUnit::Millisecond) => TimeUnit::Millisecond,
-                (TimeUnit::Microsecond, TimeUnit::Nanosecond) => TimeUnit::Microsecond,
-                (TimeUnit::Nanosecond, TimeUnit::Second) => TimeUnit::Second,
-                (TimeUnit::Nanosecond, TimeUnit::Millisecond) => TimeUnit::Millisecond,
-                (TimeUnit::Nanosecond, TimeUnit::Microsecond) => TimeUnit::Microsecond,
+                (Second, Millisecond) => Second,
+                (Second, Microsecond) => Second,
+                (Second, Nanosecond) => Second,
+                (Millisecond, Second) => Second,
+                (Millisecond, Microsecond) => Millisecond,
+                (Millisecond, Nanosecond) => Millisecond,
+                (Microsecond, Second) => Second,
+                (Microsecond, Millisecond) => Millisecond,
+                (Microsecond, Nanosecond) => Microsecond,
+                (Nanosecond, Second) => Second,
+                (Nanosecond, Millisecond) => Millisecond,
+                (Nanosecond, Microsecond) => Microsecond,
                 (l, r) => {
                     assert_eq!(l, r);
                     l.clone()