This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 869e6bc13 add Date32/Date64 support to subtract_dyn (#3974)
869e6bc13 is described below
commit 869e6bc13bdbd503dddfa620a0ae62f3c0145a97
Author: Sinan Gençoğlu <[email protected]>
AuthorDate: Wed Mar 29 15:44:34 2023 +0300
add Date32/Date64 support to subtract_dyn (#3974)
* add Date32/Date64 support to subtract_dyn
* This commit includes:
1- add subtract date utilities to types
2- add test for subtract_dyn function
* fix formatting
---
arrow-arith/src/arithmetic.rs | 190 ++++++++++++++++++++++++++++++++++++++++++
arrow-array/src/types.rs | 102 +++++++++++++++++++++++
2 files changed, 292 insertions(+)
diff --git a/arrow-arith/src/arithmetic.rs b/arrow-arith/src/arithmetic.rs
index 501878afd..6c0fd497e 100644
--- a/arrow-arith/src/arithmetic.rs
+++ b/arrow-arith/src/arithmetic.rs
@@ -923,6 +923,54 @@ pub fn subtract_dyn(left: &dyn Array, right: &dyn Array)
-> Result<ArrayRef, Arr
DataType::Dictionary(_, _) => {
typed_dict_math_op!(left, right, |a, b| a.sub_wrapping(b),
math_op_dict)
}
+ DataType::Date32 => {
+ let l = left.as_primitive::<Date32Type>();
+ match right.data_type() {
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ let r = right.as_primitive::<IntervalYearMonthType>();
+ let res = math_op(l, r, Date32Type::subtract_year_months)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ let r = right.as_primitive::<IntervalDayTimeType>();
+ let res = math_op(l, r, Date32Type::subtract_day_time)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ let r = right.as_primitive::<IntervalMonthDayNanoType>();
+ let res = math_op(l, r,
Date32Type::subtract_month_day_nano)?;
+ Ok(Arc::new(res))
+ }
+ _ => Err(ArrowError::CastError(format!(
+ "Cannot perform arithmetic operation between array of type
{} and array of type {}",
+ left.data_type(), right.data_type()
+ ))),
+ }
+ }
+ DataType::Date64 => {
+ let l = left.as_primitive::<Date64Type>();
+ match right.data_type() {
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ let r = right.as_primitive::<IntervalYearMonthType>();
+ let res = math_op(l, r, Date64Type::subtract_year_months)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ let r = right.as_primitive::<IntervalDayTimeType>();
+ let res = math_op(l, r, Date64Type::subtract_day_time)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ let r = right.as_primitive::<IntervalMonthDayNanoType>();
+ let res = math_op(l, r,
Date64Type::subtract_month_day_nano)?;
+ Ok(Arc::new(res))
+ }
+ _ => Err(ArrowError::CastError(format!(
+ "Cannot perform arithmetic operation between array of type
{} and array of type {}",
+ left.data_type(), right.data_type()
+ ))),
+ }
+ }
_ => {
downcast_primitive_array!(
(left, right) => {
@@ -955,6 +1003,54 @@ pub fn subtract_dyn_checked(
math_checked_op_dict
)
}
+ DataType::Date32 => {
+ let l = left.as_primitive::<Date32Type>();
+ match right.data_type() {
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ let r = right.as_primitive::<IntervalYearMonthType>();
+ let res = math_op(l, r, Date32Type::subtract_year_months)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ let r = right.as_primitive::<IntervalDayTimeType>();
+ let res = math_op(l, r, Date32Type::subtract_day_time)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ let r = right.as_primitive::<IntervalMonthDayNanoType>();
+ let res = math_op(l, r,
Date32Type::subtract_month_day_nano)?;
+ Ok(Arc::new(res))
+ }
+ _ => Err(ArrowError::CastError(format!(
+ "Cannot perform arithmetic operation between array of type
{} and array of type {}",
+ left.data_type(), right.data_type()
+ ))),
+ }
+ }
+ DataType::Date64 => {
+ let l = left.as_primitive::<Date64Type>();
+ match right.data_type() {
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ let r = right.as_primitive::<IntervalYearMonthType>();
+ let res = math_op(l, r, Date64Type::subtract_year_months)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ let r = right.as_primitive::<IntervalDayTimeType>();
+ let res = math_op(l, r, Date64Type::subtract_day_time)?;
+ Ok(Arc::new(res))
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ let r = right.as_primitive::<IntervalMonthDayNanoType>();
+ let res = math_op(l, r,
Date64Type::subtract_month_day_nano)?;
+ Ok(Arc::new(res))
+ }
+ _ => Err(ArrowError::CastError(format!(
+ "Cannot perform arithmetic operation between array of type
{} and array of type {}",
+ left.data_type(), right.data_type()
+ ))),
+ }
+ }
_ => {
downcast_primitive_array!(
(left, right) => {
@@ -1864,6 +1960,100 @@ mod tests {
assert_eq!(1, c.value(4));
}
+ #[test]
+ fn test_date32_month_subtract() {
+ let a = Date32Array::from(vec![Date32Type::from_naive_date(
+ NaiveDate::from_ymd_opt(2000, 7, 1).unwrap(),
+ )]);
+ let b =
+
IntervalYearMonthArray::from(vec![IntervalYearMonthType::make_value(6, 3)]);
+ let c = subtract_dyn(&a, &b).unwrap();
+ let c = c.as_any().downcast_ref::<Date32Array>().unwrap();
+ assert_eq!(
+ c.value(0),
+ Date32Type::from_naive_date(NaiveDate::from_ymd_opt(1994, 4,
1).unwrap())
+ );
+ }
+
+ #[test]
+ fn test_date32_day_time_subtract() {
+ let a = Date32Array::from(vec![Date32Type::from_naive_date(
+ NaiveDate::from_ymd_opt(2023, 3, 29).unwrap(),
+ )]);
+ let b =
+ IntervalDayTimeArray::from(vec![IntervalDayTimeType::make_value(1,
86500)]);
+ let c = subtract_dyn(&a, &b).unwrap();
+ let c = c.as_any().downcast_ref::<Date32Array>().unwrap();
+ assert_eq!(
+ c.value(0),
+ Date32Type::from_naive_date(NaiveDate::from_ymd_opt(2023, 3,
27).unwrap())
+ );
+ }
+
+ #[test]
+ fn test_date32_month_day_nano_subtract() {
+ let a = Date32Array::from(vec![Date32Type::from_naive_date(
+ NaiveDate::from_ymd_opt(2023, 3, 15).unwrap(),
+ )]);
+ let b =
+
IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::make_value(
+ 1, 2, 0,
+ )]);
+ let c = subtract_dyn(&a, &b).unwrap();
+ let c = c.as_any().downcast_ref::<Date32Array>().unwrap();
+ assert_eq!(
+ c.value(0),
+ Date32Type::from_naive_date(NaiveDate::from_ymd_opt(2023, 2,
13).unwrap())
+ );
+ }
+
+ #[test]
+ fn test_date64_month_subtract() {
+ let a = Date64Array::from(vec![Date64Type::from_naive_date(
+ NaiveDate::from_ymd_opt(2000, 7, 1).unwrap(),
+ )]);
+ let b =
+
IntervalYearMonthArray::from(vec![IntervalYearMonthType::make_value(6, 3)]);
+ let c = subtract_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_opt(1994, 4,
1).unwrap())
+ );
+ }
+
+ #[test]
+ fn test_date64_day_time_subtract() {
+ let a = Date64Array::from(vec![Date64Type::from_naive_date(
+ NaiveDate::from_ymd_opt(2023, 3, 29).unwrap(),
+ )]);
+ let b =
+ IntervalDayTimeArray::from(vec![IntervalDayTimeType::make_value(1,
86500)]);
+ let c = subtract_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_opt(2023, 3,
27).unwrap())
+ );
+ }
+
+ #[test]
+ fn test_date64_month_day_nano_subtract() {
+ let a = Date64Array::from(vec![Date64Type::from_naive_date(
+ NaiveDate::from_ymd_opt(2023, 3, 15).unwrap(),
+ )]);
+ let b =
+
IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::make_value(
+ 1, 2, 0,
+ )]);
+ let c = subtract_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_opt(2023, 2,
13).unwrap())
+ );
+ }
+
#[test]
#[cfg(feature = "dyn_arith_dict")]
fn test_primitive_array_subtract_dyn_dict() {
diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs
index 9f1965b77..60a632a06 100644
--- a/arrow-array/src/types.rs
+++ b/arrow-array/src/types.rs
@@ -511,6 +511,57 @@ impl Date32Type {
let res = res.add(Duration::nanoseconds(nanos));
Date32Type::from_naive_date(res)
}
+
+ /// Subtract the given IntervalYearMonthType to an arrow Date32Type
+ ///
+ /// # Arguments
+ ///
+ /// * `date` - The date on which to perform the operation
+ /// * `delta` - The interval to substract
+ pub fn subtract_year_months(
+ date: <Date32Type as ArrowPrimitiveType>::Native,
+ delta: <IntervalYearMonthType as ArrowPrimitiveType>::Native,
+ ) -> <Date32Type as ArrowPrimitiveType>::Native {
+ let prior = Date32Type::to_naive_date(date);
+ let months = IntervalYearMonthType::to_months(-delta);
+ let posterior = shift_months(prior, months);
+ Date32Type::from_naive_date(posterior)
+ }
+
+ /// Subtract the given IntervalDayTimeType to an arrow Date32Type
+ ///
+ /// # Arguments
+ ///
+ /// * `date` - The date on which to perform the operation
+ /// * `delta` - The interval to subtract
+ pub fn subtract_day_time(
+ date: <Date32Type as ArrowPrimitiveType>::Native,
+ delta: <IntervalDayTimeType as ArrowPrimitiveType>::Native,
+ ) -> <Date32Type as ArrowPrimitiveType>::Native {
+ let (days, ms) = IntervalDayTimeType::to_parts(-delta);
+ let res = Date32Type::to_naive_date(date);
+ let res = res.add(Duration::days(days as i64));
+ let res = res.add(Duration::milliseconds(ms as i64));
+ Date32Type::from_naive_date(res)
+ }
+
+ /// Subtract the given IntervalMonthDayNanoType to an arrow Date32Type
+ ///
+ /// # Arguments
+ ///
+ /// * `date` - The date on which to perform the operation
+ /// * `delta` - The interval to subtract
+ pub fn subtract_month_day_nano(
+ date: <Date32Type as ArrowPrimitiveType>::Native,
+ delta: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native,
+ ) -> <Date32Type as ArrowPrimitiveType>::Native {
+ let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(delta);
+ let res = Date32Type::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));
+ Date32Type::from_naive_date(res)
+ }
}
impl Date64Type {
@@ -584,6 +635,57 @@ impl Date64Type {
let res = res.add(Duration::nanoseconds(nanos));
Date64Type::from_naive_date(res)
}
+
+ /// Subtract the given IntervalYearMonthType to an arrow Date64Type
+ ///
+ /// # Arguments
+ ///
+ /// * `date` - The date on which to perform the operation
+ /// * `delta` - The interval to substract
+ pub fn subtract_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)
+ }
+
+ /// Subtract the given IntervalDayTimeType to an arrow Date64Type
+ ///
+ /// # Arguments
+ ///
+ /// * `date` - The date on which to perform the operation
+ /// * `delta` - The interval to subtract
+ pub fn subtract_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)
+ }
+
+ /// Subtract the given IntervalMonthDayNanoType to an arrow Date64Type
+ ///
+ /// # Arguments
+ ///
+ /// * `date` - The date on which to perform the operation
+ /// * `delta` - The interval to subtract
+ pub fn subtract_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)
+ }
}
/// Crate private types for Decimal Arrays
