Dyn comparison of interval arrays (apache#1106)

tustvold · Dec 29, 2021 · 3517083 · 3517083
1 parent 3dca969
commit 3517083
Showing 1 changed file with 89 additions and 1 deletion.
diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs
@@ -28,7 +28,8 @@ use crate::compute::binary_boolean_kernel;
 use crate::compute::util::combine_option_bitmap;
 use crate::datatypes::{
     ArrowNumericType, DataType, Float32Type, Float64Type, Int16Type, Int32Type,
-    Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+    Int64Type, Int8Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit,
+    IntervalYearMonthType, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
 };
 use crate::error::{ArrowError, Result};
 use crate::util::bit_util;
@@ -1136,6 +1137,42 @@ macro_rules! typed_compares {
             (DataType::LargeUtf8, DataType::LargeUtf8) => {
                 typed_cmp!($LEFT, $RIGHT, LargeStringArray, $OP_STR, i64)
             }
+            (
+                DataType::Interval(IntervalUnit::YearMonth),
+                DataType::Interval(IntervalUnit::YearMonth),
+            ) => {
+                typed_cmp!(
+                    $LEFT,
+                    $RIGHT,
+                    IntervalYearMonthArray,
+                    $OP_PRIM,
+                    IntervalYearMonthType
+                )
+            }
+            (
+                DataType::Interval(IntervalUnit::DayTime),
+                DataType::Interval(IntervalUnit::DayTime),
+            ) => {
+                typed_cmp!(
+                    $LEFT,
+                    $RIGHT,
+                    IntervalDayTimeArray,
+                    $OP_PRIM,
+                    IntervalDayTimeType
+                )
+            }
+            (
+                DataType::Interval(IntervalUnit::MonthDayNano),
+                DataType::Interval(IntervalUnit::MonthDayNano),
+            ) => {
+                typed_cmp!(
+                    $LEFT,
+                    $RIGHT,
+                    IntervalMonthDayNanoArray,
+                    $OP_PRIM,
+                    IntervalMonthDayNanoType
+                )
+            }
             (t1, t2) if t1 == t2 => Err(ArrowError::NotYetImplemented(format!(
                 "Comparing arrays of type {} is not yet implemented",
                 t1
@@ -2025,6 +2062,57 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_interval_array() {
+        let a = IntervalDayTimeArray::from(
+            vec![Some(0), Some(6), Some(834), None, Some(3), None],
+        );
+        let b = IntervalDayTimeArray::from(
+            vec![Some(70), Some(6), Some(833), Some(6), Some(3), None],
+        );
+        let res = eq(&a, &b).unwrap();
+        let res_dyn = eq_dyn(&a, &b).unwrap();
+        assert_eq!(res, res_dyn);
+        assert_eq!(
+            &res_dyn,
+            &BooleanArray::from(
+                vec![Some(false), Some(true), Some(false), None, Some(true), None]
+            )
+        );
+
+        let a = IntervalMonthDayNanoArray::from(
+            vec![Some(0), Some(6), Some(834), None, Some(3), None],
+        );
+        let b = IntervalMonthDayNanoArray::from(
+            vec![Some(86), Some(5), Some(8), Some(6), Some(3), None],
+        );
+        let res = lt(&a, &b).unwrap();
+        let res_dyn = lt_dyn(&a, &b).unwrap();
+        assert_eq!(res, res_dyn);
+        assert_eq!(
+            &res_dyn,
+            &BooleanArray::from(
+                vec![Some(true), Some(false), Some(false), None, Some(false), None]
+            )
+        );
+
+        let a = IntervalYearMonthArray::from(
+            vec![Some(0), Some(623), Some(834), None, Some(3), None],
+        );
+        let b = IntervalYearMonthArray::from(
+            vec![Some(86), Some(5), Some(834), Some(6), Some(86), None],
+        );
+        let res = gt_eq(&a, &b).unwrap();
+        let res_dyn = gt_eq_dyn(&a, &b).unwrap();
+        assert_eq!(res, res_dyn);
+        assert_eq!(
+            &res_dyn,
+            &BooleanArray::from(
+                vec![Some(false), Some(true), Some(true), None, Some(false), None]
+            )
+        );
+    }
+
     // Expected behaviour:
     // contains("ab", ["ab", "cd", null]) = true
     // contains("ef", ["ab", "cd", null]) = false