1use std::cmp::Ordering;
21use std::fmt::Formatter;
22use std::sync::Arc;
23
24use arrow_array::cast::AsArray;
25use arrow_array::timezone::Tz;
26use arrow_array::types::*;
27use arrow_array::*;
28use arrow_buffer::{ArrowNativeType, IntervalDayTime, IntervalMonthDayNano};
29use arrow_schema::{ArrowError, DataType, IntervalUnit, TimeUnit};
30
31use crate::arity::{binary, try_binary};
32
33pub fn add(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
35 arithmetic_op(Op::Add, lhs, rhs)
36}
37
38pub fn add_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
40 arithmetic_op(Op::AddWrapping, lhs, rhs)
41}
42
43pub fn sub(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
45 arithmetic_op(Op::Sub, lhs, rhs)
46}
47
48pub fn sub_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
50 arithmetic_op(Op::SubWrapping, lhs, rhs)
51}
52
53pub fn mul(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
55 arithmetic_op(Op::Mul, lhs, rhs)
56}
57
58pub fn mul_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
60 arithmetic_op(Op::MulWrapping, lhs, rhs)
61}
62
63pub fn div(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
68 arithmetic_op(Op::Div, lhs, rhs)
69}
70
71pub fn rem(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
78 arithmetic_op(Op::Rem, lhs, rhs)
79}
80
81macro_rules! neg_checked {
82 ($t:ty, $a:ident) => {{
83 let array = $a
84 .as_primitive::<$t>()
85 .try_unary::<_, $t, _>(|x| x.neg_checked())?;
86 Ok(Arc::new(array))
87 }};
88}
89
90macro_rules! neg_wrapping {
91 ($t:ty, $a:ident) => {{
92 let array = $a.as_primitive::<$t>().unary::<_, $t>(|x| x.neg_wrapping());
93 Ok(Arc::new(array))
94 }};
95}
96
97pub fn neg(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
102 use DataType::*;
103 use IntervalUnit::*;
104 use TimeUnit::*;
105
106 match array.data_type() {
107 Int8 => neg_checked!(Int8Type, array),
108 Int16 => neg_checked!(Int16Type, array),
109 Int32 => neg_checked!(Int32Type, array),
110 Int64 => neg_checked!(Int64Type, array),
111 Float16 => neg_wrapping!(Float16Type, array),
112 Float32 => neg_wrapping!(Float32Type, array),
113 Float64 => neg_wrapping!(Float64Type, array),
114 Decimal32(p, s) => {
115 let a = array
116 .as_primitive::<Decimal32Type>()
117 .try_unary::<_, Decimal32Type, _>(|x| x.neg_checked())?;
118
119 Ok(Arc::new(a.with_precision_and_scale(*p, *s)?))
120 }
121 Decimal64(p, s) => {
122 let a = array
123 .as_primitive::<Decimal64Type>()
124 .try_unary::<_, Decimal64Type, _>(|x| x.neg_checked())?;
125
126 Ok(Arc::new(a.with_precision_and_scale(*p, *s)?))
127 }
128 Decimal128(p, s) => {
129 let a = array
130 .as_primitive::<Decimal128Type>()
131 .try_unary::<_, Decimal128Type, _>(|x| x.neg_checked())?;
132
133 Ok(Arc::new(a.with_precision_and_scale(*p, *s)?))
134 }
135 Decimal256(p, s) => {
136 let a = array
137 .as_primitive::<Decimal256Type>()
138 .try_unary::<_, Decimal256Type, _>(|x| x.neg_checked())?;
139
140 Ok(Arc::new(a.with_precision_and_scale(*p, *s)?))
141 }
142 Duration(Second) => neg_checked!(DurationSecondType, array),
143 Duration(Millisecond) => neg_checked!(DurationMillisecondType, array),
144 Duration(Microsecond) => neg_checked!(DurationMicrosecondType, array),
145 Duration(Nanosecond) => neg_checked!(DurationNanosecondType, array),
146 Interval(YearMonth) => neg_checked!(IntervalYearMonthType, array),
147 Interval(DayTime) => {
148 let a = array
149 .as_primitive::<IntervalDayTimeType>()
150 .try_unary::<_, IntervalDayTimeType, ArrowError>(|x| {
151 let (days, ms) = IntervalDayTimeType::to_parts(x);
152 Ok(IntervalDayTimeType::make_value(
153 days.neg_checked()?,
154 ms.neg_checked()?,
155 ))
156 })?;
157 Ok(Arc::new(a))
158 }
159 Interval(MonthDayNano) => {
160 let a = array
161 .as_primitive::<IntervalMonthDayNanoType>()
162 .try_unary::<_, IntervalMonthDayNanoType, ArrowError>(|x| {
163 let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(x);
164 Ok(IntervalMonthDayNanoType::make_value(
165 months.neg_checked()?,
166 days.neg_checked()?,
167 nanos.neg_checked()?,
168 ))
169 })?;
170 Ok(Arc::new(a))
171 }
172 t => Err(ArrowError::InvalidArgumentError(format!(
173 "Invalid arithmetic operation: !{t}"
174 ))),
175 }
176}
177
178pub fn neg_wrapping(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
180 downcast_integer! {
181 array.data_type() => (neg_wrapping, array),
182 _ => neg(array),
183 }
184}
185
186#[derive(Debug, Copy, Clone)]
190enum Op {
191 AddWrapping,
192 Add,
193 SubWrapping,
194 Sub,
195 MulWrapping,
196 Mul,
197 Div,
198 Rem,
199}
200
201impl std::fmt::Display for Op {
202 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
203 match self {
204 Op::AddWrapping | Op::Add => write!(f, "+"),
205 Op::SubWrapping | Op::Sub => write!(f, "-"),
206 Op::MulWrapping | Op::Mul => write!(f, "*"),
207 Op::Div => write!(f, "/"),
208 Op::Rem => write!(f, "%"),
209 }
210 }
211}
212
213impl Op {
214 fn commutative(&self) -> bool {
215 matches!(self, Self::Add | Self::AddWrapping)
216 }
217}
218
219fn arithmetic_op(op: Op, lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
221 use DataType::*;
222 use IntervalUnit::*;
223 use TimeUnit::*;
224
225 macro_rules! integer_helper {
226 ($t:ty, $op:ident, $l:ident, $l_scalar:ident, $r:ident, $r_scalar:ident) => {
227 integer_op::<$t>($op, $l, $l_scalar, $r, $r_scalar)
228 };
229 }
230
231 let (l, l_scalar) = lhs.get();
232 let (r, r_scalar) = rhs.get();
233 downcast_integer! {
234 l.data_type(), r.data_type() => (integer_helper, op, l, l_scalar, r, r_scalar),
235 (Float16, Float16) => float_op::<Float16Type>(op, l, l_scalar, r, r_scalar),
236 (Float32, Float32) => float_op::<Float32Type>(op, l, l_scalar, r, r_scalar),
237 (Float64, Float64) => float_op::<Float64Type>(op, l, l_scalar, r, r_scalar),
238 (Timestamp(Second, _), _) => timestamp_op::<TimestampSecondType>(op, l, l_scalar, r, r_scalar),
239 (Timestamp(Millisecond, _), _) => timestamp_op::<TimestampMillisecondType>(op, l, l_scalar, r, r_scalar),
240 (Timestamp(Microsecond, _), _) => timestamp_op::<TimestampMicrosecondType>(op, l, l_scalar, r, r_scalar),
241 (Timestamp(Nanosecond, _), _) => timestamp_op::<TimestampNanosecondType>(op, l, l_scalar, r, r_scalar),
242 (Duration(Second), Duration(Second)) => duration_op::<DurationSecondType>(op, l, l_scalar, r, r_scalar),
243 (Duration(Millisecond), Duration(Millisecond)) => duration_op::<DurationMillisecondType>(op, l, l_scalar, r, r_scalar),
244 (Duration(Microsecond), Duration(Microsecond)) => duration_op::<DurationMicrosecondType>(op, l, l_scalar, r, r_scalar),
245 (Duration(Nanosecond), Duration(Nanosecond)) => duration_op::<DurationNanosecondType>(op, l, l_scalar, r, r_scalar),
246 (Interval(YearMonth), Interval(YearMonth)) => interval_op::<IntervalYearMonthType>(op, l, l_scalar, r, r_scalar),
247 (Interval(DayTime), Interval(DayTime)) => interval_op::<IntervalDayTimeType>(op, l, l_scalar, r, r_scalar),
248 (Interval(MonthDayNano), Interval(MonthDayNano)) => interval_op::<IntervalMonthDayNanoType>(op, l, l_scalar, r, r_scalar),
249 (Date32, _) => date_op::<Date32Type>(op, l, l_scalar, r, r_scalar),
250 (Date64, _) => date_op::<Date64Type>(op, l, l_scalar, r, r_scalar),
251 (Decimal32(_, _), Decimal32(_, _)) => decimal_op::<Decimal32Type>(op, l, l_scalar, r, r_scalar),
252 (Decimal64(_, _), Decimal64(_, _)) => decimal_op::<Decimal64Type>(op, l, l_scalar, r, r_scalar),
253 (Decimal128(_, _), Decimal128(_, _)) => decimal_op::<Decimal128Type>(op, l, l_scalar, r, r_scalar),
254 (Decimal256(_, _), Decimal256(_, _)) => decimal_op::<Decimal256Type>(op, l, l_scalar, r, r_scalar),
255 (l_t, r_t) => match (l_t, r_t) {
256 (Duration(_) | Interval(_), Date32 | Date64 | Timestamp(_, _)) if op.commutative() => {
257 arithmetic_op(op, rhs, lhs)
258 }
259 _ => Err(ArrowError::InvalidArgumentError(
260 format!("Invalid arithmetic operation: {l_t} {op} {r_t}")
261 ))
262 }
263 }
264}
265
266macro_rules! op {
268 ($l:ident, $l_s:expr, $r:ident, $r_s:expr, $op:expr) => {
269 match ($l_s, $r_s) {
270 (true, true) | (false, false) => binary($l, $r, |$l, $r| $op)?,
271 (true, false) => match ($l.null_count() == 0).then(|| $l.value(0)) {
272 None => PrimitiveArray::new_null($r.len()),
273 Some($l) => $r.unary(|$r| $op),
274 },
275 (false, true) => match ($r.null_count() == 0).then(|| $r.value(0)) {
276 None => PrimitiveArray::new_null($l.len()),
277 Some($r) => $l.unary(|$l| $op),
278 },
279 }
280 };
281}
282
283macro_rules! op_ref {
285 ($t:ty, $l:ident, $l_s:expr, $r:ident, $r_s:expr, $op:expr) => {{
286 let array: PrimitiveArray<$t> = op!($l, $l_s, $r, $r_s, $op);
287 Arc::new(array)
288 }};
289}
290
291macro_rules! try_op {
293 ($l:ident, $l_s:expr, $r:ident, $r_s:expr, $op:expr) => {
294 match ($l_s, $r_s) {
295 (true, true) | (false, false) => try_binary($l, $r, |$l, $r| $op)?,
296 (true, false) => match ($l.null_count() == 0).then(|| $l.value(0)) {
297 None => PrimitiveArray::new_null($r.len()),
298 Some($l) => $r.try_unary(|$r| $op)?,
299 },
300 (false, true) => match ($r.null_count() == 0).then(|| $r.value(0)) {
301 None => PrimitiveArray::new_null($l.len()),
302 Some($r) => $l.try_unary(|$l| $op)?,
303 },
304 }
305 };
306}
307
308macro_rules! try_op_ref {
310 ($t:ty, $l:ident, $l_s:expr, $r:ident, $r_s:expr, $op:expr) => {{
311 let array: PrimitiveArray<$t> = try_op!($l, $l_s, $r, $r_s, $op);
312 Arc::new(array)
313 }};
314}
315
316fn integer_op<T: ArrowPrimitiveType>(
318 op: Op,
319 l: &dyn Array,
320 l_s: bool,
321 r: &dyn Array,
322 r_s: bool,
323) -> Result<ArrayRef, ArrowError> {
324 let l = l.as_primitive::<T>();
325 let r = r.as_primitive::<T>();
326 let array: PrimitiveArray<T> = match op {
327 Op::AddWrapping => op!(l, l_s, r, r_s, l.add_wrapping(r)),
328 Op::Add => try_op!(l, l_s, r, r_s, l.add_checked(r)),
329 Op::SubWrapping => op!(l, l_s, r, r_s, l.sub_wrapping(r)),
330 Op::Sub => try_op!(l, l_s, r, r_s, l.sub_checked(r)),
331 Op::MulWrapping => op!(l, l_s, r, r_s, l.mul_wrapping(r)),
332 Op::Mul => try_op!(l, l_s, r, r_s, l.mul_checked(r)),
333 Op::Div => try_op!(l, l_s, r, r_s, l.div_checked(r)),
334 Op::Rem => try_op!(l, l_s, r, r_s, {
335 if r.is_zero() {
336 Err(ArrowError::DivideByZero)
337 } else {
338 Ok(l.mod_wrapping(r))
339 }
340 }),
341 };
342 Ok(Arc::new(array))
343}
344
345fn float_op<T: ArrowPrimitiveType>(
347 op: Op,
348 l: &dyn Array,
349 l_s: bool,
350 r: &dyn Array,
351 r_s: bool,
352) -> Result<ArrayRef, ArrowError> {
353 let l = l.as_primitive::<T>();
354 let r = r.as_primitive::<T>();
355 let array: PrimitiveArray<T> = match op {
356 Op::AddWrapping | Op::Add => op!(l, l_s, r, r_s, l.add_wrapping(r)),
357 Op::SubWrapping | Op::Sub => op!(l, l_s, r, r_s, l.sub_wrapping(r)),
358 Op::MulWrapping | Op::Mul => op!(l, l_s, r, r_s, l.mul_wrapping(r)),
359 Op::Div => op!(l, l_s, r, r_s, l.div_wrapping(r)),
360 Op::Rem => op!(l, l_s, r, r_s, l.mod_wrapping(r)),
361 };
362 Ok(Arc::new(array))
363}
364
365trait TimestampOp: ArrowTimestampType {
367 type Duration: ArrowPrimitiveType<Native = i64>;
368
369 fn add_year_month(timestamp: i64, delta: i32, tz: Tz) -> Option<i64>;
370 fn add_day_time(timestamp: i64, delta: IntervalDayTime, tz: Tz) -> Option<i64>;
371 fn add_month_day_nano(timestamp: i64, delta: IntervalMonthDayNano, tz: Tz) -> Option<i64>;
372
373 fn sub_year_month(timestamp: i64, delta: i32, tz: Tz) -> Option<i64>;
374 fn sub_day_time(timestamp: i64, delta: IntervalDayTime, tz: Tz) -> Option<i64>;
375 fn sub_month_day_nano(timestamp: i64, delta: IntervalMonthDayNano, tz: Tz) -> Option<i64>;
376}
377
378macro_rules! timestamp {
379 ($t:ty, $d:ty) => {
380 impl TimestampOp for $t {
381 type Duration = $d;
382
383 fn add_year_month(left: i64, right: i32, tz: Tz) -> Option<i64> {
384 Self::add_year_months(left, right, tz)
385 }
386
387 fn add_day_time(left: i64, right: IntervalDayTime, tz: Tz) -> Option<i64> {
388 Self::add_day_time(left, right, tz)
389 }
390
391 fn add_month_day_nano(left: i64, right: IntervalMonthDayNano, tz: Tz) -> Option<i64> {
392 Self::add_month_day_nano(left, right, tz)
393 }
394
395 fn sub_year_month(left: i64, right: i32, tz: Tz) -> Option<i64> {
396 Self::subtract_year_months(left, right, tz)
397 }
398
399 fn sub_day_time(left: i64, right: IntervalDayTime, tz: Tz) -> Option<i64> {
400 Self::subtract_day_time(left, right, tz)
401 }
402
403 fn sub_month_day_nano(left: i64, right: IntervalMonthDayNano, tz: Tz) -> Option<i64> {
404 Self::subtract_month_day_nano(left, right, tz)
405 }
406 }
407 };
408}
409timestamp!(TimestampSecondType, DurationSecondType);
410timestamp!(TimestampMillisecondType, DurationMillisecondType);
411timestamp!(TimestampMicrosecondType, DurationMicrosecondType);
412timestamp!(TimestampNanosecondType, DurationNanosecondType);
413
414fn timestamp_op<T: TimestampOp>(
416 op: Op,
417 l: &dyn Array,
418 l_s: bool,
419 r: &dyn Array,
420 r_s: bool,
421) -> Result<ArrayRef, ArrowError> {
422 use DataType::*;
423 use IntervalUnit::*;
424
425 let l = l.as_primitive::<T>();
426 let l_tz: Tz = l.timezone().unwrap_or("+00:00").parse()?;
427
428 let array: PrimitiveArray<T> = match (op, r.data_type()) {
429 (Op::Sub | Op::SubWrapping, Timestamp(unit, _)) if unit == &T::UNIT => {
430 let r = r.as_primitive::<T>();
431 return Ok(try_op_ref!(T::Duration, l, l_s, r, r_s, l.sub_checked(r)));
432 }
433
434 (Op::Add | Op::AddWrapping, Duration(unit)) if unit == &T::UNIT => {
435 let r = r.as_primitive::<T::Duration>();
436 try_op!(l, l_s, r, r_s, l.add_checked(r))
437 }
438 (Op::Sub | Op::SubWrapping, Duration(unit)) if unit == &T::UNIT => {
439 let r = r.as_primitive::<T::Duration>();
440 try_op!(l, l_s, r, r_s, l.sub_checked(r))
441 }
442
443 (Op::Add | Op::AddWrapping, Interval(YearMonth)) => {
444 let r = r.as_primitive::<IntervalYearMonthType>();
445 try_op!(
446 l,
447 l_s,
448 r,
449 r_s,
450 T::add_year_month(l, r, l_tz).ok_or(ArrowError::ComputeError(
451 "Timestamp out of range".to_string()
452 ))
453 )
454 }
455 (Op::Sub | Op::SubWrapping, Interval(YearMonth)) => {
456 let r = r.as_primitive::<IntervalYearMonthType>();
457 try_op!(
458 l,
459 l_s,
460 r,
461 r_s,
462 T::sub_year_month(l, r, l_tz).ok_or(ArrowError::ComputeError(
463 "Timestamp out of range".to_string()
464 ))
465 )
466 }
467
468 (Op::Add | Op::AddWrapping, Interval(DayTime)) => {
469 let r = r.as_primitive::<IntervalDayTimeType>();
470 try_op!(
471 l,
472 l_s,
473 r,
474 r_s,
475 T::add_day_time(l, r, l_tz).ok_or(ArrowError::ComputeError(
476 "Timestamp out of range".to_string()
477 ))
478 )
479 }
480 (Op::Sub | Op::SubWrapping, Interval(DayTime)) => {
481 let r = r.as_primitive::<IntervalDayTimeType>();
482 try_op!(
483 l,
484 l_s,
485 r,
486 r_s,
487 T::sub_day_time(l, r, l_tz).ok_or(ArrowError::ComputeError(
488 "Timestamp out of range".to_string()
489 ))
490 )
491 }
492
493 (Op::Add | Op::AddWrapping, Interval(MonthDayNano)) => {
494 let r = r.as_primitive::<IntervalMonthDayNanoType>();
495 try_op!(
496 l,
497 l_s,
498 r,
499 r_s,
500 T::add_month_day_nano(l, r, l_tz).ok_or(ArrowError::ComputeError(
501 "Timestamp out of range".to_string()
502 ))
503 )
504 }
505 (Op::Sub | Op::SubWrapping, Interval(MonthDayNano)) => {
506 let r = r.as_primitive::<IntervalMonthDayNanoType>();
507 try_op!(
508 l,
509 l_s,
510 r,
511 r_s,
512 T::sub_month_day_nano(l, r, l_tz).ok_or(ArrowError::ComputeError(
513 "Timestamp out of range".to_string()
514 ))
515 )
516 }
517 _ => {
518 return Err(ArrowError::InvalidArgumentError(format!(
519 "Invalid timestamp arithmetic operation: {} {op} {}",
520 l.data_type(),
521 r.data_type()
522 )));
523 }
524 };
525 Ok(Arc::new(array.with_timezone_opt(l.timezone())))
526}
527
528trait DateOp: ArrowTemporalType {
530 fn add_year_month(timestamp: Self::Native, delta: i32) -> Result<Self::Native, ArrowError>;
531 fn add_day_time(
532 timestamp: Self::Native,
533 delta: IntervalDayTime,
534 ) -> Result<Self::Native, ArrowError>;
535 fn add_month_day_nano(
536 timestamp: Self::Native,
537 delta: IntervalMonthDayNano,
538 ) -> Result<Self::Native, ArrowError>;
539
540 fn sub_year_month(timestamp: Self::Native, delta: i32) -> Result<Self::Native, ArrowError>;
541 fn sub_day_time(
542 timestamp: Self::Native,
543 delta: IntervalDayTime,
544 ) -> Result<Self::Native, ArrowError>;
545 fn sub_month_day_nano(
546 timestamp: Self::Native,
547 delta: IntervalMonthDayNano,
548 ) -> Result<Self::Native, ArrowError>;
549}
550
551impl DateOp for Date32Type {
552 fn add_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> {
553 Ok(Self::add_year_months(left, right))
555 }
556
557 fn add_day_time(
558 left: Self::Native,
559 right: IntervalDayTime,
560 ) -> Result<Self::Native, ArrowError> {
561 Ok(Self::add_day_time(left, right))
562 }
563
564 fn add_month_day_nano(
565 left: Self::Native,
566 right: IntervalMonthDayNano,
567 ) -> Result<Self::Native, ArrowError> {
568 Ok(Self::add_month_day_nano(left, right))
569 }
570
571 fn sub_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> {
572 Ok(Self::subtract_year_months(left, right))
573 }
574
575 fn sub_day_time(
576 left: Self::Native,
577 right: IntervalDayTime,
578 ) -> Result<Self::Native, ArrowError> {
579 Ok(Self::subtract_day_time(left, right))
580 }
581
582 fn sub_month_day_nano(
583 left: Self::Native,
584 right: IntervalMonthDayNano,
585 ) -> Result<Self::Native, ArrowError> {
586 Ok(Self::subtract_month_day_nano(left, right))
587 }
588}
589
590impl DateOp for Date64Type {
591 fn add_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> {
592 Self::add_year_months_opt(left, right).ok_or_else(|| {
593 ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right} months",))
594 })
595 }
596
597 fn add_day_time(
598 left: Self::Native,
599 right: IntervalDayTime,
600 ) -> Result<Self::Native, ArrowError> {
601 Self::add_day_time_opt(left, right).ok_or_else(|| {
602 ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right:?}"))
603 })
604 }
605
606 fn add_month_day_nano(
607 left: Self::Native,
608 right: IntervalMonthDayNano,
609 ) -> Result<Self::Native, ArrowError> {
610 Self::add_month_day_nano_opt(left, right).ok_or_else(|| {
611 ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right:?}"))
612 })
613 }
614
615 fn sub_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> {
616 Self::subtract_year_months_opt(left, right).ok_or_else(|| {
617 ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right} months",))
618 })
619 }
620
621 fn sub_day_time(
622 left: Self::Native,
623 right: IntervalDayTime,
624 ) -> Result<Self::Native, ArrowError> {
625 Self::subtract_day_time_opt(left, right).ok_or_else(|| {
626 ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right:?}"))
627 })
628 }
629
630 fn sub_month_day_nano(
631 left: Self::Native,
632 right: IntervalMonthDayNano,
633 ) -> Result<Self::Native, ArrowError> {
634 Self::subtract_month_day_nano_opt(left, right).ok_or_else(|| {
635 ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right:?}"))
636 })
637 }
638}
639
640trait IntervalOp: ArrowPrimitiveType {
642 fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError>;
643 fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError>;
644}
645
646impl IntervalOp for IntervalYearMonthType {
647 fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
648 left.add_checked(right)
649 }
650
651 fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
652 left.sub_checked(right)
653 }
654}
655
656impl IntervalOp for IntervalDayTimeType {
657 fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
658 let (l_days, l_ms) = Self::to_parts(left);
659 let (r_days, r_ms) = Self::to_parts(right);
660 let days = l_days.add_checked(r_days)?;
661 let ms = l_ms.add_checked(r_ms)?;
662 Ok(Self::make_value(days, ms))
663 }
664
665 fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
666 let (l_days, l_ms) = Self::to_parts(left);
667 let (r_days, r_ms) = Self::to_parts(right);
668 let days = l_days.sub_checked(r_days)?;
669 let ms = l_ms.sub_checked(r_ms)?;
670 Ok(Self::make_value(days, ms))
671 }
672}
673
674impl IntervalOp for IntervalMonthDayNanoType {
675 fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
676 let (l_months, l_days, l_nanos) = Self::to_parts(left);
677 let (r_months, r_days, r_nanos) = Self::to_parts(right);
678 let months = l_months.add_checked(r_months)?;
679 let days = l_days.add_checked(r_days)?;
680 let nanos = l_nanos.add_checked(r_nanos)?;
681 Ok(Self::make_value(months, days, nanos))
682 }
683
684 fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
685 let (l_months, l_days, l_nanos) = Self::to_parts(left);
686 let (r_months, r_days, r_nanos) = Self::to_parts(right);
687 let months = l_months.sub_checked(r_months)?;
688 let days = l_days.sub_checked(r_days)?;
689 let nanos = l_nanos.sub_checked(r_nanos)?;
690 Ok(Self::make_value(months, days, nanos))
691 }
692}
693
694fn interval_op<T: IntervalOp>(
696 op: Op,
697 l: &dyn Array,
698 l_s: bool,
699 r: &dyn Array,
700 r_s: bool,
701) -> Result<ArrayRef, ArrowError> {
702 let l = l.as_primitive::<T>();
703 let r = r.as_primitive::<T>();
704 match op {
705 Op::Add | Op::AddWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, T::add(l, r))),
706 Op::Sub | Op::SubWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub(l, r))),
707 _ => Err(ArrowError::InvalidArgumentError(format!(
708 "Invalid interval arithmetic operation: {} {op} {}",
709 l.data_type(),
710 r.data_type()
711 ))),
712 }
713}
714
715fn duration_op<T: ArrowPrimitiveType>(
716 op: Op,
717 l: &dyn Array,
718 l_s: bool,
719 r: &dyn Array,
720 r_s: bool,
721) -> Result<ArrayRef, ArrowError> {
722 let l = l.as_primitive::<T>();
723 let r = r.as_primitive::<T>();
724 match op {
725 Op::Add | Op::AddWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, l.add_checked(r))),
726 Op::Sub | Op::SubWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, l.sub_checked(r))),
727 _ => Err(ArrowError::InvalidArgumentError(format!(
728 "Invalid duration arithmetic operation: {} {op} {}",
729 l.data_type(),
730 r.data_type()
731 ))),
732 }
733}
734
735fn date_op<T: DateOp>(
737 op: Op,
738 l: &dyn Array,
739 l_s: bool,
740 r: &dyn Array,
741 r_s: bool,
742) -> Result<ArrayRef, ArrowError> {
743 use DataType::*;
744 use IntervalUnit::*;
745
746 const NUM_SECONDS_IN_DAY: i64 = 60 * 60 * 24;
747
748 let r_t = r.data_type();
749 match (T::DATA_TYPE, op, r_t) {
750 (Date32, Op::Sub | Op::SubWrapping, Date32) => {
751 let l = l.as_primitive::<Date32Type>();
752 let r = r.as_primitive::<Date32Type>();
753 return Ok(op_ref!(
754 DurationSecondType,
755 l,
756 l_s,
757 r,
758 r_s,
759 ((l as i64) - (r as i64)) * NUM_SECONDS_IN_DAY
760 ));
761 }
762 (Date64, Op::Sub | Op::SubWrapping, Date64) => {
763 let l = l.as_primitive::<Date64Type>();
764 let r = r.as_primitive::<Date64Type>();
765 let result = try_op_ref!(DurationMillisecondType, l, l_s, r, r_s, l.sub_checked(r));
766 return Ok(result);
767 }
768 _ => {}
769 }
770
771 let l = l.as_primitive::<T>();
772 match (op, r_t) {
773 (Op::Add | Op::AddWrapping, Interval(YearMonth)) => {
774 let r = r.as_primitive::<IntervalYearMonthType>();
775 Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_year_month(l, r)))
776 }
777 (Op::Sub | Op::SubWrapping, Interval(YearMonth)) => {
778 let r = r.as_primitive::<IntervalYearMonthType>();
779 Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_year_month(l, r)))
780 }
781
782 (Op::Add | Op::AddWrapping, Interval(DayTime)) => {
783 let r = r.as_primitive::<IntervalDayTimeType>();
784 Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_day_time(l, r)))
785 }
786 (Op::Sub | Op::SubWrapping, Interval(DayTime)) => {
787 let r = r.as_primitive::<IntervalDayTimeType>();
788 Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_day_time(l, r)))
789 }
790
791 (Op::Add | Op::AddWrapping, Interval(MonthDayNano)) => {
792 let r = r.as_primitive::<IntervalMonthDayNanoType>();
793 Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_month_day_nano(l, r)))
794 }
795 (Op::Sub | Op::SubWrapping, Interval(MonthDayNano)) => {
796 let r = r.as_primitive::<IntervalMonthDayNanoType>();
797 Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_month_day_nano(l, r)))
798 }
799
800 _ => Err(ArrowError::InvalidArgumentError(format!(
801 "Invalid date arithmetic operation: {} {op} {}",
802 l.data_type(),
803 r.data_type()
804 ))),
805 }
806}
807
808fn decimal_op<T: DecimalType>(
810 op: Op,
811 l: &dyn Array,
812 l_s: bool,
813 r: &dyn Array,
814 r_s: bool,
815) -> Result<ArrayRef, ArrowError> {
816 let l = l.as_primitive::<T>();
817 let r = r.as_primitive::<T>();
818
819 let (p1, s1, p2, s2) = match (l.data_type(), r.data_type()) {
820 (DataType::Decimal32(p1, s1), DataType::Decimal32(p2, s2)) => (p1, s1, p2, s2),
821 (DataType::Decimal64(p1, s1), DataType::Decimal64(p2, s2)) => (p1, s1, p2, s2),
822 (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => (p1, s1, p2, s2),
823 (DataType::Decimal256(p1, s1), DataType::Decimal256(p2, s2)) => (p1, s1, p2, s2),
824 _ => unreachable!(),
825 };
826
827 let array: PrimitiveArray<T> = match op {
830 Op::Add | Op::AddWrapping | Op::Sub | Op::SubWrapping => {
831 let result_scale = *s1.max(s2);
833
834 let result_precision =
836 (result_scale.saturating_add((*p1 as i8 - s1).max(*p2 as i8 - s2)) as u8)
837 .saturating_add(1)
838 .min(T::MAX_PRECISION);
839
840 let l_mul = T::Native::usize_as(10).pow_checked((result_scale - s1) as _)?;
841 let r_mul = T::Native::usize_as(10).pow_checked((result_scale - s2) as _)?;
842
843 match op {
844 Op::Add | Op::AddWrapping => {
845 try_op!(
846 l,
847 l_s,
848 r,
849 r_s,
850 l.mul_checked(l_mul)?.add_checked(r.mul_checked(r_mul)?)
851 )
852 }
853 Op::Sub | Op::SubWrapping => {
854 try_op!(
855 l,
856 l_s,
857 r,
858 r_s,
859 l.mul_checked(l_mul)?.sub_checked(r.mul_checked(r_mul)?)
860 )
861 }
862 _ => unreachable!(),
863 }
864 .with_precision_and_scale(result_precision, result_scale)?
865 }
866 Op::Mul | Op::MulWrapping => {
867 let result_precision = p1.saturating_add(p2 + 1).min(T::MAX_PRECISION);
868 let result_scale = s1.saturating_add(*s2);
869 if result_scale > T::MAX_SCALE {
870 return Err(ArrowError::InvalidArgumentError(format!(
873 "Output scale of {} {op} {} would exceed max scale of {}",
874 l.data_type(),
875 r.data_type(),
876 T::MAX_SCALE
877 )));
878 }
879
880 try_op!(l, l_s, r, r_s, l.mul_checked(r))
881 .with_precision_and_scale(result_precision, result_scale)?
882 }
883
884 Op::Div => {
885 let result_scale = s1.saturating_add(4).min(T::MAX_SCALE);
888 let mul_pow = result_scale - s1 + s2;
889
890 let result_precision = (mul_pow.saturating_add(*p1 as i8) as u8).min(T::MAX_PRECISION);
892
893 let (l_mul, r_mul) = match mul_pow.cmp(&0) {
894 Ordering::Greater => (
895 T::Native::usize_as(10).pow_checked(mul_pow as _)?,
896 T::Native::ONE,
897 ),
898 Ordering::Equal => (T::Native::ONE, T::Native::ONE),
899 Ordering::Less => (
900 T::Native::ONE,
901 T::Native::usize_as(10).pow_checked(mul_pow.neg_wrapping() as _)?,
902 ),
903 };
904
905 try_op!(
906 l,
907 l_s,
908 r,
909 r_s,
910 l.mul_checked(l_mul)?.div_checked(r.mul_checked(r_mul)?)
911 )
912 .with_precision_and_scale(result_precision, result_scale)?
913 }
914
915 Op::Rem => {
916 let result_scale = *s1.max(s2);
918 let result_precision =
920 (result_scale.saturating_add((*p1 as i8 - s1).min(*p2 as i8 - s2)) as u8)
921 .min(T::MAX_PRECISION);
922
923 let l_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s1) as _);
924 let r_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s2) as _);
925
926 try_op!(
927 l,
928 l_s,
929 r,
930 r_s,
931 l.mul_checked(l_mul)?.mod_checked(r.mul_checked(r_mul)?)
932 )
933 .with_precision_and_scale(result_precision, result_scale)?
934 }
935 };
936
937 Ok(Arc::new(array))
938}
939
940#[cfg(test)]
941mod tests {
942 use super::*;
943 use arrow_array::temporal_conversions::{as_date, as_datetime};
944 use arrow_buffer::{ScalarBuffer, i256};
945 use chrono::{DateTime, NaiveDate};
946
947 fn test_neg_primitive<T: ArrowPrimitiveType>(
948 input: &[T::Native],
949 out: Result<&[T::Native], &str>,
950 ) {
951 let a = PrimitiveArray::<T>::new(ScalarBuffer::from(input.to_vec()), None);
952 match out {
953 Ok(expected) => {
954 let result = neg(&a).unwrap();
955 assert_eq!(result.as_primitive::<T>().values(), expected);
956 }
957 Err(e) => {
958 let err = neg(&a).unwrap_err().to_string();
959 assert_eq!(e, err);
960 }
961 }
962 }
963
964 #[test]
965 fn test_neg() {
966 let input = &[1, -5, 2, 693, 3929];
967 let output = &[-1, 5, -2, -693, -3929];
968 test_neg_primitive::<Int32Type>(input, Ok(output));
969
970 let input = &[1, -5, 2, 693, 3929];
971 let output = &[-1, 5, -2, -693, -3929];
972 test_neg_primitive::<Int64Type>(input, Ok(output));
973 test_neg_primitive::<DurationSecondType>(input, Ok(output));
974 test_neg_primitive::<DurationMillisecondType>(input, Ok(output));
975 test_neg_primitive::<DurationMicrosecondType>(input, Ok(output));
976 test_neg_primitive::<DurationNanosecondType>(input, Ok(output));
977
978 let input = &[f32::MAX, f32::MIN, f32::INFINITY, 1.3, 0.5];
979 let output = &[f32::MIN, f32::MAX, f32::NEG_INFINITY, -1.3, -0.5];
980 test_neg_primitive::<Float32Type>(input, Ok(output));
981
982 test_neg_primitive::<Int32Type>(
983 &[i32::MIN],
984 Err("Arithmetic overflow: Overflow happened on: - -2147483648"),
985 );
986 test_neg_primitive::<Int64Type>(
987 &[i64::MIN],
988 Err("Arithmetic overflow: Overflow happened on: - -9223372036854775808"),
989 );
990 test_neg_primitive::<DurationSecondType>(
991 &[i64::MIN],
992 Err("Arithmetic overflow: Overflow happened on: - -9223372036854775808"),
993 );
994
995 let r = neg_wrapping(&Int32Array::from(vec![i32::MIN])).unwrap();
996 assert_eq!(r.as_primitive::<Int32Type>().value(0), i32::MIN);
997
998 let r = neg_wrapping(&Int64Array::from(vec![i64::MIN])).unwrap();
999 assert_eq!(r.as_primitive::<Int64Type>().value(0), i64::MIN);
1000
1001 let err = neg_wrapping(&DurationSecondArray::from(vec![i64::MIN]))
1002 .unwrap_err()
1003 .to_string();
1004
1005 assert_eq!(
1006 err,
1007 "Arithmetic overflow: Overflow happened on: - -9223372036854775808"
1008 );
1009
1010 let a = Decimal32Array::from(vec![1, 3, -44, 2, 4])
1011 .with_precision_and_scale(9, 6)
1012 .unwrap();
1013
1014 let r = neg(&a).unwrap();
1015 assert_eq!(r.data_type(), a.data_type());
1016 assert_eq!(
1017 r.as_primitive::<Decimal32Type>().values(),
1018 &[-1, -3, 44, -2, -4]
1019 );
1020
1021 let a = Decimal64Array::from(vec![1, 3, -44, 2, 4])
1022 .with_precision_and_scale(9, 6)
1023 .unwrap();
1024
1025 let r = neg(&a).unwrap();
1026 assert_eq!(r.data_type(), a.data_type());
1027 assert_eq!(
1028 r.as_primitive::<Decimal64Type>().values(),
1029 &[-1, -3, 44, -2, -4]
1030 );
1031
1032 let a = Decimal128Array::from(vec![1, 3, -44, 2, 4])
1033 .with_precision_and_scale(9, 6)
1034 .unwrap();
1035
1036 let r = neg(&a).unwrap();
1037 assert_eq!(r.data_type(), a.data_type());
1038 assert_eq!(
1039 r.as_primitive::<Decimal128Type>().values(),
1040 &[-1, -3, 44, -2, -4]
1041 );
1042
1043 let a = Decimal256Array::from(vec![
1044 i256::from_i128(342),
1045 i256::from_i128(-4949),
1046 i256::from_i128(3),
1047 ])
1048 .with_precision_and_scale(9, 6)
1049 .unwrap();
1050
1051 let r = neg(&a).unwrap();
1052 assert_eq!(r.data_type(), a.data_type());
1053 assert_eq!(
1054 r.as_primitive::<Decimal256Type>().values(),
1055 &[
1056 i256::from_i128(-342),
1057 i256::from_i128(4949),
1058 i256::from_i128(-3),
1059 ]
1060 );
1061
1062 let a = IntervalYearMonthArray::from(vec![
1063 IntervalYearMonthType::make_value(2, 4),
1064 IntervalYearMonthType::make_value(2, -4),
1065 IntervalYearMonthType::make_value(-3, -5),
1066 ]);
1067 let r = neg(&a).unwrap();
1068 assert_eq!(
1069 r.as_primitive::<IntervalYearMonthType>().values(),
1070 &[
1071 IntervalYearMonthType::make_value(-2, -4),
1072 IntervalYearMonthType::make_value(-2, 4),
1073 IntervalYearMonthType::make_value(3, 5),
1074 ]
1075 );
1076
1077 let a = IntervalDayTimeArray::from(vec![
1078 IntervalDayTimeType::make_value(2, 4),
1079 IntervalDayTimeType::make_value(2, -4),
1080 IntervalDayTimeType::make_value(-3, -5),
1081 ]);
1082 let r = neg(&a).unwrap();
1083 assert_eq!(
1084 r.as_primitive::<IntervalDayTimeType>().values(),
1085 &[
1086 IntervalDayTimeType::make_value(-2, -4),
1087 IntervalDayTimeType::make_value(-2, 4),
1088 IntervalDayTimeType::make_value(3, 5),
1089 ]
1090 );
1091
1092 let a = IntervalMonthDayNanoArray::from(vec![
1093 IntervalMonthDayNanoType::make_value(2, 4, 5953394),
1094 IntervalMonthDayNanoType::make_value(2, -4, -45839),
1095 IntervalMonthDayNanoType::make_value(-3, -5, 6944),
1096 ]);
1097 let r = neg(&a).unwrap();
1098 assert_eq!(
1099 r.as_primitive::<IntervalMonthDayNanoType>().values(),
1100 &[
1101 IntervalMonthDayNanoType::make_value(-2, -4, -5953394),
1102 IntervalMonthDayNanoType::make_value(-2, 4, 45839),
1103 IntervalMonthDayNanoType::make_value(3, 5, -6944),
1104 ]
1105 );
1106 }
1107
1108 #[test]
1109 fn test_integer() {
1110 let a = Int32Array::from(vec![4, 3, 5, -6, 100]);
1111 let b = Int32Array::from(vec![6, 2, 5, -7, 3]);
1112 let result = add(&a, &b).unwrap();
1113 assert_eq!(
1114 result.as_ref(),
1115 &Int32Array::from(vec![10, 5, 10, -13, 103])
1116 );
1117 let result = sub(&a, &b).unwrap();
1118 assert_eq!(result.as_ref(), &Int32Array::from(vec![-2, 1, 0, 1, 97]));
1119 let result = div(&a, &b).unwrap();
1120 assert_eq!(result.as_ref(), &Int32Array::from(vec![0, 1, 1, 0, 33]));
1121 let result = mul(&a, &b).unwrap();
1122 assert_eq!(result.as_ref(), &Int32Array::from(vec![24, 6, 25, 42, 300]));
1123 let result = rem(&a, &b).unwrap();
1124 assert_eq!(result.as_ref(), &Int32Array::from(vec![4, 1, 0, -6, 1]));
1125
1126 let a = Int8Array::from(vec![Some(2), None, Some(45)]);
1127 let b = Int8Array::from(vec![Some(5), Some(3), None]);
1128 let result = add(&a, &b).unwrap();
1129 assert_eq!(result.as_ref(), &Int8Array::from(vec![Some(7), None, None]));
1130
1131 let a = UInt8Array::from(vec![56, 5, 3]);
1132 let b = UInt8Array::from(vec![200, 2, 5]);
1133 let err = add(&a, &b).unwrap_err().to_string();
1134 assert_eq!(err, "Arithmetic overflow: Overflow happened on: 56 + 200");
1135 let result = add_wrapping(&a, &b).unwrap();
1136 assert_eq!(result.as_ref(), &UInt8Array::from(vec![0, 7, 8]));
1137
1138 let a = UInt8Array::from(vec![34, 5, 3]);
1139 let b = UInt8Array::from(vec![200, 2, 5]);
1140 let err = sub(&a, &b).unwrap_err().to_string();
1141 assert_eq!(err, "Arithmetic overflow: Overflow happened on: 34 - 200");
1142 let result = sub_wrapping(&a, &b).unwrap();
1143 assert_eq!(result.as_ref(), &UInt8Array::from(vec![90, 3, 254]));
1144
1145 let a = UInt8Array::from(vec![34, 5, 3]);
1146 let b = UInt8Array::from(vec![200, 2, 5]);
1147 let err = mul(&a, &b).unwrap_err().to_string();
1148 assert_eq!(err, "Arithmetic overflow: Overflow happened on: 34 * 200");
1149 let result = mul_wrapping(&a, &b).unwrap();
1150 assert_eq!(result.as_ref(), &UInt8Array::from(vec![144, 10, 15]));
1151
1152 let a = Int16Array::from(vec![i16::MIN]);
1153 let b = Int16Array::from(vec![-1]);
1154 let err = div(&a, &b).unwrap_err().to_string();
1155 assert_eq!(
1156 err,
1157 "Arithmetic overflow: Overflow happened on: -32768 / -1"
1158 );
1159
1160 let a = Int16Array::from(vec![i16::MIN]);
1161 let b = Int16Array::from(vec![-1]);
1162 let result = rem(&a, &b).unwrap();
1163 assert_eq!(result.as_ref(), &Int16Array::from(vec![0]));
1164
1165 let a = Int16Array::from(vec![21]);
1166 let b = Int16Array::from(vec![0]);
1167 let err = div(&a, &b).unwrap_err().to_string();
1168 assert_eq!(err, "Divide by zero error");
1169
1170 let a = Int16Array::from(vec![21]);
1171 let b = Int16Array::from(vec![0]);
1172 let err = rem(&a, &b).unwrap_err().to_string();
1173 assert_eq!(err, "Divide by zero error");
1174 }
1175
1176 #[test]
1177 fn test_float() {
1178 let a = Float32Array::from(vec![1., f32::MAX, 6., -4., -1., 0.]);
1179 let b = Float32Array::from(vec![1., f32::MAX, f32::MAX, -3., 45., 0.]);
1180 let result = add(&a, &b).unwrap();
1181 assert_eq!(
1182 result.as_ref(),
1183 &Float32Array::from(vec![2., f32::INFINITY, f32::MAX, -7., 44.0, 0.])
1184 );
1185
1186 let result = sub(&a, &b).unwrap();
1187 assert_eq!(
1188 result.as_ref(),
1189 &Float32Array::from(vec![0., 0., f32::MIN, -1., -46., 0.])
1190 );
1191
1192 let result = mul(&a, &b).unwrap();
1193 assert_eq!(
1194 result.as_ref(),
1195 &Float32Array::from(vec![1., f32::INFINITY, f32::INFINITY, 12., -45., 0.])
1196 );
1197
1198 let result = div(&a, &b).unwrap();
1199 let r = result.as_primitive::<Float32Type>();
1200 assert_eq!(r.value(0), 1.);
1201 assert_eq!(r.value(1), 1.);
1202 assert!(r.value(2) < f32::EPSILON);
1203 assert_eq!(r.value(3), -4. / -3.);
1204 assert!(r.value(5).is_nan());
1205
1206 let result = rem(&a, &b).unwrap();
1207 let r = result.as_primitive::<Float32Type>();
1208 assert_eq!(&r.values()[..5], &[0., 0., 6., -1., -1.]);
1209 assert!(r.value(5).is_nan());
1210 }
1211
1212 #[test]
1213 fn test_decimal() {
1214 let a = Decimal128Array::from(vec![15, 0, -577, 334, -78, 3])
1216 .with_precision_and_scale(12, 3)
1217 .unwrap();
1218
1219 let b = Decimal128Array::from(vec![54, 34, -356, 3, 6, 745])
1221 .with_precision_and_scale(12, 1)
1222 .unwrap();
1223
1224 let result = add(&a, &b).unwrap();
1225 assert_eq!(result.data_type(), &DataType::Decimal128(15, 3));
1226 assert_eq!(
1227 result.as_primitive::<Decimal128Type>().values(),
1228 &[5415, 3400, -36177, 634, 522, 74503]
1229 );
1230
1231 let result = sub(&a, &b).unwrap();
1232 assert_eq!(result.data_type(), &DataType::Decimal128(15, 3));
1233 assert_eq!(
1234 result.as_primitive::<Decimal128Type>().values(),
1235 &[-5385, -3400, 35023, 34, -678, -74497]
1236 );
1237
1238 let result = mul(&a, &b).unwrap();
1239 assert_eq!(result.data_type(), &DataType::Decimal128(25, 4));
1240 assert_eq!(
1241 result.as_primitive::<Decimal128Type>().values(),
1242 &[810, 0, 205412, 1002, -468, 2235]
1243 );
1244
1245 let result = div(&a, &b).unwrap();
1246 assert_eq!(result.data_type(), &DataType::Decimal128(17, 7));
1247 assert_eq!(
1248 result.as_primitive::<Decimal128Type>().values(),
1249 &[27777, 0, 162078, 11133333, -1300000, 402]
1250 );
1251
1252 let result = rem(&a, &b).unwrap();
1253 assert_eq!(result.data_type(), &DataType::Decimal128(12, 3));
1254 assert_eq!(
1255 result.as_primitive::<Decimal128Type>().values(),
1256 &[15, 0, -577, 34, -78, 3]
1257 );
1258
1259 let a = Decimal128Array::from(vec![1])
1260 .with_precision_and_scale(3, 3)
1261 .unwrap();
1262 let b = Decimal128Array::from(vec![1])
1263 .with_precision_and_scale(37, 37)
1264 .unwrap();
1265 let err = mul(&a, &b).unwrap_err().to_string();
1266 assert_eq!(
1267 err,
1268 "Invalid argument error: Output scale of Decimal128(3, 3) * Decimal128(37, 37) would exceed max scale of 38"
1269 );
1270
1271 let a = Decimal128Array::from(vec![1])
1272 .with_precision_and_scale(3, -2)
1273 .unwrap();
1274 let err = add(&a, &b).unwrap_err().to_string();
1275 assert_eq!(err, "Arithmetic overflow: Overflow happened on: 10 ^ 39");
1276
1277 let a = Decimal128Array::from(vec![10])
1278 .with_precision_and_scale(3, -1)
1279 .unwrap();
1280 let err = add(&a, &b).unwrap_err().to_string();
1281 assert_eq!(
1282 err,
1283 "Arithmetic overflow: Overflow happened on: 10 * 100000000000000000000000000000000000000"
1284 );
1285
1286 let b = Decimal128Array::from(vec![0])
1287 .with_precision_and_scale(1, 1)
1288 .unwrap();
1289 let err = div(&a, &b).unwrap_err().to_string();
1290 assert_eq!(err, "Divide by zero error");
1291 let err = rem(&a, &b).unwrap_err().to_string();
1292 assert_eq!(err, "Divide by zero error");
1293 }
1294
1295 fn test_timestamp_impl<T: TimestampOp>() {
1296 let a = PrimitiveArray::<T>::new(vec![2000000, 434030324, 53943340].into(), None);
1297 let b = PrimitiveArray::<T>::new(vec![329593, 59349, 694994].into(), None);
1298
1299 let result = sub(&a, &b).unwrap();
1300 assert_eq!(
1301 result.as_primitive::<T::Duration>().values(),
1302 &[1670407, 433970975, 53248346]
1303 );
1304
1305 let r2 = add(&b, &result.as_ref()).unwrap();
1306 assert_eq!(r2.as_ref(), &a);
1307
1308 let r3 = add(&result.as_ref(), &b).unwrap();
1309 assert_eq!(r3.as_ref(), &a);
1310
1311 let format_array = |x: &dyn Array| -> Vec<String> {
1312 x.as_primitive::<T>()
1313 .values()
1314 .into_iter()
1315 .map(|x| as_datetime::<T>(*x).unwrap().to_string())
1316 .collect()
1317 };
1318
1319 let values = vec![
1320 "1970-01-01T00:00:00Z",
1321 "2010-04-01T04:00:20Z",
1322 "1960-01-30T04:23:20Z",
1323 ]
1324 .into_iter()
1325 .map(|x| T::make_value(DateTime::parse_from_rfc3339(x).unwrap().naive_utc()).unwrap())
1326 .collect();
1327
1328 let a = PrimitiveArray::<T>::new(values, None);
1329 let b = IntervalYearMonthArray::from(vec![
1330 IntervalYearMonthType::make_value(5, 34),
1331 IntervalYearMonthType::make_value(-2, 4),
1332 IntervalYearMonthType::make_value(7, -4),
1333 ]);
1334 let r4 = add(&a, &b).unwrap();
1335 assert_eq!(
1336 &format_array(r4.as_ref()),
1337 &[
1338 "1977-11-01 00:00:00".to_string(),
1339 "2008-08-01 04:00:20".to_string(),
1340 "1966-09-30 04:23:20".to_string()
1341 ]
1342 );
1343
1344 let r5 = sub(&r4, &b).unwrap();
1345 assert_eq!(r5.as_ref(), &a);
1346
1347 let b = IntervalDayTimeArray::from(vec![
1348 IntervalDayTimeType::make_value(5, 454000),
1349 IntervalDayTimeType::make_value(-34, 0),
1350 IntervalDayTimeType::make_value(7, -4000),
1351 ]);
1352 let r6 = add(&a, &b).unwrap();
1353 assert_eq!(
1354 &format_array(r6.as_ref()),
1355 &[
1356 "1970-01-06 00:07:34".to_string(),
1357 "2010-02-26 04:00:20".to_string(),
1358 "1960-02-06 04:23:16".to_string()
1359 ]
1360 );
1361
1362 let r7 = sub(&r6, &b).unwrap();
1363 assert_eq!(r7.as_ref(), &a);
1364
1365 let b = IntervalMonthDayNanoArray::from(vec![
1366 IntervalMonthDayNanoType::make_value(344, 34, -43_000_000_000),
1367 IntervalMonthDayNanoType::make_value(-593, -33, 13_000_000_000),
1368 IntervalMonthDayNanoType::make_value(5, 2, 493_000_000_000),
1369 ]);
1370 let r8 = add(&a, &b).unwrap();
1371 assert_eq!(
1372 &format_array(r8.as_ref()),
1373 &[
1374 "1998-10-04 23:59:17".to_string(),
1375 "1960-09-29 04:00:33".to_string(),
1376 "1960-07-02 04:31:33".to_string()
1377 ]
1378 );
1379
1380 let r9 = sub(&r8, &b).unwrap();
1381 assert_eq!(
1383 &format_array(r9.as_ref()),
1384 &[
1385 "1970-01-02 00:00:00".to_string(),
1386 "2010-04-02 04:00:20".to_string(),
1387 "1960-01-31 04:23:20".to_string()
1388 ]
1389 );
1390 }
1391
1392 #[test]
1393 fn test_timestamp() {
1394 test_timestamp_impl::<TimestampSecondType>();
1395 test_timestamp_impl::<TimestampMillisecondType>();
1396 test_timestamp_impl::<TimestampMicrosecondType>();
1397 test_timestamp_impl::<TimestampNanosecondType>();
1398 }
1399
1400 #[test]
1401 fn test_interval() {
1402 let a = IntervalYearMonthArray::from(vec![
1403 IntervalYearMonthType::make_value(32, 4),
1404 IntervalYearMonthType::make_value(32, 4),
1405 ]);
1406 let b = IntervalYearMonthArray::from(vec![
1407 IntervalYearMonthType::make_value(-4, 6),
1408 IntervalYearMonthType::make_value(-3, 23),
1409 ]);
1410 let result = add(&a, &b).unwrap();
1411 assert_eq!(
1412 result.as_ref(),
1413 &IntervalYearMonthArray::from(vec![
1414 IntervalYearMonthType::make_value(28, 10),
1415 IntervalYearMonthType::make_value(29, 27)
1416 ])
1417 );
1418 let result = sub(&a, &b).unwrap();
1419 assert_eq!(
1420 result.as_ref(),
1421 &IntervalYearMonthArray::from(vec![
1422 IntervalYearMonthType::make_value(36, -2),
1423 IntervalYearMonthType::make_value(35, -19)
1424 ])
1425 );
1426
1427 let a = IntervalDayTimeArray::from(vec![
1428 IntervalDayTimeType::make_value(32, 4),
1429 IntervalDayTimeType::make_value(32, 4),
1430 ]);
1431 let b = IntervalDayTimeArray::from(vec![
1432 IntervalDayTimeType::make_value(-4, 6),
1433 IntervalDayTimeType::make_value(-3, 23),
1434 ]);
1435 let result = add(&a, &b).unwrap();
1436 assert_eq!(
1437 result.as_ref(),
1438 &IntervalDayTimeArray::from(vec![
1439 IntervalDayTimeType::make_value(28, 10),
1440 IntervalDayTimeType::make_value(29, 27)
1441 ])
1442 );
1443 let result = sub(&a, &b).unwrap();
1444 assert_eq!(
1445 result.as_ref(),
1446 &IntervalDayTimeArray::from(vec![
1447 IntervalDayTimeType::make_value(36, -2),
1448 IntervalDayTimeType::make_value(35, -19)
1449 ])
1450 );
1451 let a = IntervalMonthDayNanoArray::from(vec![
1452 IntervalMonthDayNanoType::make_value(32, 4, 4000000000000),
1453 IntervalMonthDayNanoType::make_value(32, 4, 45463000000000000),
1454 ]);
1455 let b = IntervalMonthDayNanoArray::from(vec![
1456 IntervalMonthDayNanoType::make_value(-4, 6, 46000000000000),
1457 IntervalMonthDayNanoType::make_value(-3, 23, 3564000000000000),
1458 ]);
1459 let result = add(&a, &b).unwrap();
1460 assert_eq!(
1461 result.as_ref(),
1462 &IntervalMonthDayNanoArray::from(vec![
1463 IntervalMonthDayNanoType::make_value(28, 10, 50000000000000),
1464 IntervalMonthDayNanoType::make_value(29, 27, 49027000000000000)
1465 ])
1466 );
1467 let result = sub(&a, &b).unwrap();
1468 assert_eq!(
1469 result.as_ref(),
1470 &IntervalMonthDayNanoArray::from(vec![
1471 IntervalMonthDayNanoType::make_value(36, -2, -42000000000000),
1472 IntervalMonthDayNanoType::make_value(35, -19, 41899000000000000)
1473 ])
1474 );
1475 let a = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNano::MAX]);
1476 let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNano::ONE]);
1477 let err = add(&a, &b).unwrap_err().to_string();
1478 assert_eq!(
1479 err,
1480 "Arithmetic overflow: Overflow happened on: 2147483647 + 1"
1481 );
1482 }
1483
1484 fn test_duration_impl<T: ArrowPrimitiveType<Native = i64>>() {
1485 let a = PrimitiveArray::<T>::new(vec![1000, 4394, -3944].into(), None);
1486 let b = PrimitiveArray::<T>::new(vec![4, -5, -243].into(), None);
1487
1488 let result = add(&a, &b).unwrap();
1489 assert_eq!(result.as_primitive::<T>().values(), &[1004, 4389, -4187]);
1490 let result = sub(&a, &b).unwrap();
1491 assert_eq!(result.as_primitive::<T>().values(), &[996, 4399, -3701]);
1492
1493 let err = mul(&a, &b).unwrap_err().to_string();
1494 assert!(
1495 err.contains("Invalid duration arithmetic operation"),
1496 "{err}"
1497 );
1498
1499 let err = div(&a, &b).unwrap_err().to_string();
1500 assert!(
1501 err.contains("Invalid duration arithmetic operation"),
1502 "{err}"
1503 );
1504
1505 let err = rem(&a, &b).unwrap_err().to_string();
1506 assert!(
1507 err.contains("Invalid duration arithmetic operation"),
1508 "{err}"
1509 );
1510
1511 let a = PrimitiveArray::<T>::new(vec![i64::MAX].into(), None);
1512 let b = PrimitiveArray::<T>::new(vec![1].into(), None);
1513 let err = add(&a, &b).unwrap_err().to_string();
1514 assert_eq!(
1515 err,
1516 "Arithmetic overflow: Overflow happened on: 9223372036854775807 + 1"
1517 );
1518 }
1519
1520 #[test]
1521 fn test_duration() {
1522 test_duration_impl::<DurationSecondType>();
1523 test_duration_impl::<DurationMillisecondType>();
1524 test_duration_impl::<DurationMicrosecondType>();
1525 test_duration_impl::<DurationNanosecondType>();
1526 }
1527
1528 fn test_date_impl<T: ArrowPrimitiveType, F>(f: F)
1529 where
1530 F: Fn(NaiveDate) -> T::Native,
1531 T::Native: TryInto<i64>,
1532 {
1533 let a = PrimitiveArray::<T>::new(
1534 vec![
1535 f(NaiveDate::from_ymd_opt(1979, 1, 30).unwrap()),
1536 f(NaiveDate::from_ymd_opt(2010, 4, 3).unwrap()),
1537 f(NaiveDate::from_ymd_opt(2008, 2, 29).unwrap()),
1538 ]
1539 .into(),
1540 None,
1541 );
1542
1543 let b = IntervalYearMonthArray::from(vec![
1544 IntervalYearMonthType::make_value(34, 2),
1545 IntervalYearMonthType::make_value(3, -3),
1546 IntervalYearMonthType::make_value(-12, 4),
1547 ]);
1548
1549 let format_array = |x: &dyn Array| -> Vec<String> {
1550 x.as_primitive::<T>()
1551 .values()
1552 .into_iter()
1553 .map(|x| {
1554 as_date::<T>((*x).try_into().ok().unwrap())
1555 .unwrap()
1556 .to_string()
1557 })
1558 .collect()
1559 };
1560
1561 let result = add(&a, &b).unwrap();
1562 assert_eq!(
1563 &format_array(result.as_ref()),
1564 &[
1565 "2013-03-30".to_string(),
1566 "2013-01-03".to_string(),
1567 "1996-06-29".to_string(),
1568 ]
1569 );
1570 let result = sub(&result, &b).unwrap();
1571 assert_eq!(result.as_ref(), &a);
1572
1573 let b = IntervalDayTimeArray::from(vec![
1574 IntervalDayTimeType::make_value(34, 2),
1575 IntervalDayTimeType::make_value(3, -3),
1576 IntervalDayTimeType::make_value(-12, 4),
1577 ]);
1578
1579 let result = add(&a, &b).unwrap();
1580 assert_eq!(
1581 &format_array(result.as_ref()),
1582 &[
1583 "1979-03-05".to_string(),
1584 "2010-04-06".to_string(),
1585 "2008-02-17".to_string(),
1586 ]
1587 );
1588 let result = sub(&result, &b).unwrap();
1589 assert_eq!(result.as_ref(), &a);
1590
1591 let b = IntervalMonthDayNanoArray::from(vec![
1592 IntervalMonthDayNanoType::make_value(34, 2, -34353534),
1593 IntervalMonthDayNanoType::make_value(3, -3, 2443),
1594 IntervalMonthDayNanoType::make_value(-12, 4, 2323242423232),
1595 ]);
1596
1597 let result = add(&a, &b).unwrap();
1598 assert_eq!(
1599 &format_array(result.as_ref()),
1600 &[
1601 "1981-12-02".to_string(),
1602 "2010-06-30".to_string(),
1603 "2007-03-04".to_string(),
1604 ]
1605 );
1606 let result = sub(&result, &b).unwrap();
1607 assert_eq!(
1608 &format_array(result.as_ref()),
1609 &[
1610 "1979-01-31".to_string(),
1611 "2010-04-02".to_string(),
1612 "2008-02-29".to_string(),
1613 ]
1614 );
1615 }
1616
1617 #[test]
1618 fn test_date() {
1619 test_date_impl::<Date32Type, _>(Date32Type::from_naive_date);
1620 test_date_impl::<Date64Type, _>(Date64Type::from_naive_date);
1621
1622 let a = Date32Array::from(vec![i32::MIN, i32::MAX, 23, 7684]);
1623 let b = Date32Array::from(vec![i32::MIN, i32::MIN, -2, 45]);
1624 let result = sub(&a, &b).unwrap();
1625 assert_eq!(
1626 result.as_primitive::<DurationSecondType>().values(),
1627 &[0, 371085174288000, 2160000, 660009600]
1628 );
1629
1630 let a = Date64Array::from(vec![4343, 76676, 3434]);
1631 let b = Date64Array::from(vec![3, -5, 5]);
1632 let result = sub(&a, &b).unwrap();
1633 assert_eq!(
1634 result.as_primitive::<DurationMillisecondType>().values(),
1635 &[4340, 76681, 3429]
1636 );
1637
1638 let a = Date64Array::from(vec![i64::MAX]);
1639 let b = Date64Array::from(vec![-1]);
1640 let err = sub(&a, &b).unwrap_err().to_string();
1641 assert_eq!(
1642 err,
1643 "Arithmetic overflow: Overflow happened on: 9223372036854775807 - -1"
1644 );
1645 }
1646
1647 #[test]
1648 fn test_date64_to_naive_date_opt_boundary_values() {
1649 use arrow_array::types::Date64Type;
1650
1651 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
1655 let ms_per_day = 24 * 60 * 60 * 1000i64;
1656
1657 let max_valid_date = NaiveDate::from_ymd_opt(262142, 12, 31).unwrap();
1659 let min_valid_date = NaiveDate::from_ymd_opt(-262143, 1, 1).unwrap();
1660
1661 let max_valid_millis = (max_valid_date - epoch).num_milliseconds();
1663 let min_valid_millis = (min_valid_date - epoch).num_milliseconds();
1664
1665 assert_eq!(
1667 max_valid_millis, 8210266790400000i64,
1668 "December 31, 262142 should be 8210266790400000 ms from epoch"
1669 );
1670 assert_eq!(
1671 min_valid_millis, -8334601228800000i64,
1672 "January 1, -262143 should be -8334601228800000 ms from epoch"
1673 );
1674
1675 assert!(
1677 Date64Type::to_naive_date_opt(max_valid_millis).is_some(),
1678 "December 31, 262142 should return Some"
1679 );
1680 assert!(
1681 Date64Type::to_naive_date_opt(min_valid_millis).is_some(),
1682 "January 1, -262143 should return Some"
1683 );
1684
1685 assert!(
1687 Date64Type::to_naive_date_opt(max_valid_millis + ms_per_day).is_none(),
1688 "January 1, 262143 should return None"
1689 );
1690 assert!(
1691 Date64Type::to_naive_date_opt(min_valid_millis - ms_per_day).is_none(),
1692 "December 31, -262144 should return None"
1693 );
1694
1695 assert!(
1697 Date64Type::to_naive_date_opt(0).is_some(),
1698 "Epoch (1970-01-01) should return Some"
1699 );
1700 let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
1701 let year_2000_millis = (year_2000 - epoch).num_milliseconds();
1702 assert!(
1703 Date64Type::to_naive_date_opt(year_2000_millis).is_some(),
1704 "Year 2000 should return Some"
1705 );
1706
1707 assert!(
1709 Date64Type::to_naive_date_opt(i64::MAX).is_none(),
1710 "i64::MAX should return None"
1711 );
1712 assert!(
1713 Date64Type::to_naive_date_opt(i64::MIN).is_none(),
1714 "i64::MIN should return None"
1715 );
1716 }
1717
1718 #[test]
1719 fn test_date64_add_year_months_opt_boundary_values() {
1720 use arrow_array::types::Date64Type;
1721
1722 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
1723
1724 let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
1726 let year_2000_millis = (year_2000 - epoch).num_milliseconds();
1727 assert!(
1728 Date64Type::add_year_months_opt(year_2000_millis, 120).is_some(),
1729 "Adding 10 years to year 2000 should succeed"
1730 );
1731
1732 let large_year = NaiveDate::from_ymd_opt(5000, 1, 1).unwrap();
1734 let large_year_millis = (large_year - epoch).num_milliseconds();
1735 assert!(
1736 Date64Type::add_year_months_opt(large_year_millis, 12).is_some(),
1737 "Adding 12 months to year 5000 should succeed"
1738 );
1739
1740 let neg_year = NaiveDate::from_ymd_opt(-5000, 12, 31).unwrap();
1741 let neg_year_millis = (neg_year - epoch).num_milliseconds();
1742 assert!(
1743 Date64Type::add_year_months_opt(neg_year_millis, -12).is_some(),
1744 "Subtracting 12 months from year -5000 should succeed"
1745 );
1746
1747 assert!(
1749 Date64Type::add_year_months_opt(i64::MAX, 1).is_none(),
1750 "Adding months to i64::MAX should fail"
1751 );
1752 assert!(
1753 Date64Type::add_year_months_opt(i64::MIN, -1).is_none(),
1754 "Subtracting months from i64::MIN should fail"
1755 );
1756
1757 assert!(
1759 Date64Type::add_year_months_opt(year_2000_millis, 0).is_some(),
1760 "Adding zero months should always succeed for valid dates"
1761 );
1762 }
1763
1764 #[test]
1765 fn test_date64_add_day_time_opt_boundary_values() {
1766 use arrow_array::types::Date64Type;
1767 use arrow_buffer::IntervalDayTime;
1768
1769 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
1770
1771 let near_max_date = NaiveDate::from_ymd_opt(200000, 12, 1).unwrap();
1773 let near_max_millis = (near_max_date - epoch).num_milliseconds();
1774
1775 let interval_30_days = IntervalDayTime::new(30, 0);
1777 assert!(
1778 Date64Type::add_day_time_opt(near_max_millis, interval_30_days).is_some(),
1779 "Adding 30 days to large year should succeed"
1780 );
1781
1782 let interval_large_days = IntervalDayTime::new(100000000, 0);
1784 assert!(
1785 Date64Type::add_day_time_opt(near_max_millis, interval_large_days).is_none(),
1786 "Adding 100M days to large year should fail"
1787 );
1788
1789 let near_min_date = NaiveDate::from_ymd_opt(-200000, 2, 1).unwrap();
1791 let near_min_millis = (near_min_date - epoch).num_milliseconds();
1792
1793 let interval_minus_30_days = IntervalDayTime::new(-30, 0);
1795 assert!(
1796 Date64Type::add_day_time_opt(near_min_millis, interval_minus_30_days).is_some(),
1797 "Subtracting 30 days from large negative year should succeed"
1798 );
1799
1800 let interval_minus_large_days = IntervalDayTime::new(-100000000, 0);
1802 assert!(
1803 Date64Type::add_day_time_opt(near_min_millis, interval_minus_large_days).is_none(),
1804 "Subtracting 100M days from large negative year should fail"
1805 );
1806
1807 let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
1809 let year_2000_millis = (year_2000 - epoch).num_milliseconds();
1810 let interval_1000_days = IntervalDayTime::new(1000, 12345);
1811 assert!(
1812 Date64Type::add_day_time_opt(year_2000_millis, interval_1000_days).is_some(),
1813 "Adding 1000 days and time to year 2000 should succeed"
1814 );
1815
1816 let interval_one_day = IntervalDayTime::new(1, 0);
1818 assert!(
1819 Date64Type::add_day_time_opt(i64::MAX, interval_one_day).is_none(),
1820 "Adding interval to i64::MAX should fail"
1821 );
1822 assert!(
1823 Date64Type::add_day_time_opt(i64::MIN, IntervalDayTime::new(-1, 0)).is_none(),
1824 "Subtracting interval from i64::MIN should fail"
1825 );
1826
1827 let max_interval = IntervalDayTime::new(i32::MAX, i32::MAX);
1829 assert!(
1830 Date64Type::add_day_time_opt(0, max_interval).is_none(),
1831 "Adding extreme interval should fail"
1832 );
1833
1834 let min_interval = IntervalDayTime::new(i32::MIN, i32::MIN);
1835 assert!(
1836 Date64Type::add_day_time_opt(0, min_interval).is_none(),
1837 "Adding extreme negative interval should fail"
1838 );
1839
1840 let large_ms_interval = IntervalDayTime::new(0, i32::MAX);
1842 assert!(
1843 Date64Type::add_day_time_opt(year_2000_millis, large_ms_interval).is_some(),
1844 "Adding large milliseconds within valid range should succeed"
1845 );
1846 }
1847
1848 #[test]
1849 fn test_date64_add_month_day_nano_opt_boundary_values() {
1850 use arrow_array::types::Date64Type;
1851 use arrow_buffer::IntervalMonthDayNano;
1852
1853 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
1854
1855 let near_max_date = NaiveDate::from_ymd_opt(5000, 11, 1).unwrap();
1857 let near_max_millis = (near_max_date - epoch).num_milliseconds();
1858
1859 let interval_safe = IntervalMonthDayNano::new(1, 30, 0);
1861 assert!(
1862 Date64Type::add_month_day_nano_opt(near_max_millis, interval_safe).is_some(),
1863 "Adding 1 month 30 days to large year should succeed"
1864 );
1865
1866 let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
1868 let year_2000_millis = (year_2000 - epoch).num_milliseconds();
1869
1870 let zero_interval = IntervalMonthDayNano::new(0, 0, 0);
1872 assert!(
1873 Date64Type::add_month_day_nano_opt(year_2000_millis, zero_interval).is_some(),
1874 "Adding zero interval should always succeed for valid dates"
1875 );
1876
1877 let near_min_date = NaiveDate::from_ymd_opt(-5000, 2, 28).unwrap();
1879 let near_min_millis = (near_min_date - epoch).num_milliseconds();
1880
1881 let interval_safe_neg = IntervalMonthDayNano::new(-1, -30, 0);
1883 assert!(
1884 Date64Type::add_month_day_nano_opt(near_min_millis, interval_safe_neg).is_some(),
1885 "Subtracting 1 month 30 days from large negative year should succeed"
1886 );
1887
1888 assert!(
1890 Date64Type::add_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(1, 0, 0))
1891 .is_none(),
1892 "Adding interval to i64::MAX should fail"
1893 );
1894
1895 let interval_normal = IntervalMonthDayNano::new(2, 10, 123_456_789_000);
1896 assert!(
1897 Date64Type::add_month_day_nano_opt(year_2000_millis, interval_normal).is_some(),
1898 "Adding 2 months, 10 days, and nanos to year 2000 should succeed"
1899 );
1900
1901 assert!(
1903 Date64Type::add_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(1, 0, 0))
1904 .is_none(),
1905 "Adding interval to i64::MAX should fail"
1906 );
1907 assert!(
1908 Date64Type::add_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(-1, 0, 0))
1909 .is_none(),
1910 "Subtracting interval from i64::MIN should fail"
1911 );
1912
1913 let nano_interval = IntervalMonthDayNano::new(0, 0, 999_999_999);
1917 assert!(
1918 Date64Type::add_month_day_nano_opt(year_2000_millis, nano_interval).is_some(),
1919 "Adding nanoseconds within valid range should succeed"
1920 );
1921
1922 let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); assert!(
1925 Date64Type::add_month_day_nano_opt(year_2000_millis, large_nano_interval).is_some(),
1926 "Adding 1 day worth of nanoseconds should succeed"
1927 );
1928 }
1929
1930 #[test]
1931 fn test_date64_subtract_year_months_opt_boundary_values() {
1932 use arrow_array::types::Date64Type;
1933
1934 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
1935
1936 let near_min_date = NaiveDate::from_ymd_opt(-5000, 12, 31).unwrap();
1938 let near_min_millis = (near_min_date - epoch).num_milliseconds();
1939
1940 assert!(
1942 Date64Type::subtract_year_months_opt(near_min_millis, 12).is_some(),
1943 "Subtracting 12 months from year -5000 should succeed"
1944 );
1945
1946 let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
1948 let year_2000_millis = (year_2000 - epoch).num_milliseconds();
1949
1950 assert!(
1952 Date64Type::subtract_year_months_opt(year_2000_millis, 0).is_some(),
1953 "Subtracting zero months should always succeed for valid dates"
1954 );
1955
1956 let near_max_date = NaiveDate::from_ymd_opt(5000, 1, 1).unwrap();
1958 let near_max_millis = (near_max_date - epoch).num_milliseconds();
1959
1960 assert!(
1962 Date64Type::subtract_year_months_opt(near_max_millis, -12).is_some(),
1963 "Adding 12 months to year 5000 should succeed"
1964 );
1965
1966 assert!(
1968 Date64Type::subtract_year_months_opt(i64::MAX, -1).is_none(),
1969 "Adding months to i64::MAX should fail"
1970 );
1971
1972 assert!(
1973 Date64Type::subtract_year_months_opt(year_2000_millis, 12).is_some(),
1974 "Subtracting 1 year from year 2000 should succeed"
1975 );
1976
1977 assert!(
1979 Date64Type::subtract_year_months_opt(i64::MAX, -1).is_none(),
1980 "Adding months to i64::MAX should fail"
1981 );
1982 assert!(
1983 Date64Type::subtract_year_months_opt(i64::MIN, 1).is_none(),
1984 "Subtracting months from i64::MIN should fail"
1985 );
1986
1987 let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap();
1989 let valid_millis = (valid_date - epoch).num_milliseconds();
1990 assert!(
1991 Date64Type::subtract_year_months_opt(valid_millis, 0).is_some(),
1992 "Subtracting zero months should always succeed for valid dates"
1993 );
1994 }
1995
1996 #[test]
1997 fn test_date64_subtract_day_time_opt_boundary_values() {
1998 use arrow_array::types::Date64Type;
1999 use arrow_buffer::IntervalDayTime;
2000
2001 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
2002
2003 let near_min_date = NaiveDate::from_ymd_opt(-200000, 2, 1).unwrap();
2005 let near_min_millis = (near_min_date - epoch).num_milliseconds();
2006
2007 let interval_30_days = IntervalDayTime::new(30, 0);
2009 assert!(
2010 Date64Type::subtract_day_time_opt(near_min_millis, interval_30_days).is_some(),
2011 "Subtracting 30 days from large negative year should succeed"
2012 );
2013
2014 let interval_large_days = IntervalDayTime::new(100000000, 0);
2016 assert!(
2017 Date64Type::subtract_day_time_opt(near_min_millis, interval_large_days).is_none(),
2018 "Subtracting 100M days from large negative year should fail"
2019 );
2020
2021 let near_max_date = NaiveDate::from_ymd_opt(200000, 12, 1).unwrap();
2023 let near_max_millis = (near_max_date - epoch).num_milliseconds();
2024
2025 let interval_minus_30_days = IntervalDayTime::new(-30, 0);
2027 assert!(
2028 Date64Type::subtract_day_time_opt(near_max_millis, interval_minus_30_days).is_some(),
2029 "Adding 30 days to large year should succeed"
2030 );
2031
2032 let interval_minus_large_days = IntervalDayTime::new(-100000000, 0);
2034 assert!(
2035 Date64Type::subtract_day_time_opt(near_max_millis, interval_minus_large_days).is_none(),
2036 "Adding 100M days to large year should fail"
2037 );
2038
2039 let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
2041 let year_2000_millis = (year_2000 - epoch).num_milliseconds();
2042 let interval_1000_days = IntervalDayTime::new(1000, 12345);
2043 assert!(
2044 Date64Type::subtract_day_time_opt(year_2000_millis, interval_1000_days).is_some(),
2045 "Subtracting 1000 days and time from year 2000 should succeed"
2046 );
2047
2048 let interval_one_day = IntervalDayTime::new(1, 0);
2050 assert!(
2051 Date64Type::subtract_day_time_opt(i64::MIN, interval_one_day).is_none(),
2052 "Subtracting interval from i64::MIN should fail"
2053 );
2054 assert!(
2055 Date64Type::subtract_day_time_opt(i64::MAX, IntervalDayTime::new(-1, 0)).is_none(),
2056 "Adding interval to i64::MAX should fail"
2057 );
2058
2059 let max_interval = IntervalDayTime::new(i32::MAX, i32::MAX);
2061 assert!(
2062 Date64Type::subtract_day_time_opt(0, max_interval).is_none(),
2063 "Subtracting extreme interval should fail"
2064 );
2065
2066 let min_interval = IntervalDayTime::new(i32::MIN, i32::MIN);
2067 assert!(
2068 Date64Type::subtract_day_time_opt(0, min_interval).is_none(),
2069 "Subtracting extreme negative interval should fail"
2070 );
2071
2072 let large_ms_interval = IntervalDayTime::new(0, i32::MAX);
2074 assert!(
2075 Date64Type::subtract_day_time_opt(year_2000_millis, large_ms_interval).is_some(),
2076 "Subtracting large milliseconds within valid range should succeed"
2077 );
2078
2079 let zero_interval = IntervalDayTime::new(0, 0);
2081 let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap();
2082 let valid_millis = (valid_date - epoch).num_milliseconds();
2083 assert!(
2084 Date64Type::subtract_day_time_opt(valid_millis, zero_interval).is_some(),
2085 "Subtracting zero interval should always succeed for valid dates"
2086 );
2087 }
2088
2089 #[test]
2090 fn test_date64_subtract_month_day_nano_opt_boundary_values() {
2091 use arrow_array::types::Date64Type;
2092 use arrow_buffer::IntervalMonthDayNano;
2093
2094 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
2095
2096 let near_min_date = NaiveDate::from_ymd_opt(-5000, 2, 28).unwrap();
2098 let near_min_millis = (near_min_date - epoch).num_milliseconds();
2099
2100 let interval_safe = IntervalMonthDayNano::new(1, 30, 0);
2102 assert!(
2103 Date64Type::subtract_month_day_nano_opt(near_min_millis, interval_safe).is_some(),
2104 "Subtracting 1 month 30 days from large negative year should succeed"
2105 );
2106
2107 let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
2109 let year_2000_millis = (year_2000 - epoch).num_milliseconds();
2110
2111 let zero_interval = IntervalMonthDayNano::new(0, 0, 0);
2113 assert!(
2114 Date64Type::subtract_month_day_nano_opt(year_2000_millis, zero_interval).is_some(),
2115 "Subtracting zero interval should always succeed for valid dates"
2116 );
2117
2118 let near_max_date = NaiveDate::from_ymd_opt(5000, 11, 1).unwrap();
2120 let near_max_millis = (near_max_date - epoch).num_milliseconds();
2121
2122 let interval_safe_neg = IntervalMonthDayNano::new(-1, -30, 0);
2124 assert!(
2125 Date64Type::subtract_month_day_nano_opt(near_max_millis, interval_safe_neg).is_some(),
2126 "Adding 1 month 30 days to large year should succeed"
2127 );
2128
2129 assert!(
2131 Date64Type::subtract_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(1, 0, 0))
2132 .is_none(),
2133 "Subtracting interval from i64::MIN should fail"
2134 );
2135
2136 let interval_normal = IntervalMonthDayNano::new(2, 10, 123_456_789_000);
2137 assert!(
2138 Date64Type::subtract_month_day_nano_opt(year_2000_millis, interval_normal).is_some(),
2139 "Subtracting 2 months, 10 days, and nanos from year 2000 should succeed"
2140 );
2141
2142 assert!(
2144 Date64Type::subtract_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(1, 0, 0))
2145 .is_none(),
2146 "Subtracting interval from i64::MIN should fail"
2147 );
2148 assert!(
2149 Date64Type::subtract_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(-1, 0, 0))
2150 .is_none(),
2151 "Adding interval to i64::MAX should fail"
2152 );
2153
2154 let nano_interval = IntervalMonthDayNano::new(0, 0, 999_999_999);
2156 assert!(
2157 Date64Type::subtract_month_day_nano_opt(year_2000_millis, nano_interval).is_some(),
2158 "Subtracting nanoseconds within valid range should succeed"
2159 );
2160
2161 let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); assert!(
2164 Date64Type::subtract_month_day_nano_opt(year_2000_millis, large_nano_interval)
2165 .is_some(),
2166 "Subtracting 1 day worth of nanoseconds should succeed"
2167 );
2168
2169 let zero_interval = IntervalMonthDayNano::new(0, 0, 0);
2171 let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap();
2172 let valid_millis = (valid_date - epoch).num_milliseconds();
2173 assert!(
2174 Date64Type::subtract_month_day_nano_opt(valid_millis, zero_interval).is_some(),
2175 "Subtracting zero interval should always succeed for valid dates"
2176 );
2177 }
2178}