Duration support

apache · Jul 4, 2023 · f402dd3 · f402dd3
1 parent ac4fd3f
commit f402dd3
Showing 1 changed file with 34 additions and 0 deletions.
diff --git a/arrow-arith/src/operation.rs b/arrow-arith/src/operation.rs
@@ -126,6 +126,10 @@ fn arithmetic_op(
         (Timestamp(Millisecond, _), _) => timestamp_op::<TimestampMillisecondType>(op, l, l_scalar, r, r_scalar),
         (Timestamp(Microsecond, _), _) => timestamp_op::<TimestampMicrosecondType>(op, l, l_scalar, r, r_scalar),
         (Timestamp(Nanosecond, _), _) => timestamp_op::<TimestampNanosecondType>(op, l, l_scalar, r, r_scalar),
+        (Duration(Second), Duration(Second)) => duration_op::<DurationSecondType>(op, l, l_scalar, r, r_scalar),
+        (Duration(Millisecond), Duration(Millisecond)) => duration_op::<DurationMillisecondType>(op, l, l_scalar, r, r_scalar),
+        (Duration(Microsecond), Duration(Microsecond)) => duration_op::<DurationMicrosecondType>(op, l, l_scalar, r, r_scalar),
+        (Duration(Nanosecond), Duration(Nanosecond)) => duration_op::<DurationNanosecondType>(op, l, l_scalar, r, r_scalar),
         (Interval(YearMonth), Interval(YearMonth)) => interval_op::<IntervalYearMonthType>(op, l, l_scalar, r, r_scalar),
         (Interval(DayTime), Interval(DayTime)) => interval_op::<IntervalDayTimeType>(op, l, l_scalar, r, r_scalar),
         (Interval(MonthDayNano), Interval(MonthDayNano)) => interval_op::<IntervalMonthDayNanoType>(op, l, l_scalar, r, r_scalar),
@@ -299,6 +303,16 @@ fn timestamp_op<T: TimestampOp>(
             let r = r.as_primitive::<T>();
             return Ok(try_op_ref!(T::Duration, l, l_s, r, r_s, l.sub_checked(r)));
         }
+
+        (Op::Add | Op::AddWrapping, Duration(unit)) if unit == &T::UNIT => {
+            let r = r.as_primitive::<T::Duration>();
+            try_op!(l, l_s, r, r_s, l.add_checked(r))
+        }
+        (Op::Sub | Op::SubWrapping, Duration(unit)) if unit == &T::UNIT => {
+            let r = r.as_primitive::<T::Duration>();
+            try_op!(l, l_s, r, r_s, l.sub_checked(r))
+        }
+
         (Op::Add | Op::AddWrapping, Interval(YearMonth)) => {
             let r = r.as_primitive::<IntervalYearMonthType>();
             try_op!(l, l_s, r, r_s, T::add_year_month(l, r))
@@ -456,6 +470,26 @@ fn interval_op<T: IntervalOp>(
     }
 }
 
+fn duration_op<T: ArrowPrimitiveType>(
+    op: Op,
+    l: &dyn Array,
+    l_s: bool,
+    r: &dyn Array,
+    r_s: bool,
+) -> Result<ArrayRef, ArrowError> {
+    let l = l.as_primitive::<T>();
+    let r = r.as_primitive::<T>();
+    match op {
+        Op::Add | Op::AddWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, l.add_checked(r))),
+        Op::Sub | Op::SubWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, l.sub_checked(r))),
+        _ => Err(ArrowError::InvalidArgumentError(format!(
+            "Invalid duration arithmetic operation: {} {op:?} {}",
+            l.data_type(),
+            r.data_type()
+        ))),
+    }
+}
+
 /// Perform arithmetic operation on a date array
 fn date_op<T: DateOp>(
     op: Op,