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
551impl DateOp for Date32Type {
552    fn add_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> {
553        // Date32Type functions don't have _opt variants and should be safe
554        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
640/// Arithmetic trait for interval arrays
641trait 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
694/// Perform arithmetic operation on an interval array
695fn 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
735/// Perform arithmetic operation on a date array
736fn 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
808/// Perform arithmetic operation on decimal arrays
809fn 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    // Follow the Hive decimal arithmetic rules
828    // https://cwiki.apache.org/confluence/download/attachments/27362075/Hive_Decimal_Precision_Scale_Support.pdf
829    let array: PrimitiveArray<T> = match op {
830        Op::Add | Op::AddWrapping | Op::Sub | Op::SubWrapping => {
831            // max(s1, s2)
832            let result_scale = *s1.max(s2);
833
834            // max(s1, s2) + max(p1-s1, p2-s2) + 1
835            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                // SQL standard says that if the resulting scale of a multiply operation goes
871                // beyond the maximum, rounding is not acceptable and thus an error occurs
872                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            // Follow postgres and MySQL adding a fixed scale increment of 4
886            // s1 + 4
887            let result_scale = s1.saturating_add(4).min(T::MAX_SCALE);
888            let mul_pow = result_scale - s1 + s2;
889
890            // p1 - s1 + s2 + result_scale
891            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            // max(s1, s2)
917            let result_scale = *s1.max(s2);
918            // min(p1-s1, p2 -s2) + max( s1,s2 )
919            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        // 0.015 7.842 -0.577 0.334 -0.078 0.003
1215        let a = Decimal128Array::from(vec![15, 0, -577, 334, -78, 3])
1216            .with_precision_and_scale(12, 3)
1217            .unwrap();
1218
1219        // 5.4 0 -35.6 0.3 0.6 7.45
1220        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        // Note: subtraction is not the inverse of addition for intervals
1382        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        // Date64Type::to_naive_date_opt has boundaries determined by NaiveDate's supported range.
1652        // The valid date range is from January 1, -262143 to December 31, 262142 (Gregorian calendar).
1653
1654        let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
1655        let ms_per_day = 24 * 60 * 60 * 1000i64;
1656
1657        // Define the boundary dates using NaiveDate::from_ymd_opt
1658        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        // Calculate their millisecond values from epoch
1662        let max_valid_millis = (max_valid_date - epoch).num_milliseconds();
1663        let min_valid_millis = (min_valid_date - epoch).num_milliseconds();
1664
1665        // Verify these match the expected boundaries in milliseconds
1666        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        // Test that the boundary dates work
1676        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        // Test that one day beyond the boundaries fails
1686        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        // Test some values well within the valid range
1696        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        // Test extreme values that definitely fail due to Duration constraints
1708        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        // Test normal case within valid range
1725        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        // Test with moderate years that are within chrono's safe range
1733        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        // Test with extreme input values that would cause overflow
1748        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        // Test edge case: adding zero should always work for valid dates
1758        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        // Test with a date far from the boundary but still testing the function
1772        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        // Adding 30 days should succeed
1776        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        // Adding a very large number of days should fail
1783        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        // Test with a date far from the boundary in the negative direction
1790        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        // Subtracting 30 days should succeed
1794        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        // Subtracting a very large number of days should fail
1801        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        // Test normal case within valid range
1808        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        // Test with extreme input values that would cause overflow
1817        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        // Test with extreme interval values
1828        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        // Test millisecond overflow within a day
1841        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        // Test with a large year that is still within chrono's safe range
1856        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        // Adding 1 month and 30 days should succeed
1860        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        // Test normal case within valid range
1867        let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
1868        let year_2000_millis = (year_2000 - epoch).num_milliseconds();
1869
1870        // Test edge case: adding zero should always work for valid dates
1871        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        // Test with a negative year that is still within chrono's safe range
1878        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        // Subtracting 1 month and 30 days should succeed
1882        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        // Test with extreme input values that would cause overflow
1889        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        // Test with extreme input values that would cause overflow
1902        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        // Test with invalid timestamp input (the _opt function should handle these gracefully)
1914
1915        // Test nanosecond precision (should not affect boundary since it's < 1ms)
1916        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        // Test large nanosecond values that convert to milliseconds
1923        let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); // 1 day in nanos
1924        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        // Test with a negative year that is still within chrono's safe range
1937        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        // Subtracting 12 months should succeed
1941        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        // Test normal case within valid range
1947        let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
1948        let year_2000_millis = (year_2000 - epoch).num_milliseconds();
1949
1950        // Test edge case: subtracting zero should always work for valid dates
1951        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        // Test with a large year that is still within chrono's safe range
1957        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        // Adding 12 months (subtracting negative) should succeed
1961        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        // Test with extreme input values that would cause overflow
1967        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        // Test with extreme input values that would cause overflow
1978        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        // Test edge case: subtracting zero should always work for valid dates
1988        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        // Test with a date far from the boundary in the negative direction
2004        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        // Subtracting 30 days should succeed
2008        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        // Subtracting a very large number of days should fail
2015        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        // Test with a date far from the boundary but still testing the function
2022        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        // Adding 30 days (subtracting negative) should succeed
2026        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        // Adding a very large number of days should fail
2033        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        // Test normal case within valid range
2040        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        // Test with extreme input values that would cause overflow
2049        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        // Test with extreme interval values
2060        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        // Test millisecond precision
2073        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        // Test edge case: subtracting zero should always work for valid dates
2080        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        // Test with a negative year that is still within chrono's safe range
2097        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        // Subtracting 1 month and 30 days should succeed
2101        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        // Test normal case within valid range
2108        let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
2109        let year_2000_millis = (year_2000 - epoch).num_milliseconds();
2110
2111        // Test edge case: subtracting zero should always work for valid dates
2112        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        // Test with a large year that is still within chrono's safe range
2119        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        // Adding 1 month and 30 days (subtracting negative) should succeed
2123        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        // Test with extreme input values that would cause overflow
2130        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        // Test with extreme input values that would cause overflow
2143        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        // Test nanosecond precision (should not affect boundary since it's < 1ms)
2155        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        // Test large nanosecond values that convert to milliseconds
2162        let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); // 1 day in nanos
2163        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        // Test edge case: subtracting zero should always work for valid dates
2170        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}