Add Date64

apache · Jul 8, 2022 · 3093692 · 3093692
1 parent 17a4626
commit 3093692
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs
@@ -949,6 +949,7 @@ mod tests {
     use chrono::NaiveDate;
     use super::*;
     use crate::array::Int32Array;
+    use crate::datatypes::Date64Type;
 
     #[test]
     fn test_primitive_array_add() {
@@ -989,6 +990,33 @@ mod tests {
         assert_eq!(c.value(0), Date32Type::from_naive_date(NaiveDate::from_ymd(2001, 02, 01)));
     }
 
+    #[test]
+    fn test_date64_month_add() {
+        let a = Date64Array::from(vec![Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
+        let b = IntervalYearMonthArray::from(vec![IntervalYearMonthType::from(1, 1)]);
+        let c = add_dyn(&a, &b).unwrap();
+        let c = c.as_any().downcast_ref::<Date64Array>().unwrap();
+        assert_eq!(c.value(0), Date64Type::from_naive_date(NaiveDate::from_ymd(2001, 02, 01)));
+    }
+
+    #[test]
+    fn test_date64_day_time_add() {
+        let a = Date64Array::from(vec![Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
+        let b = IntervalDayTimeArray::from(vec![IntervalDayTimeType::from(1, 1)]);
+        let c = add_dyn(&a, &b).unwrap();
+        let c = c.as_any().downcast_ref::<Date64Array>().unwrap();
+        assert_eq!(c.value(0), Date64Type::from_naive_date(NaiveDate::from_ymd(2001, 01, 02)));
+    }
+
+    #[test]
+    fn test_date64_month_day_nano_add() {
+        let a = Date64Array::from(vec![Date64Type::from_naive_date(NaiveDate::from_ymd(2000, 01, 01))]);
+        let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::from(1, 1, 1)]);
+        let c = add_dyn(&a, &b).unwrap();
+        let c = c.as_any().downcast_ref::<Date64Array>().unwrap();
+        assert_eq!(c.value(0), Date64Type::from_naive_date(NaiveDate::from_ymd(2001, 02, 01)));
+    }
+
     #[test]
     fn test_primitive_array_add_dyn() {
         let a = Int32Array::from(vec![Some(5), Some(6), Some(7), Some(8), Some(9)]);

diff --git a/arrow/src/datatypes/types.rs b/arrow/src/datatypes/types.rs
@@ -291,4 +291,37 @@ impl Date64Type {
         let epoch = NaiveDate::from_ymd(1970, 1, 1);
         d.sub(epoch).num_milliseconds() as <Date64Type as ArrowPrimitiveType>::Native
     }
+
+    pub fn add_year_months(
+        date: <Date64Type as ArrowPrimitiveType>::Native,
+        delta: <IntervalYearMonthType as ArrowPrimitiveType>::Native
+    ) -> <Date64Type as ArrowPrimitiveType>::Native {
+        let prior = Date64Type::to_naive_date(date);
+        let months = IntervalYearMonthType::to_months(delta);
+        let posterior = shift_months(prior, months);
+        Date64Type::from_naive_date(posterior)
+    }
+
+    pub fn add_day_time(
+        date: <Date64Type as ArrowPrimitiveType>::Native,
+        delta: <IntervalDayTimeType as ArrowPrimitiveType>::Native
+    ) -> <Date64Type as ArrowPrimitiveType>::Native {
+        let (days, ms) = IntervalDayTimeType::to_parts(delta);
+        let res = Date64Type::to_naive_date(date);
+        let res = res.add(Duration::days(days as i64));
+        let res = res.add(Duration::milliseconds(ms as i64));
+        Date64Type::from_naive_date(res)
+    }
+
+    pub fn add_month_day_nano(
+        date: <Date64Type as ArrowPrimitiveType>::Native,
+        delta: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native
+    ) -> <Date64Type as ArrowPrimitiveType>::Native {
+        let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(delta);
+        let res = Date64Type::to_naive_date(date);
+        let res = shift_months(res, months);
+        let res = res.add(Duration::days(days as i64));
+        let res = res.add(Duration::nanoseconds(nanos));
+        Date64Type::from_naive_date(res)
+    }
 }