Skip to content

Commit

Permalink
Support Date - Date (apache#4383)
Browse files Browse the repository at this point in the history
  • Loading branch information
tustvold committed Jul 10, 2023
1 parent 8da2f97 commit 5070272
Showing 1 changed file with 49 additions and 1 deletion.
50 changes: 49 additions & 1 deletion arrow-arith/src/numeric.rs
Original file line number Diff line number Diff line change
Expand Up @@ -606,8 +606,32 @@ fn date_op<T: DateOp>(
use DataType::*;
use IntervalUnit::*;

let r_t = r.data_type();
match (T::DATA_TYPE, op, r_t) {
(Date32, Op::Sub | Op::SubWrapping, Date32) => {
let l = l.as_primitive::<Date32Type>();
let r = r.as_primitive::<Date32Type>();
return Ok(op_ref!(
DurationSecondType,
l,
l_s,
r,
r_s,
((l as i64) - (r as i64)) * 86400
));
}
(Date64, Op::Sub | Op::SubWrapping, Date64) => {
let l = l.as_primitive::<Date64Type>();
let r = r.as_primitive::<Date64Type>();
let result =
try_op_ref!(DurationMillisecondType, l, l_s, r, r_s, l.sub_checked(r));
return Ok(result);
}
_ => {}
}

let l = l.as_primitive::<T>();
match (op, r.data_type()) {
match (op, r_t) {
(Op::Add | Op::AddWrapping, Interval(YearMonth)) => {
let r = r.as_primitive::<IntervalYearMonthType>();
Ok(op_ref!(T, l, l_s, r, r_s, T::add_year_month(l, r)))
Expand Down Expand Up @@ -1421,5 +1445,29 @@ mod tests {
fn test_date() {
test_date_impl::<Date32Type, _>(Date32Type::from_naive_date);
test_date_impl::<Date64Type, _>(Date64Type::from_naive_date);

let a = Date32Array::from(vec![i32::MIN, i32::MAX, 23, 7684]);
let b = Date32Array::from(vec![i32::MIN, i32::MIN, -2, 45]);
let result = sub(&a, &b).unwrap();
assert_eq!(
result.as_primitive::<DurationSecondType>().values(),
&[0, 371085174288000, 2160000, 660009600]
);

let a = Date64Array::from(vec![4343, 76676, 3434]);
let b = Date64Array::from(vec![3, -5, 5]);
let result = sub(&a, &b).unwrap();
assert_eq!(
result.as_primitive::<DurationMillisecondType>().values(),
&[4340, 76681, 3429]
);

let a = Date64Array::from(vec![i64::MAX]);
let b = Date64Array::from(vec![-1]);
let err = sub(&a, &b).unwrap_err().to_string();
assert_eq!(
err,
"Compute error: Overflow happened on: 9223372036854775807 - -1"
);
}
}

0 comments on commit 5070272

Please sign in to comment.