Skip to main content

arrow_arith/
numeric.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Defines numeric arithmetic kernels on [`PrimitiveArray`], such as [`add`]
19
20use 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
33/// Perform `lhs + rhs`, returning an error on overflow
34pub fn add(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
35    arithmetic_op(Op::Add, lhs, rhs)
36}
37
38/// Perform `lhs + rhs`, wrapping on overflow for [`DataType::is_integer`]
39pub fn add_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
40    arithmetic_op(Op::AddWrapping, lhs, rhs)
41}
42
43/// Perform `lhs - rhs`, returning an error on overflow
44pub fn sub(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
45    arithmetic_op(Op::Sub, lhs, rhs)
46}
47
48/// Perform `lhs - rhs`, wrapping on overflow for [`DataType::is_integer`]
49pub fn sub_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
50    arithmetic_op(Op::SubWrapping, lhs, rhs)
51}
52
53/// Perform `lhs * rhs`, returning an error on overflow
54pub fn mul(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
55    arithmetic_op(Op::Mul, lhs, rhs)
56}
57
58/// Perform `lhs * rhs`, wrapping on overflow for [`DataType::is_integer`]
59pub fn mul_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
60    arithmetic_op(Op::MulWrapping, lhs, rhs)
61}
62
63/// Perform `lhs / rhs`
64///
65/// Overflow or division by zero will result in an error, with exception to
66/// floating point numbers, which instead follow the IEEE 754 rules
67pub fn div(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
68    arithmetic_op(Op::Div, lhs, rhs)
69}
70
71/// Perform `lhs % rhs`
72///
73/// Division by zero will result in an error, with exception to
74/// floating point numbers, which instead follow the IEEE 754 rules
75///
76/// `signed_integer::MIN % -1` will not result in an error but return 0
77pub 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
97/// Negates each element of  `array`, returning an error on overflow
98///
99/// Note: negation of unsigned arrays is not supported and will return in an error,
100/// for wrapping unsigned negation consider using [`neg_wrapping`][neg_wrapping()]
101pub 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
178/// Negates each element of  `array`, wrapping on overflow for [`DataType::is_integer`]
179pub 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/// An enumeration of arithmetic operations
187///
188/// This allows sharing the type dispatch logic across the various kernels
189#[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
219/// Dispatch the given `op` to the appropriate specialized kernel
220fn 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
266/// Perform an infallible binary operation on potentially scalar inputs
267macro_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
283/// Same as `op` but with a type hint for the returned array
284macro_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
291/// Perform a fallible binary operation on potentially scalar inputs
292macro_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
308/// Same as `try_op` but with a type hint for the returned array
309macro_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
316/// Perform an arithmetic operation on integers
317fn 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
345/// Perform an arithmetic operation on floats
346fn 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
365/// Arithmetic trait for timestamp arrays
366trait 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
414/// Perform arithmetic operation on a timestamp array
415fn 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
528/// Arithmetic trait for date arrays
529trait 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
551macro_rules! date {
552    ($t:ty) => {
553        impl DateOp for $t {
554            fn add_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> {
555                Self::add_year_months_opt(left, right).ok_or_else(|| {
556                    ArrowError::ComputeError(format!(
557                        "Date arithmetic overflow: {left} + {right} months"
558                    ))
559                })
560            }
561
562            fn add_day_time(
563                left: Self::Native,
564                right: IntervalDayTime,
565            ) -> Result<Self::Native, ArrowError> {
566                Self::add_day_time_opt(left, right).ok_or_else(|| {
567                    ArrowError::ComputeError(format!(
568                        "Date arithmetic overflow: {left} + {right:?}"
569                    ))
570                })
571            }
572
573            fn add_month_day_nano(
574                left: Self::Native,
575                right: IntervalMonthDayNano,
576            ) -> Result<Self::Native, ArrowError> {
577                Self::add_month_day_nano_opt(left, right).ok_or_else(|| {
578                    ArrowError::ComputeError(format!(
579                        "Date arithmetic overflow: {left} + {right:?}"
580                    ))
581                })
582            }
583
584            fn sub_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> {
585                Self::subtract_year_months_opt(left, right).ok_or_else(|| {
586                    ArrowError::ComputeError(format!(
587                        "Date arithmetic overflow: {left} - {right} months"
588                    ))
589                })
590            }
591
592            fn sub_day_time(
593                left: Self::Native,
594                right: IntervalDayTime,
595            ) -> Result<Self::Native, ArrowError> {
596                Self::subtract_day_time_opt(left, right).ok_or_else(|| {
597                    ArrowError::ComputeError(format!(
598                        "Date arithmetic overflow: {left} - {right:?}"
599                    ))
600                })
601            }
602
603            fn sub_month_day_nano(
604                left: Self::Native,
605                right: IntervalMonthDayNano,
606            ) -> Result<Self::Native, ArrowError> {
607                Self::subtract_month_day_nano_opt(left, right).ok_or_else(|| {
608                    ArrowError::ComputeError(format!(
609                        "Date arithmetic overflow: {left} - {right:?}"
610                    ))
611                })
612            }
613        }
614    };
615}
616
617date!(Date32Type);
618date!(Date64Type);
619
620/// Arithmetic trait for interval arrays
621trait IntervalOp: ArrowPrimitiveType {
622    fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError>;
623    fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError>;
624}
625
626impl IntervalOp for IntervalYearMonthType {
627    fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
628        left.add_checked(right)
629    }
630
631    fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
632        left.sub_checked(right)
633    }
634}
635
636impl IntervalOp for IntervalDayTimeType {
637    fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
638        let (l_days, l_ms) = Self::to_parts(left);
639        let (r_days, r_ms) = Self::to_parts(right);
640        let days = l_days.add_checked(r_days)?;
641        let ms = l_ms.add_checked(r_ms)?;
642        Ok(Self::make_value(days, ms))
643    }
644
645    fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
646        let (l_days, l_ms) = Self::to_parts(left);
647        let (r_days, r_ms) = Self::to_parts(right);
648        let days = l_days.sub_checked(r_days)?;
649        let ms = l_ms.sub_checked(r_ms)?;
650        Ok(Self::make_value(days, ms))
651    }
652}
653
654impl IntervalOp for IntervalMonthDayNanoType {
655    fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
656        let (l_months, l_days, l_nanos) = Self::to_parts(left);
657        let (r_months, r_days, r_nanos) = Self::to_parts(right);
658        let months = l_months.add_checked(r_months)?;
659        let days = l_days.add_checked(r_days)?;
660        let nanos = l_nanos.add_checked(r_nanos)?;
661        Ok(Self::make_value(months, days, nanos))
662    }
663
664    fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> {
665        let (l_months, l_days, l_nanos) = Self::to_parts(left);
666        let (r_months, r_days, r_nanos) = Self::to_parts(right);
667        let months = l_months.sub_checked(r_months)?;
668        let days = l_days.sub_checked(r_days)?;
669        let nanos = l_nanos.sub_checked(r_nanos)?;
670        Ok(Self::make_value(months, days, nanos))
671    }
672}
673
674/// Perform arithmetic operation on an interval array
675fn interval_op<T: IntervalOp>(
676    op: Op,
677    l: &dyn Array,
678    l_s: bool,
679    r: &dyn Array,
680    r_s: bool,
681) -> Result<ArrayRef, ArrowError> {
682    let l = l.as_primitive::<T>();
683    let r = r.as_primitive::<T>();
684    match op {
685        Op::Add | Op::AddWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, T::add(l, r))),
686        Op::Sub | Op::SubWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub(l, r))),
687        _ => Err(ArrowError::InvalidArgumentError(format!(
688            "Invalid interval arithmetic operation: {} {op} {}",
689            l.data_type(),
690            r.data_type()
691        ))),
692    }
693}
694
695fn duration_op<T: ArrowPrimitiveType>(
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, l.add_checked(r))),
706        Op::Sub | Op::SubWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, l.sub_checked(r))),
707        _ => Err(ArrowError::InvalidArgumentError(format!(
708            "Invalid duration arithmetic operation: {} {op} {}",
709            l.data_type(),
710            r.data_type()
711        ))),
712    }
713}
714
715/// Perform arithmetic operation on a date array
716fn date_op<T: DateOp>(
717    op: Op,
718    l: &dyn Array,
719    l_s: bool,
720    r: &dyn Array,
721    r_s: bool,
722) -> Result<ArrayRef, ArrowError> {
723    use DataType::*;
724    use IntervalUnit::*;
725
726    const NUM_SECONDS_IN_DAY: i64 = 60 * 60 * 24;
727
728    let r_t = r.data_type();
729    match (T::DATA_TYPE, op, r_t) {
730        (Date32, Op::Sub | Op::SubWrapping, Date32) => {
731            let l = l.as_primitive::<Date32Type>();
732            let r = r.as_primitive::<Date32Type>();
733            return Ok(op_ref!(
734                DurationSecondType,
735                l,
736                l_s,
737                r,
738                r_s,
739                ((l as i64) - (r as i64)) * NUM_SECONDS_IN_DAY
740            ));
741        }
742        (Date64, Op::Sub | Op::SubWrapping, Date64) => {
743            let l = l.as_primitive::<Date64Type>();
744            let r = r.as_primitive::<Date64Type>();
745            let result = try_op_ref!(DurationMillisecondType, l, l_s, r, r_s, l.sub_checked(r));
746            return Ok(result);
747        }
748        _ => {}
749    }
750
751    let l = l.as_primitive::<T>();
752    match (op, r_t) {
753        (Op::Add | Op::AddWrapping, Interval(YearMonth)) => {
754            let r = r.as_primitive::<IntervalYearMonthType>();
755            Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_year_month(l, r)))
756        }
757        (Op::Sub | Op::SubWrapping, Interval(YearMonth)) => {
758            let r = r.as_primitive::<IntervalYearMonthType>();
759            Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_year_month(l, r)))
760        }
761
762        (Op::Add | Op::AddWrapping, Interval(DayTime)) => {
763            let r = r.as_primitive::<IntervalDayTimeType>();
764            Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_day_time(l, r)))
765        }
766        (Op::Sub | Op::SubWrapping, Interval(DayTime)) => {
767            let r = r.as_primitive::<IntervalDayTimeType>();
768            Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_day_time(l, r)))
769        }
770
771        (Op::Add | Op::AddWrapping, Interval(MonthDayNano)) => {
772            let r = r.as_primitive::<IntervalMonthDayNanoType>();
773            Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_month_day_nano(l, r)))
774        }
775        (Op::Sub | Op::SubWrapping, Interval(MonthDayNano)) => {
776            let r = r.as_primitive::<IntervalMonthDayNanoType>();
777            Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_month_day_nano(l, r)))
778        }
779
780        _ => Err(ArrowError::InvalidArgumentError(format!(
781            "Invalid date arithmetic operation: {} {op} {}",
782            l.data_type(),
783            r.data_type()
784        ))),
785    }
786}
787
788/// Perform arithmetic operation on decimal arrays
789fn decimal_op<T: DecimalType>(
790    op: Op,
791    l: &dyn Array,
792    l_s: bool,
793    r: &dyn Array,
794    r_s: bool,
795) -> Result<ArrayRef, ArrowError> {
796    let l = l.as_primitive::<T>();
797    let r = r.as_primitive::<T>();
798
799    let (p1, s1, p2, s2) = match (l.data_type(), r.data_type()) {
800        (DataType::Decimal32(p1, s1), DataType::Decimal32(p2, s2)) => (p1, s1, p2, s2),
801        (DataType::Decimal64(p1, s1), DataType::Decimal64(p2, s2)) => (p1, s1, p2, s2),
802        (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => (p1, s1, p2, s2),
803        (DataType::Decimal256(p1, s1), DataType::Decimal256(p2, s2)) => (p1, s1, p2, s2),
804        _ => unreachable!(),
805    };
806
807    // Follow the Hive decimal arithmetic rules
808    // https://cwiki.apache.org/confluence/download/attachments/27362075/Hive_Decimal_Precision_Scale_Support.pdf
809    let array: PrimitiveArray<T> = match op {
810        Op::Add | Op::AddWrapping | Op::Sub | Op::SubWrapping => {
811            // max(s1, s2)
812            let result_scale = *s1.max(s2);
813
814            // max(s1, s2) + max(p1-s1, p2-s2) + 1
815            let result_precision =
816                (result_scale.saturating_add((*p1 as i8 - s1).max(*p2 as i8 - s2)) as u8)
817                    .saturating_add(1)
818                    .min(T::MAX_PRECISION);
819
820            let l_mul = T::Native::usize_as(10).pow_checked((result_scale - s1) as _)?;
821            let r_mul = T::Native::usize_as(10).pow_checked((result_scale - s2) as _)?;
822
823            match op {
824                Op::Add | Op::AddWrapping => {
825                    try_op!(
826                        l,
827                        l_s,
828                        r,
829                        r_s,
830                        l.mul_checked(l_mul)?.add_checked(r.mul_checked(r_mul)?)
831                    )
832                }
833                Op::Sub | Op::SubWrapping => {
834                    try_op!(
835                        l,
836                        l_s,
837                        r,
838                        r_s,
839                        l.mul_checked(l_mul)?.sub_checked(r.mul_checked(r_mul)?)
840                    )
841                }
842                _ => unreachable!(),
843            }
844            .with_precision_and_scale(result_precision, result_scale)?
845        }
846        Op::Mul | Op::MulWrapping => {
847            let result_precision = p1.saturating_add(p2 + 1).min(T::MAX_PRECISION);
848            let result_scale = s1.saturating_add(*s2);
849            if result_scale > T::MAX_SCALE {
850                // SQL standard says that if the resulting scale of a multiply operation goes
851                // beyond the maximum, rounding is not acceptable and thus an error occurs
852                return Err(ArrowError::InvalidArgumentError(format!(
853                    "Output scale of {} {op} {} would exceed max scale of {}",
854                    l.data_type(),
855                    r.data_type(),
856                    T::MAX_SCALE
857                )));
858            }
859
860            try_op!(l, l_s, r, r_s, l.mul_checked(r))
861                .with_precision_and_scale(result_precision, result_scale)?
862        }
863
864        Op::Div => {
865            // Follow postgres and MySQL adding a fixed scale increment of 4
866            // s1 + 4
867            let result_scale = s1.saturating_add(4).min(T::MAX_SCALE);
868            let mul_pow = result_scale - s1 + s2;
869
870            // p1 - s1 + s2 + result_scale
871            let result_precision = (mul_pow.saturating_add(*p1 as i8) as u8).min(T::MAX_PRECISION);
872
873            let (l_mul, r_mul) = match mul_pow.cmp(&0) {
874                Ordering::Greater => (
875                    T::Native::usize_as(10).pow_checked(mul_pow as _)?,
876                    T::Native::ONE,
877                ),
878                Ordering::Equal => (T::Native::ONE, T::Native::ONE),
879                Ordering::Less => (
880                    T::Native::ONE,
881                    T::Native::usize_as(10).pow_checked(mul_pow.neg_wrapping() as _)?,
882                ),
883            };
884
885            try_op!(
886                l,
887                l_s,
888                r,
889                r_s,
890                l.mul_checked(l_mul)?.div_checked(r.mul_checked(r_mul)?)
891            )
892            .with_precision_and_scale(result_precision, result_scale)?
893        }
894
895        Op::Rem => {
896            // max(s1, s2)
897            let result_scale = *s1.max(s2);
898            // min(p1-s1, p2 -s2) + max( s1,s2 )
899            let result_precision =
900                (result_scale.saturating_add((*p1 as i8 - s1).min(*p2 as i8 - s2)) as u8)
901                    .min(T::MAX_PRECISION);
902
903            let l_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s1) as _);
904            let r_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s2) as _);
905
906            try_op!(
907                l,
908                l_s,
909                r,
910                r_s,
911                l.mul_checked(l_mul)?.mod_checked(r.mul_checked(r_mul)?)
912            )
913            .with_precision_and_scale(result_precision, result_scale)?
914        }
915    };
916
917    Ok(Arc::new(array))
918}
919
920#[cfg(test)]
921mod tests {
922    use super::*;
923    use arrow_array::temporal_conversions::{as_date, as_datetime};
924    use arrow_buffer::{ScalarBuffer, i256};
925    use chrono::{DateTime, NaiveDate};
926
927    // The valid date range of NaiveDate is from January 1, -262143 to December 31, 262142 (Gregorian calendar).
928    const MAX_VALID_DATE: NaiveDate = NaiveDate::from_ymd_opt(262142, 12, 31).unwrap();
929    const MIN_VALID_DATE: NaiveDate = NaiveDate::from_ymd_opt(-262143, 1, 1).unwrap();
930    const MAX_VALID_MILLIS: i64 = date_to_millis(MAX_VALID_DATE);
931    const MIN_VALID_MILLIS: i64 = date_to_millis(MIN_VALID_DATE);
932    const MAX_VALID_DAYS: i32 = date_to_days(MAX_VALID_DATE);
933    const MIN_VALID_DAYS: i32 = date_to_days(MIN_VALID_DATE);
934    const EPOCH: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
935    const YEAR_2000: NaiveDate = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
936
937    const fn date_to_millis(date: NaiveDate) -> i64 {
938        date.signed_duration_since(EPOCH).num_milliseconds()
939    }
940
941    const fn date_to_days(date: NaiveDate) -> i32 {
942        date.signed_duration_since(EPOCH).num_days() as i32
943    }
944
945    fn test_neg_primitive<T: ArrowPrimitiveType>(
946        input: &[T::Native],
947        out: Result<&[T::Native], &str>,
948    ) {
949        let a = PrimitiveArray::<T>::new(ScalarBuffer::from(input.to_vec()), None);
950        match out {
951            Ok(expected) => {
952                let result = neg(&a).unwrap();
953                assert_eq!(result.as_primitive::<T>().values(), expected);
954            }
955            Err(e) => {
956                let err = neg(&a).unwrap_err().to_string();
957                assert_eq!(e, err);
958            }
959        }
960    }
961
962    #[test]
963    fn test_neg() {
964        let input = &[1, -5, 2, 693, 3929];
965        let output = &[-1, 5, -2, -693, -3929];
966        test_neg_primitive::<Int32Type>(input, Ok(output));
967
968        let input = &[1, -5, 2, 693, 3929];
969        let output = &[-1, 5, -2, -693, -3929];
970        test_neg_primitive::<Int64Type>(input, Ok(output));
971        test_neg_primitive::<DurationSecondType>(input, Ok(output));
972        test_neg_primitive::<DurationMillisecondType>(input, Ok(output));
973        test_neg_primitive::<DurationMicrosecondType>(input, Ok(output));
974        test_neg_primitive::<DurationNanosecondType>(input, Ok(output));
975
976        let input = &[f32::MAX, f32::MIN, f32::INFINITY, 1.3, 0.5];
977        let output = &[f32::MIN, f32::MAX, f32::NEG_INFINITY, -1.3, -0.5];
978        test_neg_primitive::<Float32Type>(input, Ok(output));
979
980        test_neg_primitive::<Int32Type>(
981            &[i32::MIN],
982            Err("Arithmetic overflow: Overflow happened on: - -2147483648"),
983        );
984        test_neg_primitive::<Int64Type>(
985            &[i64::MIN],
986            Err("Arithmetic overflow: Overflow happened on: - -9223372036854775808"),
987        );
988        test_neg_primitive::<DurationSecondType>(
989            &[i64::MIN],
990            Err("Arithmetic overflow: Overflow happened on: - -9223372036854775808"),
991        );
992
993        let r = neg_wrapping(&Int32Array::from(vec![i32::MIN])).unwrap();
994        assert_eq!(r.as_primitive::<Int32Type>().value(0), i32::MIN);
995
996        let r = neg_wrapping(&Int64Array::from(vec![i64::MIN])).unwrap();
997        assert_eq!(r.as_primitive::<Int64Type>().value(0), i64::MIN);
998
999        let err = neg_wrapping(&DurationSecondArray::from(vec![i64::MIN]))
1000            .unwrap_err()
1001            .to_string();
1002
1003        assert_eq!(
1004            err,
1005            "Arithmetic overflow: Overflow happened on: - -9223372036854775808"
1006        );
1007
1008        let a = Decimal32Array::from(vec![1, 3, -44, 2, 4])
1009            .with_precision_and_scale(9, 6)
1010            .unwrap();
1011
1012        let r = neg(&a).unwrap();
1013        assert_eq!(r.data_type(), a.data_type());
1014        assert_eq!(
1015            r.as_primitive::<Decimal32Type>().values(),
1016            &[-1, -3, 44, -2, -4]
1017        );
1018
1019        let a = Decimal64Array::from(vec![1, 3, -44, 2, 4])
1020            .with_precision_and_scale(9, 6)
1021            .unwrap();
1022
1023        let r = neg(&a).unwrap();
1024        assert_eq!(r.data_type(), a.data_type());
1025        assert_eq!(
1026            r.as_primitive::<Decimal64Type>().values(),
1027            &[-1, -3, 44, -2, -4]
1028        );
1029
1030        let a = Decimal128Array::from(vec![1, 3, -44, 2, 4])
1031            .with_precision_and_scale(9, 6)
1032            .unwrap();
1033
1034        let r = neg(&a).unwrap();
1035        assert_eq!(r.data_type(), a.data_type());
1036        assert_eq!(
1037            r.as_primitive::<Decimal128Type>().values(),
1038            &[-1, -3, 44, -2, -4]
1039        );
1040
1041        let a = Decimal256Array::from(vec![
1042            i256::from_i128(342),
1043            i256::from_i128(-4949),
1044            i256::from_i128(3),
1045        ])
1046        .with_precision_and_scale(9, 6)
1047        .unwrap();
1048
1049        let r = neg(&a).unwrap();
1050        assert_eq!(r.data_type(), a.data_type());
1051        assert_eq!(
1052            r.as_primitive::<Decimal256Type>().values(),
1053            &[
1054                i256::from_i128(-342),
1055                i256::from_i128(4949),
1056                i256::from_i128(-3),
1057            ]
1058        );
1059
1060        let a = IntervalYearMonthArray::from(vec![
1061            IntervalYearMonthType::make_value(2, 4),
1062            IntervalYearMonthType::make_value(2, -4),
1063            IntervalYearMonthType::make_value(-3, -5),
1064        ]);
1065        let r = neg(&a).unwrap();
1066        assert_eq!(
1067            r.as_primitive::<IntervalYearMonthType>().values(),
1068            &[
1069                IntervalYearMonthType::make_value(-2, -4),
1070                IntervalYearMonthType::make_value(-2, 4),
1071                IntervalYearMonthType::make_value(3, 5),
1072            ]
1073        );
1074
1075        let a = IntervalDayTimeArray::from(vec![
1076            IntervalDayTimeType::make_value(2, 4),
1077            IntervalDayTimeType::make_value(2, -4),
1078            IntervalDayTimeType::make_value(-3, -5),
1079        ]);
1080        let r = neg(&a).unwrap();
1081        assert_eq!(
1082            r.as_primitive::<IntervalDayTimeType>().values(),
1083            &[
1084                IntervalDayTimeType::make_value(-2, -4),
1085                IntervalDayTimeType::make_value(-2, 4),
1086                IntervalDayTimeType::make_value(3, 5),
1087            ]
1088        );
1089
1090        let a = IntervalMonthDayNanoArray::from(vec![
1091            IntervalMonthDayNanoType::make_value(2, 4, 5953394),
1092            IntervalMonthDayNanoType::make_value(2, -4, -45839),
1093            IntervalMonthDayNanoType::make_value(-3, -5, 6944),
1094        ]);
1095        let r = neg(&a).unwrap();
1096        assert_eq!(
1097            r.as_primitive::<IntervalMonthDayNanoType>().values(),
1098            &[
1099                IntervalMonthDayNanoType::make_value(-2, -4, -5953394),
1100                IntervalMonthDayNanoType::make_value(-2, 4, 45839),
1101                IntervalMonthDayNanoType::make_value(3, 5, -6944),
1102            ]
1103        );
1104    }
1105
1106    #[test]
1107    fn test_integer() {
1108        let a = Int32Array::from(vec![4, 3, 5, -6, 100]);
1109        let b = Int32Array::from(vec![6, 2, 5, -7, 3]);
1110        let result = add(&a, &b).unwrap();
1111        assert_eq!(
1112            result.as_ref(),
1113            &Int32Array::from(vec![10, 5, 10, -13, 103])
1114        );
1115        let result = sub(&a, &b).unwrap();
1116        assert_eq!(result.as_ref(), &Int32Array::from(vec![-2, 1, 0, 1, 97]));
1117        let result = div(&a, &b).unwrap();
1118        assert_eq!(result.as_ref(), &Int32Array::from(vec![0, 1, 1, 0, 33]));
1119        let result = mul(&a, &b).unwrap();
1120        assert_eq!(result.as_ref(), &Int32Array::from(vec![24, 6, 25, 42, 300]));
1121        let result = rem(&a, &b).unwrap();
1122        assert_eq!(result.as_ref(), &Int32Array::from(vec![4, 1, 0, -6, 1]));
1123
1124        let a = Int8Array::from(vec![Some(2), None, Some(45)]);
1125        let b = Int8Array::from(vec![Some(5), Some(3), None]);
1126        let result = add(&a, &b).unwrap();
1127        assert_eq!(result.as_ref(), &Int8Array::from(vec![Some(7), None, None]));
1128
1129        let a = UInt8Array::from(vec![56, 5, 3]);
1130        let b = UInt8Array::from(vec![200, 2, 5]);
1131        let err = add(&a, &b).unwrap_err().to_string();
1132        assert_eq!(err, "Arithmetic overflow: Overflow happened on: 56 + 200");
1133        let result = add_wrapping(&a, &b).unwrap();
1134        assert_eq!(result.as_ref(), &UInt8Array::from(vec![0, 7, 8]));
1135
1136        let a = UInt8Array::from(vec![34, 5, 3]);
1137        let b = UInt8Array::from(vec![200, 2, 5]);
1138        let err = sub(&a, &b).unwrap_err().to_string();
1139        assert_eq!(err, "Arithmetic overflow: Overflow happened on: 34 - 200");
1140        let result = sub_wrapping(&a, &b).unwrap();
1141        assert_eq!(result.as_ref(), &UInt8Array::from(vec![90, 3, 254]));
1142
1143        let a = UInt8Array::from(vec![34, 5, 3]);
1144        let b = UInt8Array::from(vec![200, 2, 5]);
1145        let err = mul(&a, &b).unwrap_err().to_string();
1146        assert_eq!(err, "Arithmetic overflow: Overflow happened on: 34 * 200");
1147        let result = mul_wrapping(&a, &b).unwrap();
1148        assert_eq!(result.as_ref(), &UInt8Array::from(vec![144, 10, 15]));
1149
1150        let a = Int16Array::from(vec![i16::MIN]);
1151        let b = Int16Array::from(vec![-1]);
1152        let err = div(&a, &b).unwrap_err().to_string();
1153        assert_eq!(
1154            err,
1155            "Arithmetic overflow: Overflow happened on: -32768 / -1"
1156        );
1157
1158        let a = Int16Array::from(vec![i16::MIN]);
1159        let b = Int16Array::from(vec![-1]);
1160        let result = rem(&a, &b).unwrap();
1161        assert_eq!(result.as_ref(), &Int16Array::from(vec![0]));
1162
1163        let a = Int16Array::from(vec![21]);
1164        let b = Int16Array::from(vec![0]);
1165        let err = div(&a, &b).unwrap_err().to_string();
1166        assert_eq!(err, "Divide by zero error");
1167
1168        let a = Int16Array::from(vec![21]);
1169        let b = Int16Array::from(vec![0]);
1170        let err = rem(&a, &b).unwrap_err().to_string();
1171        assert_eq!(err, "Divide by zero error");
1172    }
1173
1174    #[test]
1175    fn test_float() {
1176        let a = Float32Array::from(vec![1., f32::MAX, 6., -4., -1., 0.]);
1177        let b = Float32Array::from(vec![1., f32::MAX, f32::MAX, -3., 45., 0.]);
1178        let result = add(&a, &b).unwrap();
1179        assert_eq!(
1180            result.as_ref(),
1181            &Float32Array::from(vec![2., f32::INFINITY, f32::MAX, -7., 44.0, 0.])
1182        );
1183
1184        let result = sub(&a, &b).unwrap();
1185        assert_eq!(
1186            result.as_ref(),
1187            &Float32Array::from(vec![0., 0., f32::MIN, -1., -46., 0.])
1188        );
1189
1190        let result = mul(&a, &b).unwrap();
1191        assert_eq!(
1192            result.as_ref(),
1193            &Float32Array::from(vec![1., f32::INFINITY, f32::INFINITY, 12., -45., 0.])
1194        );
1195
1196        let result = div(&a, &b).unwrap();
1197        let r = result.as_primitive::<Float32Type>();
1198        assert_eq!(r.value(0), 1.);
1199        assert_eq!(r.value(1), 1.);
1200        assert!(r.value(2) < f32::EPSILON);
1201        assert_eq!(r.value(3), -4. / -3.);
1202        assert!(r.value(5).is_nan());
1203
1204        let result = rem(&a, &b).unwrap();
1205        let r = result.as_primitive::<Float32Type>();
1206        assert_eq!(&r.values()[..5], &[0., 0., 6., -1., -1.]);
1207        assert!(r.value(5).is_nan());
1208    }
1209
1210    #[test]
1211    fn test_decimal() {
1212        // 0.015 7.842 -0.577 0.334 -0.078 0.003
1213        let a = Decimal128Array::from(vec![15, 0, -577, 334, -78, 3])
1214            .with_precision_and_scale(12, 3)
1215            .unwrap();
1216
1217        // 5.4 0 -35.6 0.3 0.6 7.45
1218        let b = Decimal128Array::from(vec![54, 34, -356, 3, 6, 745])
1219            .with_precision_and_scale(12, 1)
1220            .unwrap();
1221
1222        let result = add(&a, &b).unwrap();
1223        assert_eq!(result.data_type(), &DataType::Decimal128(15, 3));
1224        assert_eq!(
1225            result.as_primitive::<Decimal128Type>().values(),
1226            &[5415, 3400, -36177, 634, 522, 74503]
1227        );
1228
1229        let result = sub(&a, &b).unwrap();
1230        assert_eq!(result.data_type(), &DataType::Decimal128(15, 3));
1231        assert_eq!(
1232            result.as_primitive::<Decimal128Type>().values(),
1233            &[-5385, -3400, 35023, 34, -678, -74497]
1234        );
1235
1236        let result = mul(&a, &b).unwrap();
1237        assert_eq!(result.data_type(), &DataType::Decimal128(25, 4));
1238        assert_eq!(
1239            result.as_primitive::<Decimal128Type>().values(),
1240            &[810, 0, 205412, 1002, -468, 2235]
1241        );
1242
1243        let result = div(&a, &b).unwrap();
1244        assert_eq!(result.data_type(), &DataType::Decimal128(17, 7));
1245        assert_eq!(
1246            result.as_primitive::<Decimal128Type>().values(),
1247            &[27777, 0, 162078, 11133333, -1300000, 402]
1248        );
1249
1250        let result = rem(&a, &b).unwrap();
1251        assert_eq!(result.data_type(), &DataType::Decimal128(12, 3));
1252        assert_eq!(
1253            result.as_primitive::<Decimal128Type>().values(),
1254            &[15, 0, -577, 34, -78, 3]
1255        );
1256
1257        let a = Decimal128Array::from(vec![1])
1258            .with_precision_and_scale(3, 3)
1259            .unwrap();
1260        let b = Decimal128Array::from(vec![1])
1261            .with_precision_and_scale(37, 37)
1262            .unwrap();
1263        let err = mul(&a, &b).unwrap_err().to_string();
1264        assert_eq!(
1265            err,
1266            "Invalid argument error: Output scale of Decimal128(3, 3) * Decimal128(37, 37) would exceed max scale of 38"
1267        );
1268
1269        let a = Decimal128Array::from(vec![1])
1270            .with_precision_and_scale(3, -2)
1271            .unwrap();
1272        let err = add(&a, &b).unwrap_err().to_string();
1273        assert_eq!(err, "Arithmetic overflow: Overflow happened on: 10 ^ 39");
1274
1275        let a = Decimal128Array::from(vec![10])
1276            .with_precision_and_scale(3, -1)
1277            .unwrap();
1278        let err = add(&a, &b).unwrap_err().to_string();
1279        assert_eq!(
1280            err,
1281            "Arithmetic overflow: Overflow happened on: 10 * 100000000000000000000000000000000000000"
1282        );
1283
1284        let b = Decimal128Array::from(vec![0])
1285            .with_precision_and_scale(1, 1)
1286            .unwrap();
1287        let err = div(&a, &b).unwrap_err().to_string();
1288        assert_eq!(err, "Divide by zero error");
1289        let err = rem(&a, &b).unwrap_err().to_string();
1290        assert_eq!(err, "Divide by zero error");
1291    }
1292
1293    fn test_timestamp_impl<T: TimestampOp>() {
1294        let a = PrimitiveArray::<T>::new(vec![2000000, 434030324, 53943340].into(), None);
1295        let b = PrimitiveArray::<T>::new(vec![329593, 59349, 694994].into(), None);
1296
1297        let result = sub(&a, &b).unwrap();
1298        assert_eq!(
1299            result.as_primitive::<T::Duration>().values(),
1300            &[1670407, 433970975, 53248346]
1301        );
1302
1303        let r2 = add(&b, &result.as_ref()).unwrap();
1304        assert_eq!(r2.as_ref(), &a);
1305
1306        let r3 = add(&result.as_ref(), &b).unwrap();
1307        assert_eq!(r3.as_ref(), &a);
1308
1309        let format_array = |x: &dyn Array| -> Vec<String> {
1310            x.as_primitive::<T>()
1311                .values()
1312                .into_iter()
1313                .map(|x| as_datetime::<T>(*x).unwrap().to_string())
1314                .collect()
1315        };
1316
1317        let values = vec![
1318            "1970-01-01T00:00:00Z",
1319            "2010-04-01T04:00:20Z",
1320            "1960-01-30T04:23:20Z",
1321        ]
1322        .into_iter()
1323        .map(|x| T::make_value(DateTime::parse_from_rfc3339(x).unwrap().naive_utc()).unwrap())
1324        .collect();
1325
1326        let a = PrimitiveArray::<T>::new(values, None);
1327        let b = IntervalYearMonthArray::from(vec![
1328            IntervalYearMonthType::make_value(5, 34),
1329            IntervalYearMonthType::make_value(-2, 4),
1330            IntervalYearMonthType::make_value(7, -4),
1331        ]);
1332        let r4 = add(&a, &b).unwrap();
1333        assert_eq!(
1334            &format_array(r4.as_ref()),
1335            &[
1336                "1977-11-01 00:00:00".to_string(),
1337                "2008-08-01 04:00:20".to_string(),
1338                "1966-09-30 04:23:20".to_string()
1339            ]
1340        );
1341
1342        let r5 = sub(&r4, &b).unwrap();
1343        assert_eq!(r5.as_ref(), &a);
1344
1345        let b = IntervalDayTimeArray::from(vec![
1346            IntervalDayTimeType::make_value(5, 454000),
1347            IntervalDayTimeType::make_value(-34, 0),
1348            IntervalDayTimeType::make_value(7, -4000),
1349        ]);
1350        let r6 = add(&a, &b).unwrap();
1351        assert_eq!(
1352            &format_array(r6.as_ref()),
1353            &[
1354                "1970-01-06 00:07:34".to_string(),
1355                "2010-02-26 04:00:20".to_string(),
1356                "1960-02-06 04:23:16".to_string()
1357            ]
1358        );
1359
1360        let r7 = sub(&r6, &b).unwrap();
1361        assert_eq!(r7.as_ref(), &a);
1362
1363        let b = IntervalMonthDayNanoArray::from(vec![
1364            IntervalMonthDayNanoType::make_value(344, 34, -43_000_000_000),
1365            IntervalMonthDayNanoType::make_value(-593, -33, 13_000_000_000),
1366            IntervalMonthDayNanoType::make_value(5, 2, 493_000_000_000),
1367        ]);
1368        let r8 = add(&a, &b).unwrap();
1369        assert_eq!(
1370            &format_array(r8.as_ref()),
1371            &[
1372                "1998-10-04 23:59:17".to_string(),
1373                "1960-09-29 04:00:33".to_string(),
1374                "1960-07-02 04:31:33".to_string()
1375            ]
1376        );
1377
1378        let r9 = sub(&r8, &b).unwrap();
1379        // Note: subtraction is not the inverse of addition for intervals
1380        assert_eq!(
1381            &format_array(r9.as_ref()),
1382            &[
1383                "1970-01-02 00:00:00".to_string(),
1384                "2010-04-02 04:00:20".to_string(),
1385                "1960-01-31 04:23:20".to_string()
1386            ]
1387        );
1388    }
1389
1390    #[test]
1391    fn test_timestamp() {
1392        test_timestamp_impl::<TimestampSecondType>();
1393        test_timestamp_impl::<TimestampMillisecondType>();
1394        test_timestamp_impl::<TimestampMicrosecondType>();
1395        test_timestamp_impl::<TimestampNanosecondType>();
1396    }
1397
1398    #[test]
1399    fn test_interval() {
1400        let a = IntervalYearMonthArray::from(vec![
1401            IntervalYearMonthType::make_value(32, 4),
1402            IntervalYearMonthType::make_value(32, 4),
1403        ]);
1404        let b = IntervalYearMonthArray::from(vec![
1405            IntervalYearMonthType::make_value(-4, 6),
1406            IntervalYearMonthType::make_value(-3, 23),
1407        ]);
1408        let result = add(&a, &b).unwrap();
1409        assert_eq!(
1410            result.as_ref(),
1411            &IntervalYearMonthArray::from(vec![
1412                IntervalYearMonthType::make_value(28, 10),
1413                IntervalYearMonthType::make_value(29, 27)
1414            ])
1415        );
1416        let result = sub(&a, &b).unwrap();
1417        assert_eq!(
1418            result.as_ref(),
1419            &IntervalYearMonthArray::from(vec![
1420                IntervalYearMonthType::make_value(36, -2),
1421                IntervalYearMonthType::make_value(35, -19)
1422            ])
1423        );
1424
1425        let a = IntervalDayTimeArray::from(vec![
1426            IntervalDayTimeType::make_value(32, 4),
1427            IntervalDayTimeType::make_value(32, 4),
1428        ]);
1429        let b = IntervalDayTimeArray::from(vec![
1430            IntervalDayTimeType::make_value(-4, 6),
1431            IntervalDayTimeType::make_value(-3, 23),
1432        ]);
1433        let result = add(&a, &b).unwrap();
1434        assert_eq!(
1435            result.as_ref(),
1436            &IntervalDayTimeArray::from(vec![
1437                IntervalDayTimeType::make_value(28, 10),
1438                IntervalDayTimeType::make_value(29, 27)
1439            ])
1440        );
1441        let result = sub(&a, &b).unwrap();
1442        assert_eq!(
1443            result.as_ref(),
1444            &IntervalDayTimeArray::from(vec![
1445                IntervalDayTimeType::make_value(36, -2),
1446                IntervalDayTimeType::make_value(35, -19)
1447            ])
1448        );
1449        let a = IntervalMonthDayNanoArray::from(vec![
1450            IntervalMonthDayNanoType::make_value(32, 4, 4000000000000),
1451            IntervalMonthDayNanoType::make_value(32, 4, 45463000000000000),
1452        ]);
1453        let b = IntervalMonthDayNanoArray::from(vec![
1454            IntervalMonthDayNanoType::make_value(-4, 6, 46000000000000),
1455            IntervalMonthDayNanoType::make_value(-3, 23, 3564000000000000),
1456        ]);
1457        let result = add(&a, &b).unwrap();
1458        assert_eq!(
1459            result.as_ref(),
1460            &IntervalMonthDayNanoArray::from(vec![
1461                IntervalMonthDayNanoType::make_value(28, 10, 50000000000000),
1462                IntervalMonthDayNanoType::make_value(29, 27, 49027000000000000)
1463            ])
1464        );
1465        let result = sub(&a, &b).unwrap();
1466        assert_eq!(
1467            result.as_ref(),
1468            &IntervalMonthDayNanoArray::from(vec![
1469                IntervalMonthDayNanoType::make_value(36, -2, -42000000000000),
1470                IntervalMonthDayNanoType::make_value(35, -19, 41899000000000000)
1471            ])
1472        );
1473        let a = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNano::MAX]);
1474        let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNano::ONE]);
1475        let err = add(&a, &b).unwrap_err().to_string();
1476        assert_eq!(
1477            err,
1478            "Arithmetic overflow: Overflow happened on: 2147483647 + 1"
1479        );
1480    }
1481
1482    fn test_duration_impl<T: ArrowPrimitiveType<Native = i64>>() {
1483        let a = PrimitiveArray::<T>::new(vec![1000, 4394, -3944].into(), None);
1484        let b = PrimitiveArray::<T>::new(vec![4, -5, -243].into(), None);
1485
1486        let result = add(&a, &b).unwrap();
1487        assert_eq!(result.as_primitive::<T>().values(), &[1004, 4389, -4187]);
1488        let result = sub(&a, &b).unwrap();
1489        assert_eq!(result.as_primitive::<T>().values(), &[996, 4399, -3701]);
1490
1491        let err = mul(&a, &b).unwrap_err().to_string();
1492        assert!(
1493            err.contains("Invalid duration arithmetic operation"),
1494            "{err}"
1495        );
1496
1497        let err = div(&a, &b).unwrap_err().to_string();
1498        assert!(
1499            err.contains("Invalid duration arithmetic operation"),
1500            "{err}"
1501        );
1502
1503        let err = rem(&a, &b).unwrap_err().to_string();
1504        assert!(
1505            err.contains("Invalid duration arithmetic operation"),
1506            "{err}"
1507        );
1508
1509        let a = PrimitiveArray::<T>::new(vec![i64::MAX].into(), None);
1510        let b = PrimitiveArray::<T>::new(vec![1].into(), None);
1511        let err = add(&a, &b).unwrap_err().to_string();
1512        assert_eq!(
1513            err,
1514            "Arithmetic overflow: Overflow happened on: 9223372036854775807 + 1"
1515        );
1516    }
1517
1518    #[test]
1519    fn test_duration() {
1520        test_duration_impl::<DurationSecondType>();
1521        test_duration_impl::<DurationMillisecondType>();
1522        test_duration_impl::<DurationMicrosecondType>();
1523        test_duration_impl::<DurationNanosecondType>();
1524    }
1525
1526    fn test_date_impl<T: ArrowPrimitiveType, F>(f: F)
1527    where
1528        F: Fn(NaiveDate) -> T::Native,
1529        T::Native: TryInto<i64>,
1530    {
1531        let a = PrimitiveArray::<T>::new(
1532            vec![
1533                f(NaiveDate::from_ymd_opt(1979, 1, 30).unwrap()),
1534                f(NaiveDate::from_ymd_opt(2010, 4, 3).unwrap()),
1535                f(NaiveDate::from_ymd_opt(2008, 2, 29).unwrap()),
1536            ]
1537            .into(),
1538            None,
1539        );
1540
1541        let b = IntervalYearMonthArray::from(vec![
1542            IntervalYearMonthType::make_value(34, 2),
1543            IntervalYearMonthType::make_value(3, -3),
1544            IntervalYearMonthType::make_value(-12, 4),
1545        ]);
1546
1547        let format_array = |x: &dyn Array| -> Vec<String> {
1548            x.as_primitive::<T>()
1549                .values()
1550                .into_iter()
1551                .map(|x| {
1552                    as_date::<T>((*x).try_into().ok().unwrap())
1553                        .unwrap()
1554                        .to_string()
1555                })
1556                .collect()
1557        };
1558
1559        let result = add(&a, &b).unwrap();
1560        assert_eq!(
1561            &format_array(result.as_ref()),
1562            &[
1563                "2013-03-30".to_string(),
1564                "2013-01-03".to_string(),
1565                "1996-06-29".to_string(),
1566            ]
1567        );
1568        let result = sub(&result, &b).unwrap();
1569        assert_eq!(result.as_ref(), &a);
1570
1571        let b = IntervalDayTimeArray::from(vec![
1572            IntervalDayTimeType::make_value(34, 2),
1573            IntervalDayTimeType::make_value(3, -3),
1574            IntervalDayTimeType::make_value(-12, 4),
1575        ]);
1576
1577        let result = add(&a, &b).unwrap();
1578        assert_eq!(
1579            &format_array(result.as_ref()),
1580            &[
1581                "1979-03-05".to_string(),
1582                "2010-04-06".to_string(),
1583                "2008-02-17".to_string(),
1584            ]
1585        );
1586        let result = sub(&result, &b).unwrap();
1587        assert_eq!(result.as_ref(), &a);
1588
1589        let b = IntervalMonthDayNanoArray::from(vec![
1590            IntervalMonthDayNanoType::make_value(34, 2, -34353534),
1591            IntervalMonthDayNanoType::make_value(3, -3, 2443),
1592            IntervalMonthDayNanoType::make_value(-12, 4, 2323242423232),
1593        ]);
1594
1595        let result = add(&a, &b).unwrap();
1596        assert_eq!(
1597            &format_array(result.as_ref()),
1598            &[
1599                "1981-12-02".to_string(),
1600                "2010-06-30".to_string(),
1601                "2007-03-04".to_string(),
1602            ]
1603        );
1604        let result = sub(&result, &b).unwrap();
1605        assert_eq!(
1606            &format_array(result.as_ref()),
1607            &[
1608                "1979-01-31".to_string(),
1609                "2010-04-02".to_string(),
1610                "2008-02-29".to_string(),
1611            ]
1612        );
1613    }
1614
1615    #[test]
1616    fn test_date() {
1617        test_date_impl::<Date32Type, _>(Date32Type::from_naive_date);
1618        test_date_impl::<Date64Type, _>(Date64Type::from_naive_date);
1619
1620        let a = Date32Array::from(vec![i32::MIN, i32::MAX, 23, 7684]);
1621        let b = Date32Array::from(vec![i32::MIN, i32::MIN, -2, 45]);
1622        let result = sub(&a, &b).unwrap();
1623        assert_eq!(
1624            result.as_primitive::<DurationSecondType>().values(),
1625            &[0, 371085174288000, 2160000, 660009600]
1626        );
1627
1628        let a = Date64Array::from(vec![4343, 76676, 3434]);
1629        let b = Date64Array::from(vec![3, -5, 5]);
1630        let result = sub(&a, &b).unwrap();
1631        assert_eq!(
1632            result.as_primitive::<DurationMillisecondType>().values(),
1633            &[4340, 76681, 3429]
1634        );
1635
1636        let a = Date64Array::from(vec![i64::MAX]);
1637        let b = Date64Array::from(vec![-1]);
1638        let err = sub(&a, &b).unwrap_err().to_string();
1639        assert_eq!(
1640            err,
1641            "Arithmetic overflow: Overflow happened on: 9223372036854775807 - -1"
1642        );
1643    }
1644
1645    #[test]
1646    fn test_date32_to_naive_date_opt_boundaries() {
1647        assert_eq!(MAX_VALID_DAYS, 95026236);
1648        assert_eq!(MIN_VALID_DAYS, -96465292);
1649
1650        // Valid boundary dates work
1651        assert!(Date32Type::to_naive_date_opt(MAX_VALID_DAYS).is_some());
1652        assert!(Date32Type::to_naive_date_opt(MIN_VALID_DAYS).is_some());
1653
1654        // Beyond boundaries fail
1655        assert!(Date32Type::to_naive_date_opt(MAX_VALID_DAYS + 1).is_none());
1656        assert!(Date32Type::to_naive_date_opt(MIN_VALID_DAYS - 1).is_none());
1657
1658        // Extreme values fail
1659        assert!(Date32Type::to_naive_date_opt(i32::MAX).is_none());
1660        assert!(Date32Type::to_naive_date_opt(i32::MIN).is_none());
1661
1662        // Common values work
1663        assert!(Date32Type::to_naive_date_opt(0).is_some());
1664        assert!(Date32Type::to_naive_date_opt(date_to_days(YEAR_2000)).is_some());
1665    }
1666
1667    #[test]
1668    fn test_date64_to_naive_date_opt_boundaries() {
1669        const MS_PER_DAY: i64 = 24 * 60 * 60 * 1000;
1670
1671        // Verify boundary millisecond values
1672        assert_eq!(MAX_VALID_MILLIS, 8210266790400000i64);
1673        assert_eq!(MIN_VALID_MILLIS, -8334601228800000i64);
1674
1675        // Valid boundary dates work
1676        assert!(Date64Type::to_naive_date_opt(MAX_VALID_MILLIS).is_some());
1677        assert!(Date64Type::to_naive_date_opt(MIN_VALID_MILLIS).is_some());
1678
1679        // Beyond boundaries fail
1680        assert!(Date64Type::to_naive_date_opt(MAX_VALID_MILLIS + MS_PER_DAY).is_none());
1681        assert!(Date64Type::to_naive_date_opt(MIN_VALID_MILLIS - MS_PER_DAY).is_none());
1682
1683        // Extreme values fail
1684        assert!(Date64Type::to_naive_date_opt(i64::MAX).is_none());
1685        assert!(Date64Type::to_naive_date_opt(i64::MIN).is_none());
1686
1687        // Common values work
1688        assert!(Date64Type::to_naive_date_opt(0).is_some());
1689        assert!(Date64Type::to_naive_date_opt(date_to_millis(YEAR_2000)).is_some());
1690    }
1691
1692    macro_rules! test_year_month_ops {
1693        ($type:ty, $date_fn:expr) => {{
1694            let date = $date_fn(YEAR_2000);
1695
1696            // Normal operations succeed
1697            assert!(
1698                <$type>::add_year_months_opt(date, 120).is_some(),
1699                "add_year_months: normal add"
1700            );
1701            assert!(
1702                <$type>::add_year_months_opt(date, 0).is_some(),
1703                "add_year_months: zero interval"
1704            );
1705            assert!(
1706                <$type>::subtract_year_months_opt(date, 120).is_some(),
1707                "subtract_year_months: normal subtract"
1708            );
1709            assert!(
1710                <$type>::subtract_year_months_opt(date, 0).is_some(),
1711                "subtract_year_months: zero interval"
1712            );
1713
1714            // Large but valid years work
1715            let large_year = $date_fn(NaiveDate::from_ymd_opt(5000, 1, 1).unwrap());
1716            let neg_year = $date_fn(NaiveDate::from_ymd_opt(-5000, 12, 31).unwrap());
1717            assert!(
1718                <$type>::add_year_months_opt(large_year, 12).is_some(),
1719                "add_year_months: large year"
1720            );
1721            assert!(
1722                <$type>::add_year_months_opt(neg_year, -12).is_some(),
1723                "add_year_months: negative year"
1724            );
1725            assert!(
1726                <$type>::subtract_year_months_opt(large_year, 12).is_some(),
1727                "subtract_year_months: large year"
1728            );
1729            assert!(
1730                <$type>::subtract_year_months_opt(neg_year, -12).is_some(),
1731                "subtract_year_months: negative year"
1732            );
1733
1734            // Overflow handling
1735            assert!(
1736                <$type>::subtract_year_months_opt($date_fn(MIN_VALID_DATE), 1).is_none(),
1737                "subtract_year_months: overflow days from min"
1738            );
1739            assert!(
1740                <$type>::subtract_year_months_opt($date_fn(MAX_VALID_DATE), -1).is_none(),
1741                "subtract_year_months: overflow neg days from max"
1742            );
1743            assert!(
1744                <$type>::add_year_months_opt($date_fn(MAX_VALID_DATE), 1).is_none(),
1745                "add_year_months: overflow days"
1746            );
1747            assert!(
1748                <$type>::add_year_months_opt($date_fn(MIN_VALID_DATE), -1).is_none(),
1749                "add_year_months: overflow neg days"
1750            );
1751        }};
1752    }
1753
1754    #[test]
1755    fn test_date_year_month_operations() {
1756        test_year_month_ops!(Date32Type, date_to_days);
1757        test_year_month_ops!(Date64Type, date_to_millis);
1758    }
1759
1760    macro_rules! test_day_time_ops {
1761        ($type:ty, $date_fn:expr) => {{
1762            let date = $date_fn(YEAR_2000);
1763
1764            // Moderate intervals succeed
1765            assert!(
1766                <$type>::add_day_time_opt(date, IntervalDayTime::new(30, 0)).is_some(),
1767                "add_day_time: +30 days"
1768            );
1769            assert!(
1770                <$type>::add_day_time_opt(date, IntervalDayTime::new(-30, 0)).is_some(),
1771                "add_day_time: -30 days"
1772            );
1773            assert!(
1774                <$type>::add_day_time_opt(date, IntervalDayTime::new(1000, 12345)).is_some(),
1775                "add_day_time: normal"
1776            );
1777            assert!(
1778                <$type>::subtract_day_time_opt(date, IntervalDayTime::new(30, 0)).is_some(),
1779                "subtract_day_time: +30 days"
1780            );
1781            assert!(
1782                <$type>::subtract_day_time_opt(date, IntervalDayTime::new(-30, 0)).is_some(),
1783                "subtract_day_time: -30 days"
1784            );
1785            assert!(
1786                <$type>::subtract_day_time_opt(date, IntervalDayTime::new(1000, 12345)).is_some(),
1787                "subtract_day_time: normal"
1788            );
1789
1790            // Overflow handling - subtract
1791            assert!(
1792                <$type>::subtract_day_time_opt(
1793                    $date_fn(MIN_VALID_DATE),
1794                    IntervalDayTime::new(1, 0)
1795                )
1796                .is_none(),
1797                "subtract_day_time: overflow days from min"
1798            );
1799            assert!(
1800                <$type>::subtract_day_time_opt(
1801                    $date_fn(MAX_VALID_DATE),
1802                    IntervalDayTime::new(-1, 0)
1803                )
1804                .is_none(),
1805                "subtract_day_time: overflow neg days from max"
1806            );
1807
1808            // Overflow handling - add
1809            assert!(
1810                <$type>::add_day_time_opt($date_fn(MAX_VALID_DATE), IntervalDayTime::new(1, 0))
1811                    .is_none(),
1812                "add_day_time: overflow days"
1813            );
1814            assert!(
1815                <$type>::add_day_time_opt($date_fn(MIN_VALID_DATE), IntervalDayTime::new(-1, 0))
1816                    .is_none(),
1817                "add_day_time: overflow neg days"
1818            );
1819
1820            // Extreme intervals fail
1821            assert!(
1822                <$type>::add_day_time_opt(
1823                    $date_fn(EPOCH),
1824                    IntervalDayTime::new(i32::MAX, i32::MAX)
1825                )
1826                .is_none(),
1827                "add_day_time: max interval"
1828            );
1829            assert!(
1830                <$type>::add_day_time_opt(
1831                    $date_fn(EPOCH),
1832                    IntervalDayTime::new(i32::MIN, i32::MIN)
1833                )
1834                .is_none(),
1835                "add_day_time: min interval"
1836            );
1837            assert!(
1838                <$type>::subtract_day_time_opt(
1839                    $date_fn(EPOCH),
1840                    IntervalDayTime::new(i32::MAX, i32::MAX)
1841                )
1842                .is_none(),
1843                "subtract_day_time: max interval"
1844            );
1845            assert!(
1846                <$type>::subtract_day_time_opt(
1847                    $date_fn(EPOCH),
1848                    IntervalDayTime::new(i32::MIN, i32::MIN)
1849                )
1850                .is_none(),
1851                "subtract_day_time: min interval"
1852            );
1853        }};
1854    }
1855
1856    #[test]
1857    fn test_date_day_time_operations() {
1858        test_day_time_ops!(Date32Type, date_to_days);
1859        test_day_time_ops!(Date64Type, date_to_millis);
1860    }
1861
1862    macro_rules! test_month_day_nano_ops {
1863        ($type:ty, $date_fn:expr) => {{
1864            let date = $date_fn(YEAR_2000);
1865            let zero = IntervalMonthDayNano::new(0, 0, 0);
1866
1867            // Normal operations succeed
1868            assert!(
1869                <$type>::add_month_day_nano_opt(date, IntervalMonthDayNano::new(1, 30, 0))
1870                    .is_some(),
1871                "add_month_day_nano: +1mo +30d"
1872            );
1873            assert!(
1874                <$type>::add_month_day_nano_opt(date, IntervalMonthDayNano::new(-1, -30, 0))
1875                    .is_some(),
1876                "add_month_day_nano: -1mo -30d"
1877            );
1878            assert!(
1879                <$type>::add_month_day_nano_opt(date, zero).is_some(),
1880                "add_month_day_nano: zero interval"
1881            );
1882            assert!(
1883                <$type>::add_month_day_nano_opt(
1884                    date,
1885                    IntervalMonthDayNano::new(2, 10, 123_456_789_000)
1886                )
1887                .is_some(),
1888                "add_month_day_nano: normal"
1889            );
1890            assert!(
1891                <$type>::subtract_month_day_nano_opt(date, IntervalMonthDayNano::new(1, 30, 0))
1892                    .is_some(),
1893                "subtract_month_day_nano: +1mo +30d"
1894            );
1895            assert!(
1896                <$type>::subtract_month_day_nano_opt(date, IntervalMonthDayNano::new(-1, -30, 0))
1897                    .is_some(),
1898                "subtract_month_day_nano: -1mo -30d"
1899            );
1900            assert!(
1901                <$type>::subtract_month_day_nano_opt(date, zero).is_some(),
1902                "subtract_month_day_nano: zero interval"
1903            );
1904            assert!(
1905                <$type>::subtract_month_day_nano_opt(
1906                    date,
1907                    IntervalMonthDayNano::new(2, 10, 123_456_789_000)
1908                )
1909                .is_some(),
1910                "subtract_month_day_nano: normal"
1911            );
1912
1913            // Overflow handling - subtract
1914            assert!(
1915                <$type>::subtract_month_day_nano_opt(
1916                    $date_fn(MIN_VALID_DATE),
1917                    IntervalMonthDayNano::new(0, 1, 0)
1918                )
1919                .is_none(),
1920                "subtract_month_day_nano: overflow days from min"
1921            );
1922            assert!(
1923                <$type>::subtract_month_day_nano_opt(
1924                    $date_fn(MAX_VALID_DATE),
1925                    IntervalMonthDayNano::new(0, -1, 0)
1926                )
1927                .is_none(),
1928                "subtract_month_day_nano: overflow neg days from max"
1929            );
1930
1931            // Overflow handling - add
1932            assert!(
1933                <$type>::add_month_day_nano_opt(
1934                    $date_fn(MAX_VALID_DATE),
1935                    IntervalMonthDayNano::new(0, 1, 0)
1936                )
1937                .is_none(),
1938                "add_month_day_nano: overflow days"
1939            );
1940            assert!(
1941                <$type>::add_month_day_nano_opt(
1942                    $date_fn(MIN_VALID_DATE),
1943                    IntervalMonthDayNano::new(0, -1, 0)
1944                )
1945                .is_none(),
1946                "add_month_day_nano: overflow neg days"
1947            );
1948
1949            // Nanosecond precision works
1950            assert!(
1951                <$type>::add_month_day_nano_opt(date, IntervalMonthDayNano::new(0, 0, 999_999_999))
1952                    .is_some(),
1953                "add_month_day_nano: nanos"
1954            );
1955            assert!(
1956                <$type>::subtract_month_day_nano_opt(
1957                    date,
1958                    IntervalMonthDayNano::new(0, 0, 999_999_999)
1959                )
1960                .is_some(),
1961                "subtract_month_day_nano: nanos"
1962            );
1963            // 1 day in nanos
1964            assert!(
1965                <$type>::add_month_day_nano_opt(
1966                    date,
1967                    IntervalMonthDayNano::new(0, 0, 86_400_000_000_000)
1968                )
1969                .is_some(),
1970                "add_month_day_nano: 1 day nanos"
1971            );
1972            assert!(
1973                <$type>::subtract_month_day_nano_opt(
1974                    date,
1975                    IntervalMonthDayNano::new(0, 0, 86_400_000_000_000)
1976                )
1977                .is_some(),
1978                "subtract_month_day_nano: 1 day nanos"
1979            );
1980        }};
1981    }
1982
1983    #[test]
1984    fn test_date_month_day_nano_operations() {
1985        test_month_day_nano_ops!(Date32Type, date_to_days);
1986        test_month_day_nano_ops!(Date64Type, date_to_millis);
1987    }
1988}