/Users/andrewlamb/Software/arrow-rs/arrow-arith/src/numeric.rs
Line | Count | Source |
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 | | |
20 | | use std::cmp::Ordering; |
21 | | use std::fmt::Formatter; |
22 | | use std::sync::Arc; |
23 | | |
24 | | use arrow_array::cast::AsArray; |
25 | | use arrow_array::timezone::Tz; |
26 | | use arrow_array::types::*; |
27 | | use arrow_array::*; |
28 | | use arrow_buffer::{ArrowNativeType, IntervalDayTime, IntervalMonthDayNano}; |
29 | | use arrow_schema::{ArrowError, DataType, IntervalUnit, TimeUnit}; |
30 | | |
31 | | use crate::arity::{binary, try_binary}; |
32 | | |
33 | | /// Perform `lhs + rhs`, returning an error on overflow |
34 | 0 | pub fn add(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
35 | 0 | arithmetic_op(Op::Add, lhs, rhs) |
36 | 0 | } |
37 | | |
38 | | /// Perform `lhs + rhs`, wrapping on overflow for [`DataType::is_integer`] |
39 | 0 | pub fn add_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
40 | 0 | arithmetic_op(Op::AddWrapping, lhs, rhs) |
41 | 0 | } |
42 | | |
43 | | /// Perform `lhs - rhs`, returning an error on overflow |
44 | 0 | pub fn sub(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
45 | 0 | arithmetic_op(Op::Sub, lhs, rhs) |
46 | 0 | } |
47 | | |
48 | | /// Perform `lhs - rhs`, wrapping on overflow for [`DataType::is_integer`] |
49 | 0 | pub fn sub_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
50 | 0 | arithmetic_op(Op::SubWrapping, lhs, rhs) |
51 | 0 | } |
52 | | |
53 | | /// Perform `lhs * rhs`, returning an error on overflow |
54 | 0 | pub fn mul(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
55 | 0 | arithmetic_op(Op::Mul, lhs, rhs) |
56 | 0 | } |
57 | | |
58 | | /// Perform `lhs * rhs`, wrapping on overflow for [`DataType::is_integer`] |
59 | 0 | pub fn mul_wrapping(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
60 | 0 | arithmetic_op(Op::MulWrapping, lhs, rhs) |
61 | 0 | } |
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 |
67 | 0 | pub fn div(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
68 | 0 | arithmetic_op(Op::Div, lhs, rhs) |
69 | 0 | } |
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 |
77 | 0 | pub fn rem(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> { |
78 | 0 | arithmetic_op(Op::Rem, lhs, rhs) |
79 | 0 | } |
80 | | |
81 | | macro_rules! neg_checked { |
82 | | ($t:ty, $a:ident) => {{ |
83 | | let array = $a |
84 | | .as_primitive::<$t>() |
85 | 0 | .try_unary::<_, $t, _>(|x| x.neg_checked())?; |
86 | | Ok(Arc::new(array)) |
87 | | }}; |
88 | | } |
89 | | |
90 | | macro_rules! neg_wrapping { |
91 | | ($t:ty, $a:ident) => {{ |
92 | 0 | 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()] |
101 | 0 | pub fn neg(array: &dyn Array) -> Result<ArrayRef, ArrowError> { |
102 | | use DataType::*; |
103 | | use IntervalUnit::*; |
104 | | use TimeUnit::*; |
105 | | |
106 | 0 | match array.data_type() { |
107 | 0 | Int8 => neg_checked!(Int8Type, array), |
108 | 0 | Int16 => neg_checked!(Int16Type, array), |
109 | 0 | Int32 => neg_checked!(Int32Type, array), |
110 | 0 | Int64 => neg_checked!(Int64Type, array), |
111 | 0 | Float16 => neg_wrapping!(Float16Type, array), |
112 | 0 | Float32 => neg_wrapping!(Float32Type, array), |
113 | 0 | Float64 => neg_wrapping!(Float64Type, array), |
114 | 0 | Decimal32(p, s) => { |
115 | 0 | let a = array |
116 | 0 | .as_primitive::<Decimal32Type>() |
117 | 0 | .try_unary::<_, Decimal32Type, _>(|x| x.neg_checked())?; |
118 | | |
119 | 0 | Ok(Arc::new(a.with_precision_and_scale(*p, *s)?)) |
120 | | } |
121 | 0 | Decimal64(p, s) => { |
122 | 0 | let a = array |
123 | 0 | .as_primitive::<Decimal64Type>() |
124 | 0 | .try_unary::<_, Decimal64Type, _>(|x| x.neg_checked())?; |
125 | | |
126 | 0 | Ok(Arc::new(a.with_precision_and_scale(*p, *s)?)) |
127 | | } |
128 | 0 | Decimal128(p, s) => { |
129 | 0 | let a = array |
130 | 0 | .as_primitive::<Decimal128Type>() |
131 | 0 | .try_unary::<_, Decimal128Type, _>(|x| x.neg_checked())?; |
132 | | |
133 | 0 | Ok(Arc::new(a.with_precision_and_scale(*p, *s)?)) |
134 | | } |
135 | 0 | Decimal256(p, s) => { |
136 | 0 | let a = array |
137 | 0 | .as_primitive::<Decimal256Type>() |
138 | 0 | .try_unary::<_, Decimal256Type, _>(|x| x.neg_checked())?; |
139 | | |
140 | 0 | Ok(Arc::new(a.with_precision_and_scale(*p, *s)?)) |
141 | | } |
142 | 0 | Duration(Second) => neg_checked!(DurationSecondType, array), |
143 | 0 | Duration(Millisecond) => neg_checked!(DurationMillisecondType, array), |
144 | 0 | Duration(Microsecond) => neg_checked!(DurationMicrosecondType, array), |
145 | 0 | Duration(Nanosecond) => neg_checked!(DurationNanosecondType, array), |
146 | 0 | Interval(YearMonth) => neg_checked!(IntervalYearMonthType, array), |
147 | | Interval(DayTime) => { |
148 | 0 | let a = array |
149 | 0 | .as_primitive::<IntervalDayTimeType>() |
150 | 0 | .try_unary::<_, IntervalDayTimeType, ArrowError>(|x| { |
151 | 0 | let (days, ms) = IntervalDayTimeType::to_parts(x); |
152 | 0 | Ok(IntervalDayTimeType::make_value( |
153 | 0 | days.neg_checked()?, |
154 | 0 | ms.neg_checked()?, |
155 | | )) |
156 | 0 | })?; |
157 | 0 | Ok(Arc::new(a)) |
158 | | } |
159 | | Interval(MonthDayNano) => { |
160 | 0 | let a = array |
161 | 0 | .as_primitive::<IntervalMonthDayNanoType>() |
162 | 0 | .try_unary::<_, IntervalMonthDayNanoType, ArrowError>(|x| { |
163 | 0 | let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(x); |
164 | 0 | Ok(IntervalMonthDayNanoType::make_value( |
165 | 0 | months.neg_checked()?, |
166 | 0 | days.neg_checked()?, |
167 | 0 | nanos.neg_checked()?, |
168 | | )) |
169 | 0 | })?; |
170 | 0 | Ok(Arc::new(a)) |
171 | | } |
172 | 0 | t => Err(ArrowError::InvalidArgumentError(format!( |
173 | 0 | "Invalid arithmetic operation: !{t}" |
174 | 0 | ))), |
175 | | } |
176 | 0 | } |
177 | | |
178 | | /// Negates each element of `array`, wrapping on overflow for [`DataType::is_integer`] |
179 | 0 | pub fn neg_wrapping(array: &dyn Array) -> Result<ArrayRef, ArrowError> { |
180 | 0 | downcast_integer! { |
181 | 0 | array.data_type() => (neg_wrapping, array), |
182 | 0 | _ => neg(array), |
183 | | } |
184 | 0 | } |
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)] |
190 | | enum Op { |
191 | | AddWrapping, |
192 | | Add, |
193 | | SubWrapping, |
194 | | Sub, |
195 | | MulWrapping, |
196 | | Mul, |
197 | | Div, |
198 | | Rem, |
199 | | } |
200 | | |
201 | | impl std::fmt::Display for Op { |
202 | 0 | fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { |
203 | 0 | match self { |
204 | 0 | Op::AddWrapping | Op::Add => write!(f, "+"), |
205 | 0 | Op::SubWrapping | Op::Sub => write!(f, "-"), |
206 | 0 | Op::MulWrapping | Op::Mul => write!(f, "*"), |
207 | 0 | Op::Div => write!(f, "/"), |
208 | 0 | Op::Rem => write!(f, "%"), |
209 | | } |
210 | 0 | } |
211 | | } |
212 | | |
213 | | impl Op { |
214 | 0 | fn commutative(&self) -> bool { |
215 | 0 | matches!(self, Self::Add | Self::AddWrapping) |
216 | 0 | } |
217 | | } |
218 | | |
219 | | /// Dispatch the given `op` to the appropriate specialized kernel |
220 | 0 | fn 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 | 0 | let (l, l_scalar) = lhs.get(); |
232 | 0 | let (r, r_scalar) = rhs.get(); |
233 | 0 | downcast_integer! { |
234 | 0 | l.data_type(), r.data_type() => (integer_helper, op, l, l_scalar, r, r_scalar), |
235 | 0 | (Float16, Float16) => float_op::<Float16Type>(op, l, l_scalar, r, r_scalar), |
236 | 0 | (Float32, Float32) => float_op::<Float32Type>(op, l, l_scalar, r, r_scalar), |
237 | 0 | (Float64, Float64) => float_op::<Float64Type>(op, l, l_scalar, r, r_scalar), |
238 | 0 | (Timestamp(Second, _), _) => timestamp_op::<TimestampSecondType>(op, l, l_scalar, r, r_scalar), |
239 | 0 | (Timestamp(Millisecond, _), _) => timestamp_op::<TimestampMillisecondType>(op, l, l_scalar, r, r_scalar), |
240 | 0 | (Timestamp(Microsecond, _), _) => timestamp_op::<TimestampMicrosecondType>(op, l, l_scalar, r, r_scalar), |
241 | 0 | (Timestamp(Nanosecond, _), _) => timestamp_op::<TimestampNanosecondType>(op, l, l_scalar, r, r_scalar), |
242 | 0 | (Duration(Second), Duration(Second)) => duration_op::<DurationSecondType>(op, l, l_scalar, r, r_scalar), |
243 | 0 | (Duration(Millisecond), Duration(Millisecond)) => duration_op::<DurationMillisecondType>(op, l, l_scalar, r, r_scalar), |
244 | 0 | (Duration(Microsecond), Duration(Microsecond)) => duration_op::<DurationMicrosecondType>(op, l, l_scalar, r, r_scalar), |
245 | 0 | (Duration(Nanosecond), Duration(Nanosecond)) => duration_op::<DurationNanosecondType>(op, l, l_scalar, r, r_scalar), |
246 | 0 | (Interval(YearMonth), Interval(YearMonth)) => interval_op::<IntervalYearMonthType>(op, l, l_scalar, r, r_scalar), |
247 | 0 | (Interval(DayTime), Interval(DayTime)) => interval_op::<IntervalDayTimeType>(op, l, l_scalar, r, r_scalar), |
248 | 0 | (Interval(MonthDayNano), Interval(MonthDayNano)) => interval_op::<IntervalMonthDayNanoType>(op, l, l_scalar, r, r_scalar), |
249 | 0 | (Date32, _) => date_op::<Date32Type>(op, l, l_scalar, r, r_scalar), |
250 | 0 | (Date64, _) => date_op::<Date64Type>(op, l, l_scalar, r, r_scalar), |
251 | 0 | (Decimal32(_, _), Decimal32(_, _)) => decimal_op::<Decimal32Type>(op, l, l_scalar, r, r_scalar), |
252 | 0 | (Decimal64(_, _), Decimal64(_, _)) => decimal_op::<Decimal64Type>(op, l, l_scalar, r, r_scalar), |
253 | 0 | (Decimal128(_, _), Decimal128(_, _)) => decimal_op::<Decimal128Type>(op, l, l_scalar, r, r_scalar), |
254 | 0 | (Decimal256(_, _), Decimal256(_, _)) => decimal_op::<Decimal256Type>(op, l, l_scalar, r, r_scalar), |
255 | 0 | (l_t, r_t) => match (l_t, r_t) { |
256 | 0 | (Duration(_) | Interval(_), Date32 | Date64 | Timestamp(_, _)) if op.commutative() => { |
257 | 0 | arithmetic_op(op, rhs, lhs) |
258 | | } |
259 | 0 | _ => Err(ArrowError::InvalidArgumentError( |
260 | 0 | format!("Invalid arithmetic operation: {l_t} {op} {r_t}") |
261 | 0 | )) |
262 | | } |
263 | | } |
264 | 0 | } |
265 | | |
266 | | /// Perform an infallible binary operation on potentially scalar inputs |
267 | | macro_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 | 0 | (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 | 0 | (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 |
284 | | macro_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 |
292 | | macro_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 | 0 | (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 | 0 | (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 |
309 | | macro_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 |
317 | 0 | fn integer_op<T: ArrowPrimitiveType>( |
318 | 0 | op: Op, |
319 | 0 | l: &dyn Array, |
320 | 0 | l_s: bool, |
321 | 0 | r: &dyn Array, |
322 | 0 | r_s: bool, |
323 | 0 | ) -> Result<ArrayRef, ArrowError> { |
324 | 0 | let l = l.as_primitive::<T>(); |
325 | 0 | let r = r.as_primitive::<T>(); |
326 | 0 | let array: PrimitiveArray<T> = match op { |
327 | 0 | Op::AddWrapping => op!(l, l_s, r, r_s, l.add_wrapping(r)), |
328 | 0 | Op::Add => try_op!(l, l_s, r, r_s, l.add_checked(r)), |
329 | 0 | Op::SubWrapping => op!(l, l_s, r, r_s, l.sub_wrapping(r)), |
330 | 0 | Op::Sub => try_op!(l, l_s, r, r_s, l.sub_checked(r)), |
331 | 0 | Op::MulWrapping => op!(l, l_s, r, r_s, l.mul_wrapping(r)), |
332 | 0 | Op::Mul => try_op!(l, l_s, r, r_s, l.mul_checked(r)), |
333 | 0 | Op::Div => try_op!(l, l_s, r, r_s, l.div_checked(r)), |
334 | 0 | Op::Rem => try_op!(l, l_s, r, r_s, { |
335 | 0 | if r.is_zero() { |
336 | 0 | Err(ArrowError::DivideByZero) |
337 | | } else { |
338 | 0 | Ok(l.mod_wrapping(r)) |
339 | | } |
340 | | }), |
341 | | }; |
342 | 0 | Ok(Arc::new(array)) |
343 | 0 | } |
344 | | |
345 | | /// Perform an arithmetic operation on floats |
346 | 0 | fn float_op<T: ArrowPrimitiveType>( |
347 | 0 | op: Op, |
348 | 0 | l: &dyn Array, |
349 | 0 | l_s: bool, |
350 | 0 | r: &dyn Array, |
351 | 0 | r_s: bool, |
352 | 0 | ) -> Result<ArrayRef, ArrowError> { |
353 | 0 | let l = l.as_primitive::<T>(); |
354 | 0 | let r = r.as_primitive::<T>(); |
355 | 0 | let array: PrimitiveArray<T> = match op { |
356 | 0 | Op::AddWrapping | Op::Add => op!(l, l_s, r, r_s, l.add_wrapping(r)), |
357 | 0 | Op::SubWrapping | Op::Sub => op!(l, l_s, r, r_s, l.sub_wrapping(r)), |
358 | 0 | Op::MulWrapping | Op::Mul => op!(l, l_s, r, r_s, l.mul_wrapping(r)), |
359 | 0 | Op::Div => op!(l, l_s, r, r_s, l.div_wrapping(r)), |
360 | 0 | Op::Rem => op!(l, l_s, r, r_s, l.mod_wrapping(r)), |
361 | | }; |
362 | 0 | Ok(Arc::new(array)) |
363 | 0 | } |
364 | | |
365 | | /// Arithmetic trait for timestamp arrays |
366 | | trait 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 | | |
378 | | macro_rules! timestamp { |
379 | | ($t:ty, $d:ty) => { |
380 | | impl TimestampOp for $t { |
381 | | type Duration = $d; |
382 | | |
383 | 0 | fn add_year_month(left: i64, right: i32, tz: Tz) -> Option<i64> { |
384 | 0 | Self::add_year_months(left, right, tz) |
385 | 0 | } |
386 | | |
387 | 0 | fn add_day_time(left: i64, right: IntervalDayTime, tz: Tz) -> Option<i64> { |
388 | 0 | Self::add_day_time(left, right, tz) |
389 | 0 | } |
390 | | |
391 | 0 | fn add_month_day_nano(left: i64, right: IntervalMonthDayNano, tz: Tz) -> Option<i64> { |
392 | 0 | Self::add_month_day_nano(left, right, tz) |
393 | 0 | } |
394 | | |
395 | 0 | fn sub_year_month(left: i64, right: i32, tz: Tz) -> Option<i64> { |
396 | 0 | Self::subtract_year_months(left, right, tz) |
397 | 0 | } |
398 | | |
399 | 0 | fn sub_day_time(left: i64, right: IntervalDayTime, tz: Tz) -> Option<i64> { |
400 | 0 | Self::subtract_day_time(left, right, tz) |
401 | 0 | } |
402 | | |
403 | 0 | fn sub_month_day_nano(left: i64, right: IntervalMonthDayNano, tz: Tz) -> Option<i64> { |
404 | 0 | Self::subtract_month_day_nano(left, right, tz) |
405 | 0 | } |
406 | | } |
407 | | }; |
408 | | } |
409 | | timestamp!(TimestampSecondType, DurationSecondType); |
410 | | timestamp!(TimestampMillisecondType, DurationMillisecondType); |
411 | | timestamp!(TimestampMicrosecondType, DurationMicrosecondType); |
412 | | timestamp!(TimestampNanosecondType, DurationNanosecondType); |
413 | | |
414 | | /// Perform arithmetic operation on a timestamp array |
415 | 0 | fn timestamp_op<T: TimestampOp>( |
416 | 0 | op: Op, |
417 | 0 | l: &dyn Array, |
418 | 0 | l_s: bool, |
419 | 0 | r: &dyn Array, |
420 | 0 | r_s: bool, |
421 | 0 | ) -> Result<ArrayRef, ArrowError> { |
422 | | use DataType::*; |
423 | | use IntervalUnit::*; |
424 | | |
425 | 0 | let l = l.as_primitive::<T>(); |
426 | 0 | let l_tz: Tz = l.timezone().unwrap_or("+00:00").parse()?; |
427 | | |
428 | 0 | let array: PrimitiveArray<T> = match (op, r.data_type()) { |
429 | 0 | (Op::Sub | Op::SubWrapping, Timestamp(unit, _)) if unit == &T::UNIT => { |
430 | 0 | let r = r.as_primitive::<T>(); |
431 | 0 | return Ok(try_op_ref!(T::Duration, l, l_s, r, r_s, l.sub_checked(r))); |
432 | | } |
433 | | |
434 | 0 | (Op::Add | Op::AddWrapping, Duration(unit)) if unit == &T::UNIT => { |
435 | 0 | let r = r.as_primitive::<T::Duration>(); |
436 | 0 | try_op!(l, l_s, r, r_s, l.add_checked(r)) |
437 | | } |
438 | 0 | (Op::Sub | Op::SubWrapping, Duration(unit)) if unit == &T::UNIT => { |
439 | 0 | let r = r.as_primitive::<T::Duration>(); |
440 | 0 | try_op!(l, l_s, r, r_s, l.sub_checked(r)) |
441 | | } |
442 | | |
443 | | (Op::Add | Op::AddWrapping, Interval(YearMonth)) => { |
444 | 0 | let r = r.as_primitive::<IntervalYearMonthType>(); |
445 | 0 | try_op!( |
446 | | l, |
447 | 0 | l_s, |
448 | | r, |
449 | 0 | r_s, |
450 | 0 | T::add_year_month(l, r, l_tz).ok_or(ArrowError::ComputeError( |
451 | 0 | "Timestamp out of range".to_string() |
452 | 0 | )) |
453 | | ) |
454 | | } |
455 | | (Op::Sub | Op::SubWrapping, Interval(YearMonth)) => { |
456 | 0 | let r = r.as_primitive::<IntervalYearMonthType>(); |
457 | 0 | try_op!( |
458 | | l, |
459 | 0 | l_s, |
460 | | r, |
461 | 0 | r_s, |
462 | 0 | T::sub_year_month(l, r, l_tz).ok_or(ArrowError::ComputeError( |
463 | 0 | "Timestamp out of range".to_string() |
464 | 0 | )) |
465 | | ) |
466 | | } |
467 | | |
468 | | (Op::Add | Op::AddWrapping, Interval(DayTime)) => { |
469 | 0 | let r = r.as_primitive::<IntervalDayTimeType>(); |
470 | 0 | try_op!( |
471 | | l, |
472 | 0 | l_s, |
473 | | r, |
474 | 0 | r_s, |
475 | 0 | T::add_day_time(l, r, l_tz).ok_or(ArrowError::ComputeError( |
476 | 0 | "Timestamp out of range".to_string() |
477 | 0 | )) |
478 | | ) |
479 | | } |
480 | | (Op::Sub | Op::SubWrapping, Interval(DayTime)) => { |
481 | 0 | let r = r.as_primitive::<IntervalDayTimeType>(); |
482 | 0 | try_op!( |
483 | | l, |
484 | 0 | l_s, |
485 | | r, |
486 | 0 | r_s, |
487 | 0 | T::sub_day_time(l, r, l_tz).ok_or(ArrowError::ComputeError( |
488 | 0 | "Timestamp out of range".to_string() |
489 | 0 | )) |
490 | | ) |
491 | | } |
492 | | |
493 | | (Op::Add | Op::AddWrapping, Interval(MonthDayNano)) => { |
494 | 0 | let r = r.as_primitive::<IntervalMonthDayNanoType>(); |
495 | 0 | try_op!( |
496 | | l, |
497 | 0 | l_s, |
498 | | r, |
499 | 0 | r_s, |
500 | 0 | T::add_month_day_nano(l, r, l_tz).ok_or(ArrowError::ComputeError( |
501 | 0 | "Timestamp out of range".to_string() |
502 | 0 | )) |
503 | | ) |
504 | | } |
505 | | (Op::Sub | Op::SubWrapping, Interval(MonthDayNano)) => { |
506 | 0 | let r = r.as_primitive::<IntervalMonthDayNanoType>(); |
507 | 0 | try_op!( |
508 | | l, |
509 | 0 | l_s, |
510 | | r, |
511 | 0 | r_s, |
512 | 0 | T::sub_month_day_nano(l, r, l_tz).ok_or(ArrowError::ComputeError( |
513 | 0 | "Timestamp out of range".to_string() |
514 | 0 | )) |
515 | | ) |
516 | | } |
517 | | _ => { |
518 | 0 | return Err(ArrowError::InvalidArgumentError(format!( |
519 | 0 | "Invalid timestamp arithmetic operation: {} {op} {}", |
520 | 0 | l.data_type(), |
521 | 0 | r.data_type() |
522 | 0 | ))) |
523 | | } |
524 | | }; |
525 | 0 | Ok(Arc::new(array.with_timezone_opt(l.timezone()))) |
526 | 0 | } |
527 | | |
528 | | /// Arithmetic trait for date arrays |
529 | | trait 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 | | |
551 | | impl DateOp for Date32Type { |
552 | 0 | 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 | 0 | Ok(Self::add_year_months(left, right)) |
555 | 0 | } |
556 | | |
557 | 0 | fn add_day_time( |
558 | 0 | left: Self::Native, |
559 | 0 | right: IntervalDayTime, |
560 | 0 | ) -> Result<Self::Native, ArrowError> { |
561 | 0 | Ok(Self::add_day_time(left, right)) |
562 | 0 | } |
563 | | |
564 | 0 | fn add_month_day_nano( |
565 | 0 | left: Self::Native, |
566 | 0 | right: IntervalMonthDayNano, |
567 | 0 | ) -> Result<Self::Native, ArrowError> { |
568 | 0 | Ok(Self::add_month_day_nano(left, right)) |
569 | 0 | } |
570 | | |
571 | 0 | fn sub_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> { |
572 | 0 | Ok(Self::subtract_year_months(left, right)) |
573 | 0 | } |
574 | | |
575 | 0 | fn sub_day_time( |
576 | 0 | left: Self::Native, |
577 | 0 | right: IntervalDayTime, |
578 | 0 | ) -> Result<Self::Native, ArrowError> { |
579 | 0 | Ok(Self::subtract_day_time(left, right)) |
580 | 0 | } |
581 | | |
582 | 0 | fn sub_month_day_nano( |
583 | 0 | left: Self::Native, |
584 | 0 | right: IntervalMonthDayNano, |
585 | 0 | ) -> Result<Self::Native, ArrowError> { |
586 | 0 | Ok(Self::subtract_month_day_nano(left, right)) |
587 | 0 | } |
588 | | } |
589 | | |
590 | | impl DateOp for Date64Type { |
591 | 0 | fn add_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> { |
592 | 0 | Self::add_year_months_opt(left, right).ok_or_else(|| { |
593 | 0 | ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right} months",)) |
594 | 0 | }) |
595 | 0 | } |
596 | | |
597 | 0 | fn add_day_time( |
598 | 0 | left: Self::Native, |
599 | 0 | right: IntervalDayTime, |
600 | 0 | ) -> Result<Self::Native, ArrowError> { |
601 | 0 | Self::add_day_time_opt(left, right).ok_or_else(|| { |
602 | 0 | ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right:?}")) |
603 | 0 | }) |
604 | 0 | } |
605 | | |
606 | 0 | fn add_month_day_nano( |
607 | 0 | left: Self::Native, |
608 | 0 | right: IntervalMonthDayNano, |
609 | 0 | ) -> Result<Self::Native, ArrowError> { |
610 | 0 | Self::add_month_day_nano_opt(left, right).ok_or_else(|| { |
611 | 0 | ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right:?}")) |
612 | 0 | }) |
613 | 0 | } |
614 | | |
615 | 0 | fn sub_year_month(left: Self::Native, right: i32) -> Result<Self::Native, ArrowError> { |
616 | 0 | Self::subtract_year_months_opt(left, right).ok_or_else(|| { |
617 | 0 | ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right} months",)) |
618 | 0 | }) |
619 | 0 | } |
620 | | |
621 | 0 | fn sub_day_time( |
622 | 0 | left: Self::Native, |
623 | 0 | right: IntervalDayTime, |
624 | 0 | ) -> Result<Self::Native, ArrowError> { |
625 | 0 | Self::subtract_day_time_opt(left, right).ok_or_else(|| { |
626 | 0 | ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right:?}")) |
627 | 0 | }) |
628 | 0 | } |
629 | | |
630 | 0 | fn sub_month_day_nano( |
631 | 0 | left: Self::Native, |
632 | 0 | right: IntervalMonthDayNano, |
633 | 0 | ) -> Result<Self::Native, ArrowError> { |
634 | 0 | Self::subtract_month_day_nano_opt(left, right).ok_or_else(|| { |
635 | 0 | ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right:?}")) |
636 | 0 | }) |
637 | 0 | } |
638 | | } |
639 | | |
640 | | /// Arithmetic trait for interval arrays |
641 | | trait 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 | | |
646 | | impl IntervalOp for IntervalYearMonthType { |
647 | 0 | fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> { |
648 | 0 | left.add_checked(right) |
649 | 0 | } |
650 | | |
651 | 0 | fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> { |
652 | 0 | left.sub_checked(right) |
653 | 0 | } |
654 | | } |
655 | | |
656 | | impl IntervalOp for IntervalDayTimeType { |
657 | 0 | fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> { |
658 | 0 | let (l_days, l_ms) = Self::to_parts(left); |
659 | 0 | let (r_days, r_ms) = Self::to_parts(right); |
660 | 0 | let days = l_days.add_checked(r_days)?; |
661 | 0 | let ms = l_ms.add_checked(r_ms)?; |
662 | 0 | Ok(Self::make_value(days, ms)) |
663 | 0 | } |
664 | | |
665 | 0 | fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> { |
666 | 0 | let (l_days, l_ms) = Self::to_parts(left); |
667 | 0 | let (r_days, r_ms) = Self::to_parts(right); |
668 | 0 | let days = l_days.sub_checked(r_days)?; |
669 | 0 | let ms = l_ms.sub_checked(r_ms)?; |
670 | 0 | Ok(Self::make_value(days, ms)) |
671 | 0 | } |
672 | | } |
673 | | |
674 | | impl IntervalOp for IntervalMonthDayNanoType { |
675 | 0 | fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> { |
676 | 0 | let (l_months, l_days, l_nanos) = Self::to_parts(left); |
677 | 0 | let (r_months, r_days, r_nanos) = Self::to_parts(right); |
678 | 0 | let months = l_months.add_checked(r_months)?; |
679 | 0 | let days = l_days.add_checked(r_days)?; |
680 | 0 | let nanos = l_nanos.add_checked(r_nanos)?; |
681 | 0 | Ok(Self::make_value(months, days, nanos)) |
682 | 0 | } |
683 | | |
684 | 0 | fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native, ArrowError> { |
685 | 0 | let (l_months, l_days, l_nanos) = Self::to_parts(left); |
686 | 0 | let (r_months, r_days, r_nanos) = Self::to_parts(right); |
687 | 0 | let months = l_months.sub_checked(r_months)?; |
688 | 0 | let days = l_days.sub_checked(r_days)?; |
689 | 0 | let nanos = l_nanos.sub_checked(r_nanos)?; |
690 | 0 | Ok(Self::make_value(months, days, nanos)) |
691 | 0 | } |
692 | | } |
693 | | |
694 | | /// Perform arithmetic operation on an interval array |
695 | 0 | fn interval_op<T: IntervalOp>( |
696 | 0 | op: Op, |
697 | 0 | l: &dyn Array, |
698 | 0 | l_s: bool, |
699 | 0 | r: &dyn Array, |
700 | 0 | r_s: bool, |
701 | 0 | ) -> Result<ArrayRef, ArrowError> { |
702 | 0 | let l = l.as_primitive::<T>(); |
703 | 0 | let r = r.as_primitive::<T>(); |
704 | 0 | match op { |
705 | 0 | Op::Add | Op::AddWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, T::add(l, r))), |
706 | 0 | Op::Sub | Op::SubWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub(l, r))), |
707 | 0 | _ => Err(ArrowError::InvalidArgumentError(format!( |
708 | 0 | "Invalid interval arithmetic operation: {} {op} {}", |
709 | 0 | l.data_type(), |
710 | 0 | r.data_type() |
711 | 0 | ))), |
712 | | } |
713 | 0 | } |
714 | | |
715 | 0 | fn duration_op<T: ArrowPrimitiveType>( |
716 | 0 | op: Op, |
717 | 0 | l: &dyn Array, |
718 | 0 | l_s: bool, |
719 | 0 | r: &dyn Array, |
720 | 0 | r_s: bool, |
721 | 0 | ) -> Result<ArrayRef, ArrowError> { |
722 | 0 | let l = l.as_primitive::<T>(); |
723 | 0 | let r = r.as_primitive::<T>(); |
724 | 0 | match op { |
725 | 0 | Op::Add | Op::AddWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, l.add_checked(r))), |
726 | 0 | Op::Sub | Op::SubWrapping => Ok(try_op_ref!(T, l, l_s, r, r_s, l.sub_checked(r))), |
727 | 0 | _ => Err(ArrowError::InvalidArgumentError(format!( |
728 | 0 | "Invalid duration arithmetic operation: {} {op} {}", |
729 | 0 | l.data_type(), |
730 | 0 | r.data_type() |
731 | 0 | ))), |
732 | | } |
733 | 0 | } |
734 | | |
735 | | /// Perform arithmetic operation on a date array |
736 | 0 | fn date_op<T: DateOp>( |
737 | 0 | op: Op, |
738 | 0 | l: &dyn Array, |
739 | 0 | l_s: bool, |
740 | 0 | r: &dyn Array, |
741 | 0 | r_s: bool, |
742 | 0 | ) -> Result<ArrayRef, ArrowError> { |
743 | | use DataType::*; |
744 | | use IntervalUnit::*; |
745 | | |
746 | | const NUM_SECONDS_IN_DAY: i64 = 60 * 60 * 24; |
747 | | |
748 | 0 | let r_t = r.data_type(); |
749 | 0 | match (T::DATA_TYPE, op, r_t) { |
750 | | (Date32, Op::Sub | Op::SubWrapping, Date32) => { |
751 | 0 | let l = l.as_primitive::<Date32Type>(); |
752 | 0 | let r = r.as_primitive::<Date32Type>(); |
753 | 0 | return Ok(op_ref!( |
754 | | DurationSecondType, |
755 | | l, |
756 | 0 | l_s, |
757 | | r, |
758 | 0 | r_s, |
759 | 0 | ((l as i64) - (r as i64)) * NUM_SECONDS_IN_DAY |
760 | | )); |
761 | | } |
762 | | (Date64, Op::Sub | Op::SubWrapping, Date64) => { |
763 | 0 | let l = l.as_primitive::<Date64Type>(); |
764 | 0 | let r = r.as_primitive::<Date64Type>(); |
765 | 0 | let result = try_op_ref!(DurationMillisecondType, l, l_s, r, r_s, l.sub_checked(r)); |
766 | 0 | return Ok(result); |
767 | | } |
768 | 0 | _ => {} |
769 | | } |
770 | | |
771 | 0 | let l = l.as_primitive::<T>(); |
772 | 0 | match (op, r_t) { |
773 | | (Op::Add | Op::AddWrapping, Interval(YearMonth)) => { |
774 | 0 | let r = r.as_primitive::<IntervalYearMonthType>(); |
775 | 0 | 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 | 0 | let r = r.as_primitive::<IntervalYearMonthType>(); |
779 | 0 | 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 | 0 | let r = r.as_primitive::<IntervalDayTimeType>(); |
784 | 0 | 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 | 0 | let r = r.as_primitive::<IntervalDayTimeType>(); |
788 | 0 | 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 | 0 | let r = r.as_primitive::<IntervalMonthDayNanoType>(); |
793 | 0 | 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 | 0 | let r = r.as_primitive::<IntervalMonthDayNanoType>(); |
797 | 0 | Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_month_day_nano(l, r))) |
798 | | } |
799 | | |
800 | 0 | _ => Err(ArrowError::InvalidArgumentError(format!( |
801 | 0 | "Invalid date arithmetic operation: {} {op} {}", |
802 | 0 | l.data_type(), |
803 | 0 | r.data_type() |
804 | 0 | ))), |
805 | | } |
806 | 0 | } |
807 | | |
808 | | /// Perform arithmetic operation on decimal arrays |
809 | 0 | fn decimal_op<T: DecimalType>( |
810 | 0 | op: Op, |
811 | 0 | l: &dyn Array, |
812 | 0 | l_s: bool, |
813 | 0 | r: &dyn Array, |
814 | 0 | r_s: bool, |
815 | 0 | ) -> Result<ArrayRef, ArrowError> { |
816 | 0 | let l = l.as_primitive::<T>(); |
817 | 0 | let r = r.as_primitive::<T>(); |
818 | | |
819 | 0 | let (p1, s1, p2, s2) = match (l.data_type(), r.data_type()) { |
820 | 0 | (DataType::Decimal32(p1, s1), DataType::Decimal32(p2, s2)) => (p1, s1, p2, s2), |
821 | 0 | (DataType::Decimal64(p1, s1), DataType::Decimal64(p2, s2)) => (p1, s1, p2, s2), |
822 | 0 | (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => (p1, s1, p2, s2), |
823 | 0 | (DataType::Decimal256(p1, s1), DataType::Decimal256(p2, s2)) => (p1, s1, p2, s2), |
824 | 0 | _ => 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 | 0 | let array: PrimitiveArray<T> = match op { |
830 | | Op::Add | Op::AddWrapping | Op::Sub | Op::SubWrapping => { |
831 | | // max(s1, s2) |
832 | 0 | let result_scale = *s1.max(s2); |
833 | | |
834 | | // max(s1, s2) + max(p1-s1, p2-s2) + 1 |
835 | 0 | let result_precision = |
836 | 0 | (result_scale.saturating_add((*p1 as i8 - s1).max(*p2 as i8 - s2)) as u8) |
837 | 0 | .saturating_add(1) |
838 | 0 | .min(T::MAX_PRECISION); |
839 | | |
840 | 0 | let l_mul = T::Native::usize_as(10).pow_checked((result_scale - s1) as _)?; |
841 | 0 | let r_mul = T::Native::usize_as(10).pow_checked((result_scale - s2) as _)?; |
842 | | |
843 | 0 | match op { |
844 | | Op::Add | Op::AddWrapping => { |
845 | 0 | try_op!( |
846 | | l, |
847 | 0 | l_s, |
848 | | r, |
849 | 0 | r_s, |
850 | 0 | l.mul_checked(l_mul)?.add_checked(r.mul_checked(r_mul)?) |
851 | | ) |
852 | | } |
853 | | Op::Sub | Op::SubWrapping => { |
854 | 0 | try_op!( |
855 | | l, |
856 | 0 | l_s, |
857 | | r, |
858 | 0 | r_s, |
859 | 0 | l.mul_checked(l_mul)?.sub_checked(r.mul_checked(r_mul)?) |
860 | | ) |
861 | | } |
862 | 0 | _ => unreachable!(), |
863 | | } |
864 | 0 | .with_precision_and_scale(result_precision, result_scale)? |
865 | | } |
866 | | Op::Mul | Op::MulWrapping => { |
867 | 0 | let result_precision = p1.saturating_add(p2 + 1).min(T::MAX_PRECISION); |
868 | 0 | let result_scale = s1.saturating_add(*s2); |
869 | 0 | 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 | 0 | return Err(ArrowError::InvalidArgumentError(format!( |
873 | 0 | "Output scale of {} {op} {} would exceed max scale of {}", |
874 | 0 | l.data_type(), |
875 | 0 | r.data_type(), |
876 | 0 | T::MAX_SCALE |
877 | 0 | ))); |
878 | 0 | } |
879 | | |
880 | 0 | try_op!(l, l_s, r, r_s, l.mul_checked(r)) |
881 | 0 | .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 | 0 | let result_scale = s1.saturating_add(4).min(T::MAX_SCALE); |
888 | 0 | let mul_pow = result_scale - s1 + s2; |
889 | | |
890 | | // p1 - s1 + s2 + result_scale |
891 | 0 | let result_precision = (mul_pow.saturating_add(*p1 as i8) as u8).min(T::MAX_PRECISION); |
892 | | |
893 | 0 | let (l_mul, r_mul) = match mul_pow.cmp(&0) { |
894 | | Ordering::Greater => ( |
895 | 0 | T::Native::usize_as(10).pow_checked(mul_pow as _)?, |
896 | | T::Native::ONE, |
897 | | ), |
898 | 0 | Ordering::Equal => (T::Native::ONE, T::Native::ONE), |
899 | | Ordering::Less => ( |
900 | | T::Native::ONE, |
901 | 0 | T::Native::usize_as(10).pow_checked(mul_pow.neg_wrapping() as _)?, |
902 | | ), |
903 | | }; |
904 | | |
905 | 0 | try_op!( |
906 | | l, |
907 | 0 | l_s, |
908 | | r, |
909 | 0 | r_s, |
910 | 0 | l.mul_checked(l_mul)?.div_checked(r.mul_checked(r_mul)?) |
911 | | ) |
912 | 0 | .with_precision_and_scale(result_precision, result_scale)? |
913 | | } |
914 | | |
915 | | Op::Rem => { |
916 | | // max(s1, s2) |
917 | 0 | let result_scale = *s1.max(s2); |
918 | | // min(p1-s1, p2 -s2) + max( s1,s2 ) |
919 | 0 | let result_precision = |
920 | 0 | (result_scale.saturating_add((*p1 as i8 - s1).min(*p2 as i8 - s2)) as u8) |
921 | 0 | .min(T::MAX_PRECISION); |
922 | | |
923 | 0 | let l_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s1) as _); |
924 | 0 | let r_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s2) as _); |
925 | | |
926 | 0 | try_op!( |
927 | | l, |
928 | 0 | l_s, |
929 | | r, |
930 | 0 | r_s, |
931 | 0 | l.mul_checked(l_mul)?.mod_checked(r.mul_checked(r_mul)?) |
932 | | ) |
933 | 0 | .with_precision_and_scale(result_precision, result_scale)? |
934 | | } |
935 | | }; |
936 | | |
937 | 0 | Ok(Arc::new(array)) |
938 | 0 | } |
939 | | |
940 | | #[cfg(test)] |
941 | | mod tests { |
942 | | use super::*; |
943 | | use arrow_array::temporal_conversions::{as_date, as_datetime}; |
944 | | use arrow_buffer::{i256, ScalarBuffer}; |
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!(err, "Invalid argument error: Output scale of Decimal128(3, 3) * Decimal128(37, 37) would exceed max scale of 38"); |
1267 | | |
1268 | | let a = Decimal128Array::from(vec![1]) |
1269 | | .with_precision_and_scale(3, -2) |
1270 | | .unwrap(); |
1271 | | let err = add(&a, &b).unwrap_err().to_string(); |
1272 | | assert_eq!(err, "Arithmetic overflow: Overflow happened on: 10 ^ 39"); |
1273 | | |
1274 | | let a = Decimal128Array::from(vec![10]) |
1275 | | .with_precision_and_scale(3, -1) |
1276 | | .unwrap(); |
1277 | | let err = add(&a, &b).unwrap_err().to_string(); |
1278 | | assert_eq!( |
1279 | | err, |
1280 | | "Arithmetic overflow: Overflow happened on: 10 * 100000000000000000000000000000000000000" |
1281 | | ); |
1282 | | |
1283 | | let b = Decimal128Array::from(vec![0]) |
1284 | | .with_precision_and_scale(1, 1) |
1285 | | .unwrap(); |
1286 | | let err = div(&a, &b).unwrap_err().to_string(); |
1287 | | assert_eq!(err, "Divide by zero error"); |
1288 | | let err = rem(&a, &b).unwrap_err().to_string(); |
1289 | | assert_eq!(err, "Divide by zero error"); |
1290 | | } |
1291 | | |
1292 | | fn test_timestamp_impl<T: TimestampOp>() { |
1293 | | let a = PrimitiveArray::<T>::new(vec![2000000, 434030324, 53943340].into(), None); |
1294 | | let b = PrimitiveArray::<T>::new(vec![329593, 59349, 694994].into(), None); |
1295 | | |
1296 | | let result = sub(&a, &b).unwrap(); |
1297 | | assert_eq!( |
1298 | | result.as_primitive::<T::Duration>().values(), |
1299 | | &[1670407, 433970975, 53248346] |
1300 | | ); |
1301 | | |
1302 | | let r2 = add(&b, &result.as_ref()).unwrap(); |
1303 | | assert_eq!(r2.as_ref(), &a); |
1304 | | |
1305 | | let r3 = add(&result.as_ref(), &b).unwrap(); |
1306 | | assert_eq!(r3.as_ref(), &a); |
1307 | | |
1308 | | let format_array = |x: &dyn Array| -> Vec<String> { |
1309 | | x.as_primitive::<T>() |
1310 | | .values() |
1311 | | .into_iter() |
1312 | | .map(|x| as_datetime::<T>(*x).unwrap().to_string()) |
1313 | | .collect() |
1314 | | }; |
1315 | | |
1316 | | let values = vec![ |
1317 | | "1970-01-01T00:00:00Z", |
1318 | | "2010-04-01T04:00:20Z", |
1319 | | "1960-01-30T04:23:20Z", |
1320 | | ] |
1321 | | .into_iter() |
1322 | | .map(|x| T::make_value(DateTime::parse_from_rfc3339(x).unwrap().naive_utc()).unwrap()) |
1323 | | .collect(); |
1324 | | |
1325 | | let a = PrimitiveArray::<T>::new(values, None); |
1326 | | let b = IntervalYearMonthArray::from(vec![ |
1327 | | IntervalYearMonthType::make_value(5, 34), |
1328 | | IntervalYearMonthType::make_value(-2, 4), |
1329 | | IntervalYearMonthType::make_value(7, -4), |
1330 | | ]); |
1331 | | let r4 = add(&a, &b).unwrap(); |
1332 | | assert_eq!( |
1333 | | &format_array(r4.as_ref()), |
1334 | | &[ |
1335 | | "1977-11-01 00:00:00".to_string(), |
1336 | | "2008-08-01 04:00:20".to_string(), |
1337 | | "1966-09-30 04:23:20".to_string() |
1338 | | ] |
1339 | | ); |
1340 | | |
1341 | | let r5 = sub(&r4, &b).unwrap(); |
1342 | | assert_eq!(r5.as_ref(), &a); |
1343 | | |
1344 | | let b = IntervalDayTimeArray::from(vec![ |
1345 | | IntervalDayTimeType::make_value(5, 454000), |
1346 | | IntervalDayTimeType::make_value(-34, 0), |
1347 | | IntervalDayTimeType::make_value(7, -4000), |
1348 | | ]); |
1349 | | let r6 = add(&a, &b).unwrap(); |
1350 | | assert_eq!( |
1351 | | &format_array(r6.as_ref()), |
1352 | | &[ |
1353 | | "1970-01-06 00:07:34".to_string(), |
1354 | | "2010-02-26 04:00:20".to_string(), |
1355 | | "1960-02-06 04:23:16".to_string() |
1356 | | ] |
1357 | | ); |
1358 | | |
1359 | | let r7 = sub(&r6, &b).unwrap(); |
1360 | | assert_eq!(r7.as_ref(), &a); |
1361 | | |
1362 | | let b = IntervalMonthDayNanoArray::from(vec![ |
1363 | | IntervalMonthDayNanoType::make_value(344, 34, -43_000_000_000), |
1364 | | IntervalMonthDayNanoType::make_value(-593, -33, 13_000_000_000), |
1365 | | IntervalMonthDayNanoType::make_value(5, 2, 493_000_000_000), |
1366 | | ]); |
1367 | | let r8 = add(&a, &b).unwrap(); |
1368 | | assert_eq!( |
1369 | | &format_array(r8.as_ref()), |
1370 | | &[ |
1371 | | "1998-10-04 23:59:17".to_string(), |
1372 | | "1960-09-29 04:00:33".to_string(), |
1373 | | "1960-07-02 04:31:33".to_string() |
1374 | | ] |
1375 | | ); |
1376 | | |
1377 | | let r9 = sub(&r8, &b).unwrap(); |
1378 | | // Note: subtraction is not the inverse of addition for intervals |
1379 | | assert_eq!( |
1380 | | &format_array(r9.as_ref()), |
1381 | | &[ |
1382 | | "1970-01-02 00:00:00".to_string(), |
1383 | | "2010-04-02 04:00:20".to_string(), |
1384 | | "1960-01-31 04:23:20".to_string() |
1385 | | ] |
1386 | | ); |
1387 | | } |
1388 | | |
1389 | | #[test] |
1390 | | fn test_timestamp() { |
1391 | | test_timestamp_impl::<TimestampSecondType>(); |
1392 | | test_timestamp_impl::<TimestampMillisecondType>(); |
1393 | | test_timestamp_impl::<TimestampMicrosecondType>(); |
1394 | | test_timestamp_impl::<TimestampNanosecondType>(); |
1395 | | } |
1396 | | |
1397 | | #[test] |
1398 | | fn test_interval() { |
1399 | | let a = IntervalYearMonthArray::from(vec![ |
1400 | | IntervalYearMonthType::make_value(32, 4), |
1401 | | IntervalYearMonthType::make_value(32, 4), |
1402 | | ]); |
1403 | | let b = IntervalYearMonthArray::from(vec![ |
1404 | | IntervalYearMonthType::make_value(-4, 6), |
1405 | | IntervalYearMonthType::make_value(-3, 23), |
1406 | | ]); |
1407 | | let result = add(&a, &b).unwrap(); |
1408 | | assert_eq!( |
1409 | | result.as_ref(), |
1410 | | &IntervalYearMonthArray::from(vec![ |
1411 | | IntervalYearMonthType::make_value(28, 10), |
1412 | | IntervalYearMonthType::make_value(29, 27) |
1413 | | ]) |
1414 | | ); |
1415 | | let result = sub(&a, &b).unwrap(); |
1416 | | assert_eq!( |
1417 | | result.as_ref(), |
1418 | | &IntervalYearMonthArray::from(vec![ |
1419 | | IntervalYearMonthType::make_value(36, -2), |
1420 | | IntervalYearMonthType::make_value(35, -19) |
1421 | | ]) |
1422 | | ); |
1423 | | |
1424 | | let a = IntervalDayTimeArray::from(vec![ |
1425 | | IntervalDayTimeType::make_value(32, 4), |
1426 | | IntervalDayTimeType::make_value(32, 4), |
1427 | | ]); |
1428 | | let b = IntervalDayTimeArray::from(vec![ |
1429 | | IntervalDayTimeType::make_value(-4, 6), |
1430 | | IntervalDayTimeType::make_value(-3, 23), |
1431 | | ]); |
1432 | | let result = add(&a, &b).unwrap(); |
1433 | | assert_eq!( |
1434 | | result.as_ref(), |
1435 | | &IntervalDayTimeArray::from(vec![ |
1436 | | IntervalDayTimeType::make_value(28, 10), |
1437 | | IntervalDayTimeType::make_value(29, 27) |
1438 | | ]) |
1439 | | ); |
1440 | | let result = sub(&a, &b).unwrap(); |
1441 | | assert_eq!( |
1442 | | result.as_ref(), |
1443 | | &IntervalDayTimeArray::from(vec![ |
1444 | | IntervalDayTimeType::make_value(36, -2), |
1445 | | IntervalDayTimeType::make_value(35, -19) |
1446 | | ]) |
1447 | | ); |
1448 | | let a = IntervalMonthDayNanoArray::from(vec![ |
1449 | | IntervalMonthDayNanoType::make_value(32, 4, 4000000000000), |
1450 | | IntervalMonthDayNanoType::make_value(32, 4, 45463000000000000), |
1451 | | ]); |
1452 | | let b = IntervalMonthDayNanoArray::from(vec![ |
1453 | | IntervalMonthDayNanoType::make_value(-4, 6, 46000000000000), |
1454 | | IntervalMonthDayNanoType::make_value(-3, 23, 3564000000000000), |
1455 | | ]); |
1456 | | let result = add(&a, &b).unwrap(); |
1457 | | assert_eq!( |
1458 | | result.as_ref(), |
1459 | | &IntervalMonthDayNanoArray::from(vec![ |
1460 | | IntervalMonthDayNanoType::make_value(28, 10, 50000000000000), |
1461 | | IntervalMonthDayNanoType::make_value(29, 27, 49027000000000000) |
1462 | | ]) |
1463 | | ); |
1464 | | let result = sub(&a, &b).unwrap(); |
1465 | | assert_eq!( |
1466 | | result.as_ref(), |
1467 | | &IntervalMonthDayNanoArray::from(vec![ |
1468 | | IntervalMonthDayNanoType::make_value(36, -2, -42000000000000), |
1469 | | IntervalMonthDayNanoType::make_value(35, -19, 41899000000000000) |
1470 | | ]) |
1471 | | ); |
1472 | | let a = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNano::MAX]); |
1473 | | let b = IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNano::ONE]); |
1474 | | let err = add(&a, &b).unwrap_err().to_string(); |
1475 | | assert_eq!( |
1476 | | err, |
1477 | | "Arithmetic overflow: Overflow happened on: 2147483647 + 1" |
1478 | | ); |
1479 | | } |
1480 | | |
1481 | | fn test_duration_impl<T: ArrowPrimitiveType<Native = i64>>() { |
1482 | | let a = PrimitiveArray::<T>::new(vec![1000, 4394, -3944].into(), None); |
1483 | | let b = PrimitiveArray::<T>::new(vec![4, -5, -243].into(), None); |
1484 | | |
1485 | | let result = add(&a, &b).unwrap(); |
1486 | | assert_eq!(result.as_primitive::<T>().values(), &[1004, 4389, -4187]); |
1487 | | let result = sub(&a, &b).unwrap(); |
1488 | | assert_eq!(result.as_primitive::<T>().values(), &[996, 4399, -3701]); |
1489 | | |
1490 | | let err = mul(&a, &b).unwrap_err().to_string(); |
1491 | | assert!( |
1492 | | err.contains("Invalid duration arithmetic operation"), |
1493 | | "{err}" |
1494 | | ); |
1495 | | |
1496 | | let err = div(&a, &b).unwrap_err().to_string(); |
1497 | | assert!( |
1498 | | err.contains("Invalid duration arithmetic operation"), |
1499 | | "{err}" |
1500 | | ); |
1501 | | |
1502 | | let err = rem(&a, &b).unwrap_err().to_string(); |
1503 | | assert!( |
1504 | | err.contains("Invalid duration arithmetic operation"), |
1505 | | "{err}" |
1506 | | ); |
1507 | | |
1508 | | let a = PrimitiveArray::<T>::new(vec![i64::MAX].into(), None); |
1509 | | let b = PrimitiveArray::<T>::new(vec![1].into(), None); |
1510 | | let err = add(&a, &b).unwrap_err().to_string(); |
1511 | | assert_eq!( |
1512 | | err, |
1513 | | "Arithmetic overflow: Overflow happened on: 9223372036854775807 + 1" |
1514 | | ); |
1515 | | } |
1516 | | |
1517 | | #[test] |
1518 | | fn test_duration() { |
1519 | | test_duration_impl::<DurationSecondType>(); |
1520 | | test_duration_impl::<DurationMillisecondType>(); |
1521 | | test_duration_impl::<DurationMicrosecondType>(); |
1522 | | test_duration_impl::<DurationNanosecondType>(); |
1523 | | } |
1524 | | |
1525 | | fn test_date_impl<T: ArrowPrimitiveType, F>(f: F) |
1526 | | where |
1527 | | F: Fn(NaiveDate) -> T::Native, |
1528 | | T::Native: TryInto<i64>, |
1529 | | { |
1530 | | let a = PrimitiveArray::<T>::new( |
1531 | | vec![ |
1532 | | f(NaiveDate::from_ymd_opt(1979, 1, 30).unwrap()), |
1533 | | f(NaiveDate::from_ymd_opt(2010, 4, 3).unwrap()), |
1534 | | f(NaiveDate::from_ymd_opt(2008, 2, 29).unwrap()), |
1535 | | ] |
1536 | | .into(), |
1537 | | None, |
1538 | | ); |
1539 | | |
1540 | | let b = IntervalYearMonthArray::from(vec![ |
1541 | | IntervalYearMonthType::make_value(34, 2), |
1542 | | IntervalYearMonthType::make_value(3, -3), |
1543 | | IntervalYearMonthType::make_value(-12, 4), |
1544 | | ]); |
1545 | | |
1546 | | let format_array = |x: &dyn Array| -> Vec<String> { |
1547 | | x.as_primitive::<T>() |
1548 | | .values() |
1549 | | .into_iter() |
1550 | | .map(|x| { |
1551 | | as_date::<T>((*x).try_into().ok().unwrap()) |
1552 | | .unwrap() |
1553 | | .to_string() |
1554 | | }) |
1555 | | .collect() |
1556 | | }; |
1557 | | |
1558 | | let result = add(&a, &b).unwrap(); |
1559 | | assert_eq!( |
1560 | | &format_array(result.as_ref()), |
1561 | | &[ |
1562 | | "2013-03-30".to_string(), |
1563 | | "2013-01-03".to_string(), |
1564 | | "1996-06-29".to_string(), |
1565 | | ] |
1566 | | ); |
1567 | | let result = sub(&result, &b).unwrap(); |
1568 | | assert_eq!(result.as_ref(), &a); |
1569 | | |
1570 | | let b = IntervalDayTimeArray::from(vec![ |
1571 | | IntervalDayTimeType::make_value(34, 2), |
1572 | | IntervalDayTimeType::make_value(3, -3), |
1573 | | IntervalDayTimeType::make_value(-12, 4), |
1574 | | ]); |
1575 | | |
1576 | | let result = add(&a, &b).unwrap(); |
1577 | | assert_eq!( |
1578 | | &format_array(result.as_ref()), |
1579 | | &[ |
1580 | | "1979-03-05".to_string(), |
1581 | | "2010-04-06".to_string(), |
1582 | | "2008-02-17".to_string(), |
1583 | | ] |
1584 | | ); |
1585 | | let result = sub(&result, &b).unwrap(); |
1586 | | assert_eq!(result.as_ref(), &a); |
1587 | | |
1588 | | let b = IntervalMonthDayNanoArray::from(vec![ |
1589 | | IntervalMonthDayNanoType::make_value(34, 2, -34353534), |
1590 | | IntervalMonthDayNanoType::make_value(3, -3, 2443), |
1591 | | IntervalMonthDayNanoType::make_value(-12, 4, 2323242423232), |
1592 | | ]); |
1593 | | |
1594 | | let result = add(&a, &b).unwrap(); |
1595 | | assert_eq!( |
1596 | | &format_array(result.as_ref()), |
1597 | | &[ |
1598 | | "1981-12-02".to_string(), |
1599 | | "2010-06-30".to_string(), |
1600 | | "2007-03-04".to_string(), |
1601 | | ] |
1602 | | ); |
1603 | | let result = sub(&result, &b).unwrap(); |
1604 | | assert_eq!( |
1605 | | &format_array(result.as_ref()), |
1606 | | &[ |
1607 | | "1979-01-31".to_string(), |
1608 | | "2010-04-02".to_string(), |
1609 | | "2008-02-29".to_string(), |
1610 | | ] |
1611 | | ); |
1612 | | } |
1613 | | |
1614 | | #[test] |
1615 | | fn test_date() { |
1616 | | test_date_impl::<Date32Type, _>(Date32Type::from_naive_date); |
1617 | | test_date_impl::<Date64Type, _>(Date64Type::from_naive_date); |
1618 | | |
1619 | | let a = Date32Array::from(vec![i32::MIN, i32::MAX, 23, 7684]); |
1620 | | let b = Date32Array::from(vec![i32::MIN, i32::MIN, -2, 45]); |
1621 | | let result = sub(&a, &b).unwrap(); |
1622 | | assert_eq!( |
1623 | | result.as_primitive::<DurationSecondType>().values(), |
1624 | | &[0, 371085174288000, 2160000, 660009600] |
1625 | | ); |
1626 | | |
1627 | | let a = Date64Array::from(vec![4343, 76676, 3434]); |
1628 | | let b = Date64Array::from(vec![3, -5, 5]); |
1629 | | let result = sub(&a, &b).unwrap(); |
1630 | | assert_eq!( |
1631 | | result.as_primitive::<DurationMillisecondType>().values(), |
1632 | | &[4340, 76681, 3429] |
1633 | | ); |
1634 | | |
1635 | | let a = Date64Array::from(vec![i64::MAX]); |
1636 | | let b = Date64Array::from(vec![-1]); |
1637 | | let err = sub(&a, &b).unwrap_err().to_string(); |
1638 | | assert_eq!( |
1639 | | err, |
1640 | | "Arithmetic overflow: Overflow happened on: 9223372036854775807 - -1" |
1641 | | ); |
1642 | | } |
1643 | | |
1644 | | #[test] |
1645 | | fn test_date64_to_naive_date_opt_boundary_values() { |
1646 | | use arrow_array::types::Date64Type; |
1647 | | |
1648 | | // Date64Type::to_naive_date_opt has boundaries determined by NaiveDate's supported range. |
1649 | | // The valid date range is from January 1, -262143 to December 31, 262142 (Gregorian calendar). |
1650 | | |
1651 | | let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); |
1652 | | let ms_per_day = 24 * 60 * 60 * 1000i64; |
1653 | | |
1654 | | // Define the boundary dates using NaiveDate::from_ymd_opt |
1655 | | let max_valid_date = NaiveDate::from_ymd_opt(262142, 12, 31).unwrap(); |
1656 | | let min_valid_date = NaiveDate::from_ymd_opt(-262143, 1, 1).unwrap(); |
1657 | | |
1658 | | // Calculate their millisecond values from epoch |
1659 | | let max_valid_millis = (max_valid_date - epoch).num_milliseconds(); |
1660 | | let min_valid_millis = (min_valid_date - epoch).num_milliseconds(); |
1661 | | |
1662 | | // Verify these match the expected boundaries in milliseconds |
1663 | | assert_eq!( |
1664 | | max_valid_millis, 8210266790400000i64, |
1665 | | "December 31, 262142 should be 8210266790400000 ms from epoch" |
1666 | | ); |
1667 | | assert_eq!( |
1668 | | min_valid_millis, -8334601228800000i64, |
1669 | | "January 1, -262143 should be -8334601228800000 ms from epoch" |
1670 | | ); |
1671 | | |
1672 | | // Test that the boundary dates work |
1673 | | assert!( |
1674 | | Date64Type::to_naive_date_opt(max_valid_millis).is_some(), |
1675 | | "December 31, 262142 should return Some" |
1676 | | ); |
1677 | | assert!( |
1678 | | Date64Type::to_naive_date_opt(min_valid_millis).is_some(), |
1679 | | "January 1, -262143 should return Some" |
1680 | | ); |
1681 | | |
1682 | | // Test that one day beyond the boundaries fails |
1683 | | assert!( |
1684 | | Date64Type::to_naive_date_opt(max_valid_millis + ms_per_day).is_none(), |
1685 | | "January 1, 262143 should return None" |
1686 | | ); |
1687 | | assert!( |
1688 | | Date64Type::to_naive_date_opt(min_valid_millis - ms_per_day).is_none(), |
1689 | | "December 31, -262144 should return None" |
1690 | | ); |
1691 | | |
1692 | | // Test some values well within the valid range |
1693 | | assert!( |
1694 | | Date64Type::to_naive_date_opt(0).is_some(), |
1695 | | "Epoch (1970-01-01) should return Some" |
1696 | | ); |
1697 | | let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); |
1698 | | let year_2000_millis = (year_2000 - epoch).num_milliseconds(); |
1699 | | assert!( |
1700 | | Date64Type::to_naive_date_opt(year_2000_millis).is_some(), |
1701 | | "Year 2000 should return Some" |
1702 | | ); |
1703 | | |
1704 | | // Test extreme values that definitely fail due to Duration constraints |
1705 | | assert!( |
1706 | | Date64Type::to_naive_date_opt(i64::MAX).is_none(), |
1707 | | "i64::MAX should return None" |
1708 | | ); |
1709 | | assert!( |
1710 | | Date64Type::to_naive_date_opt(i64::MIN).is_none(), |
1711 | | "i64::MIN should return None" |
1712 | | ); |
1713 | | } |
1714 | | |
1715 | | #[test] |
1716 | | fn test_date64_add_year_months_opt_boundary_values() { |
1717 | | use arrow_array::types::Date64Type; |
1718 | | |
1719 | | let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); |
1720 | | |
1721 | | // Test normal case within valid range |
1722 | | let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); |
1723 | | let year_2000_millis = (year_2000 - epoch).num_milliseconds(); |
1724 | | assert!( |
1725 | | Date64Type::add_year_months_opt(year_2000_millis, 120).is_some(), |
1726 | | "Adding 10 years to year 2000 should succeed" |
1727 | | ); |
1728 | | |
1729 | | // Test with moderate years that are within chrono's safe range |
1730 | | let large_year = NaiveDate::from_ymd_opt(5000, 1, 1).unwrap(); |
1731 | | let large_year_millis = (large_year - epoch).num_milliseconds(); |
1732 | | assert!( |
1733 | | Date64Type::add_year_months_opt(large_year_millis, 12).is_some(), |
1734 | | "Adding 12 months to year 5000 should succeed" |
1735 | | ); |
1736 | | |
1737 | | let neg_year = NaiveDate::from_ymd_opt(-5000, 12, 31).unwrap(); |
1738 | | let neg_year_millis = (neg_year - epoch).num_milliseconds(); |
1739 | | assert!( |
1740 | | Date64Type::add_year_months_opt(neg_year_millis, -12).is_some(), |
1741 | | "Subtracting 12 months from year -5000 should succeed" |
1742 | | ); |
1743 | | |
1744 | | // Test with extreme input values that would cause overflow |
1745 | | assert!( |
1746 | | Date64Type::add_year_months_opt(i64::MAX, 1).is_none(), |
1747 | | "Adding months to i64::MAX should fail" |
1748 | | ); |
1749 | | assert!( |
1750 | | Date64Type::add_year_months_opt(i64::MIN, -1).is_none(), |
1751 | | "Subtracting months from i64::MIN should fail" |
1752 | | ); |
1753 | | |
1754 | | // Test edge case: adding zero should always work for valid dates |
1755 | | assert!( |
1756 | | Date64Type::add_year_months_opt(year_2000_millis, 0).is_some(), |
1757 | | "Adding zero months should always succeed for valid dates" |
1758 | | ); |
1759 | | } |
1760 | | |
1761 | | #[test] |
1762 | | fn test_date64_add_day_time_opt_boundary_values() { |
1763 | | use arrow_array::types::Date64Type; |
1764 | | use arrow_buffer::IntervalDayTime; |
1765 | | |
1766 | | let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); |
1767 | | |
1768 | | // Test with a date far from the boundary but still testing the function |
1769 | | let near_max_date = NaiveDate::from_ymd_opt(200000, 12, 1).unwrap(); |
1770 | | let near_max_millis = (near_max_date - epoch).num_milliseconds(); |
1771 | | |
1772 | | // Adding 30 days should succeed |
1773 | | let interval_30_days = IntervalDayTime::new(30, 0); |
1774 | | assert!( |
1775 | | Date64Type::add_day_time_opt(near_max_millis, interval_30_days).is_some(), |
1776 | | "Adding 30 days to large year should succeed" |
1777 | | ); |
1778 | | |
1779 | | // Adding a very large number of days should fail |
1780 | | let interval_large_days = IntervalDayTime::new(100000000, 0); |
1781 | | assert!( |
1782 | | Date64Type::add_day_time_opt(near_max_millis, interval_large_days).is_none(), |
1783 | | "Adding 100M days to large year should fail" |
1784 | | ); |
1785 | | |
1786 | | // Test with a date far from the boundary in the negative direction |
1787 | | let near_min_date = NaiveDate::from_ymd_opt(-200000, 2, 1).unwrap(); |
1788 | | let near_min_millis = (near_min_date - epoch).num_milliseconds(); |
1789 | | |
1790 | | // Subtracting 30 days should succeed |
1791 | | let interval_minus_30_days = IntervalDayTime::new(-30, 0); |
1792 | | assert!( |
1793 | | Date64Type::add_day_time_opt(near_min_millis, interval_minus_30_days).is_some(), |
1794 | | "Subtracting 30 days from large negative year should succeed" |
1795 | | ); |
1796 | | |
1797 | | // Subtracting a very large number of days should fail |
1798 | | let interval_minus_large_days = IntervalDayTime::new(-100000000, 0); |
1799 | | assert!( |
1800 | | Date64Type::add_day_time_opt(near_min_millis, interval_minus_large_days).is_none(), |
1801 | | "Subtracting 100M days from large negative year should fail" |
1802 | | ); |
1803 | | |
1804 | | // Test normal case within valid range |
1805 | | let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); |
1806 | | let year_2000_millis = (year_2000 - epoch).num_milliseconds(); |
1807 | | let interval_1000_days = IntervalDayTime::new(1000, 12345); |
1808 | | assert!( |
1809 | | Date64Type::add_day_time_opt(year_2000_millis, interval_1000_days).is_some(), |
1810 | | "Adding 1000 days and time to year 2000 should succeed" |
1811 | | ); |
1812 | | |
1813 | | // Test with extreme input values that would cause overflow |
1814 | | let interval_one_day = IntervalDayTime::new(1, 0); |
1815 | | assert!( |
1816 | | Date64Type::add_day_time_opt(i64::MAX, interval_one_day).is_none(), |
1817 | | "Adding interval to i64::MAX should fail" |
1818 | | ); |
1819 | | assert!( |
1820 | | Date64Type::add_day_time_opt(i64::MIN, IntervalDayTime::new(-1, 0)).is_none(), |
1821 | | "Subtracting interval from i64::MIN should fail" |
1822 | | ); |
1823 | | |
1824 | | // Test with extreme interval values |
1825 | | let max_interval = IntervalDayTime::new(i32::MAX, i32::MAX); |
1826 | | assert!( |
1827 | | Date64Type::add_day_time_opt(0, max_interval).is_none(), |
1828 | | "Adding extreme interval should fail" |
1829 | | ); |
1830 | | |
1831 | | let min_interval = IntervalDayTime::new(i32::MIN, i32::MIN); |
1832 | | assert!( |
1833 | | Date64Type::add_day_time_opt(0, min_interval).is_none(), |
1834 | | "Adding extreme negative interval should fail" |
1835 | | ); |
1836 | | |
1837 | | // Test millisecond overflow within a day |
1838 | | let large_ms_interval = IntervalDayTime::new(0, i32::MAX); |
1839 | | assert!( |
1840 | | Date64Type::add_day_time_opt(year_2000_millis, large_ms_interval).is_some(), |
1841 | | "Adding large milliseconds within valid range should succeed" |
1842 | | ); |
1843 | | } |
1844 | | |
1845 | | #[test] |
1846 | | fn test_date64_add_month_day_nano_opt_boundary_values() { |
1847 | | use arrow_array::types::Date64Type; |
1848 | | use arrow_buffer::IntervalMonthDayNano; |
1849 | | |
1850 | | let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); |
1851 | | |
1852 | | // Test with a large year that is still within chrono's safe range |
1853 | | let near_max_date = NaiveDate::from_ymd_opt(5000, 11, 1).unwrap(); |
1854 | | let near_max_millis = (near_max_date - epoch).num_milliseconds(); |
1855 | | |
1856 | | // Adding 1 month and 30 days should succeed |
1857 | | let interval_safe = IntervalMonthDayNano::new(1, 30, 0); |
1858 | | assert!( |
1859 | | Date64Type::add_month_day_nano_opt(near_max_millis, interval_safe).is_some(), |
1860 | | "Adding 1 month 30 days to large year should succeed" |
1861 | | ); |
1862 | | |
1863 | | // Test normal case within valid range |
1864 | | let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); |
1865 | | let year_2000_millis = (year_2000 - epoch).num_milliseconds(); |
1866 | | |
1867 | | // Test edge case: adding zero should always work for valid dates |
1868 | | let zero_interval = IntervalMonthDayNano::new(0, 0, 0); |
1869 | | assert!( |
1870 | | Date64Type::add_month_day_nano_opt(year_2000_millis, zero_interval).is_some(), |
1871 | | "Adding zero interval should always succeed for valid dates" |
1872 | | ); |
1873 | | |
1874 | | // Test with a negative year that is still within chrono's safe range |
1875 | | let near_min_date = NaiveDate::from_ymd_opt(-5000, 2, 28).unwrap(); |
1876 | | let near_min_millis = (near_min_date - epoch).num_milliseconds(); |
1877 | | |
1878 | | // Subtracting 1 month and 30 days should succeed |
1879 | | let interval_safe_neg = IntervalMonthDayNano::new(-1, -30, 0); |
1880 | | assert!( |
1881 | | Date64Type::add_month_day_nano_opt(near_min_millis, interval_safe_neg).is_some(), |
1882 | | "Subtracting 1 month 30 days from large negative year should succeed" |
1883 | | ); |
1884 | | |
1885 | | // Test with extreme input values that would cause overflow |
1886 | | assert!( |
1887 | | Date64Type::add_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(1, 0, 0)) |
1888 | | .is_none(), |
1889 | | "Adding interval to i64::MAX should fail" |
1890 | | ); |
1891 | | |
1892 | | let interval_normal = IntervalMonthDayNano::new(2, 10, 123_456_789_000); |
1893 | | assert!( |
1894 | | Date64Type::add_month_day_nano_opt(year_2000_millis, interval_normal).is_some(), |
1895 | | "Adding 2 months, 10 days, and nanos to year 2000 should succeed" |
1896 | | ); |
1897 | | |
1898 | | // Test with extreme input values that would cause overflow |
1899 | | assert!( |
1900 | | Date64Type::add_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(1, 0, 0)) |
1901 | | .is_none(), |
1902 | | "Adding interval to i64::MAX should fail" |
1903 | | ); |
1904 | | assert!( |
1905 | | Date64Type::add_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(-1, 0, 0)) |
1906 | | .is_none(), |
1907 | | "Subtracting interval from i64::MIN should fail" |
1908 | | ); |
1909 | | |
1910 | | // Test with invalid timestamp input (the _opt function should handle these gracefully) |
1911 | | |
1912 | | // Test nanosecond precision (should not affect boundary since it's < 1ms) |
1913 | | let nano_interval = IntervalMonthDayNano::new(0, 0, 999_999_999); |
1914 | | assert!( |
1915 | | Date64Type::add_month_day_nano_opt(year_2000_millis, nano_interval).is_some(), |
1916 | | "Adding nanoseconds within valid range should succeed" |
1917 | | ); |
1918 | | |
1919 | | // Test large nanosecond values that convert to milliseconds |
1920 | | let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); // 1 day in nanos |
1921 | | assert!( |
1922 | | Date64Type::add_month_day_nano_opt(year_2000_millis, large_nano_interval).is_some(), |
1923 | | "Adding 1 day worth of nanoseconds should succeed" |
1924 | | ); |
1925 | | } |
1926 | | |
1927 | | #[test] |
1928 | | fn test_date64_subtract_year_months_opt_boundary_values() { |
1929 | | use arrow_array::types::Date64Type; |
1930 | | |
1931 | | let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); |
1932 | | |
1933 | | // Test with a negative year that is still within chrono's safe range |
1934 | | let near_min_date = NaiveDate::from_ymd_opt(-5000, 12, 31).unwrap(); |
1935 | | let near_min_millis = (near_min_date - epoch).num_milliseconds(); |
1936 | | |
1937 | | // Subtracting 12 months should succeed |
1938 | | assert!( |
1939 | | Date64Type::subtract_year_months_opt(near_min_millis, 12).is_some(), |
1940 | | "Subtracting 12 months from year -5000 should succeed" |
1941 | | ); |
1942 | | |
1943 | | // Test normal case within valid range |
1944 | | let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); |
1945 | | let year_2000_millis = (year_2000 - epoch).num_milliseconds(); |
1946 | | |
1947 | | // Test edge case: subtracting zero should always work for valid dates |
1948 | | assert!( |
1949 | | Date64Type::subtract_year_months_opt(year_2000_millis, 0).is_some(), |
1950 | | "Subtracting zero months should always succeed for valid dates" |
1951 | | ); |
1952 | | |
1953 | | // Test with a large year that is still within chrono's safe range |
1954 | | let near_max_date = NaiveDate::from_ymd_opt(5000, 1, 1).unwrap(); |
1955 | | let near_max_millis = (near_max_date - epoch).num_milliseconds(); |
1956 | | |
1957 | | // Adding 12 months (subtracting negative) should succeed |
1958 | | assert!( |
1959 | | Date64Type::subtract_year_months_opt(near_max_millis, -12).is_some(), |
1960 | | "Adding 12 months to year 5000 should succeed" |
1961 | | ); |
1962 | | |
1963 | | // Test with extreme input values that would cause overflow |
1964 | | assert!( |
1965 | | Date64Type::subtract_year_months_opt(i64::MAX, -1).is_none(), |
1966 | | "Adding months to i64::MAX should fail" |
1967 | | ); |
1968 | | |
1969 | | assert!( |
1970 | | Date64Type::subtract_year_months_opt(year_2000_millis, 12).is_some(), |
1971 | | "Subtracting 1 year from year 2000 should succeed" |
1972 | | ); |
1973 | | |
1974 | | // Test with extreme input values that would cause overflow |
1975 | | assert!( |
1976 | | Date64Type::subtract_year_months_opt(i64::MAX, -1).is_none(), |
1977 | | "Adding months to i64::MAX should fail" |
1978 | | ); |
1979 | | assert!( |
1980 | | Date64Type::subtract_year_months_opt(i64::MIN, 1).is_none(), |
1981 | | "Subtracting months from i64::MIN should fail" |
1982 | | ); |
1983 | | |
1984 | | // Test edge case: subtracting zero should always work for valid dates |
1985 | | let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap(); |
1986 | | let valid_millis = (valid_date - epoch).num_milliseconds(); |
1987 | | assert!( |
1988 | | Date64Type::subtract_year_months_opt(valid_millis, 0).is_some(), |
1989 | | "Subtracting zero months should always succeed for valid dates" |
1990 | | ); |
1991 | | } |
1992 | | |
1993 | | #[test] |
1994 | | fn test_date64_subtract_day_time_opt_boundary_values() { |
1995 | | use arrow_array::types::Date64Type; |
1996 | | use arrow_buffer::IntervalDayTime; |
1997 | | |
1998 | | let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); |
1999 | | |
2000 | | // Test with a date far from the boundary in the negative direction |
2001 | | let near_min_date = NaiveDate::from_ymd_opt(-200000, 2, 1).unwrap(); |
2002 | | let near_min_millis = (near_min_date - epoch).num_milliseconds(); |
2003 | | |
2004 | | // Subtracting 30 days should succeed |
2005 | | let interval_30_days = IntervalDayTime::new(30, 0); |
2006 | | assert!( |
2007 | | Date64Type::subtract_day_time_opt(near_min_millis, interval_30_days).is_some(), |
2008 | | "Subtracting 30 days from large negative year should succeed" |
2009 | | ); |
2010 | | |
2011 | | // Subtracting a very large number of days should fail |
2012 | | let interval_large_days = IntervalDayTime::new(100000000, 0); |
2013 | | assert!( |
2014 | | Date64Type::subtract_day_time_opt(near_min_millis, interval_large_days).is_none(), |
2015 | | "Subtracting 100M days from large negative year should fail" |
2016 | | ); |
2017 | | |
2018 | | // Test with a date far from the boundary but still testing the function |
2019 | | let near_max_date = NaiveDate::from_ymd_opt(200000, 12, 1).unwrap(); |
2020 | | let near_max_millis = (near_max_date - epoch).num_milliseconds(); |
2021 | | |
2022 | | // Adding 30 days (subtracting negative) should succeed |
2023 | | let interval_minus_30_days = IntervalDayTime::new(-30, 0); |
2024 | | assert!( |
2025 | | Date64Type::subtract_day_time_opt(near_max_millis, interval_minus_30_days).is_some(), |
2026 | | "Adding 30 days to large year should succeed" |
2027 | | ); |
2028 | | |
2029 | | // Adding a very large number of days should fail |
2030 | | let interval_minus_large_days = IntervalDayTime::new(-100000000, 0); |
2031 | | assert!( |
2032 | | Date64Type::subtract_day_time_opt(near_max_millis, interval_minus_large_days).is_none(), |
2033 | | "Adding 100M days to large year should fail" |
2034 | | ); |
2035 | | |
2036 | | // Test normal case within valid range |
2037 | | let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); |
2038 | | let year_2000_millis = (year_2000 - epoch).num_milliseconds(); |
2039 | | let interval_1000_days = IntervalDayTime::new(1000, 12345); |
2040 | | assert!( |
2041 | | Date64Type::subtract_day_time_opt(year_2000_millis, interval_1000_days).is_some(), |
2042 | | "Subtracting 1000 days and time from year 2000 should succeed" |
2043 | | ); |
2044 | | |
2045 | | // Test with extreme input values that would cause overflow |
2046 | | let interval_one_day = IntervalDayTime::new(1, 0); |
2047 | | assert!( |
2048 | | Date64Type::subtract_day_time_opt(i64::MIN, interval_one_day).is_none(), |
2049 | | "Subtracting interval from i64::MIN should fail" |
2050 | | ); |
2051 | | assert!( |
2052 | | Date64Type::subtract_day_time_opt(i64::MAX, IntervalDayTime::new(-1, 0)).is_none(), |
2053 | | "Adding interval to i64::MAX should fail" |
2054 | | ); |
2055 | | |
2056 | | // Test with extreme interval values |
2057 | | let max_interval = IntervalDayTime::new(i32::MAX, i32::MAX); |
2058 | | assert!( |
2059 | | Date64Type::subtract_day_time_opt(0, max_interval).is_none(), |
2060 | | "Subtracting extreme interval should fail" |
2061 | | ); |
2062 | | |
2063 | | let min_interval = IntervalDayTime::new(i32::MIN, i32::MIN); |
2064 | | assert!( |
2065 | | Date64Type::subtract_day_time_opt(0, min_interval).is_none(), |
2066 | | "Subtracting extreme negative interval should fail" |
2067 | | ); |
2068 | | |
2069 | | // Test millisecond precision |
2070 | | let large_ms_interval = IntervalDayTime::new(0, i32::MAX); |
2071 | | assert!( |
2072 | | Date64Type::subtract_day_time_opt(year_2000_millis, large_ms_interval).is_some(), |
2073 | | "Subtracting large milliseconds within valid range should succeed" |
2074 | | ); |
2075 | | |
2076 | | // Test edge case: subtracting zero should always work for valid dates |
2077 | | let zero_interval = IntervalDayTime::new(0, 0); |
2078 | | let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap(); |
2079 | | let valid_millis = (valid_date - epoch).num_milliseconds(); |
2080 | | assert!( |
2081 | | Date64Type::subtract_day_time_opt(valid_millis, zero_interval).is_some(), |
2082 | | "Subtracting zero interval should always succeed for valid dates" |
2083 | | ); |
2084 | | } |
2085 | | |
2086 | | #[test] |
2087 | | fn test_date64_subtract_month_day_nano_opt_boundary_values() { |
2088 | | use arrow_array::types::Date64Type; |
2089 | | use arrow_buffer::IntervalMonthDayNano; |
2090 | | |
2091 | | let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); |
2092 | | |
2093 | | // Test with a negative year that is still within chrono's safe range |
2094 | | let near_min_date = NaiveDate::from_ymd_opt(-5000, 2, 28).unwrap(); |
2095 | | let near_min_millis = (near_min_date - epoch).num_milliseconds(); |
2096 | | |
2097 | | // Subtracting 1 month and 30 days should succeed |
2098 | | let interval_safe = IntervalMonthDayNano::new(1, 30, 0); |
2099 | | assert!( |
2100 | | Date64Type::subtract_month_day_nano_opt(near_min_millis, interval_safe).is_some(), |
2101 | | "Subtracting 1 month 30 days from large negative year should succeed" |
2102 | | ); |
2103 | | |
2104 | | // Test normal case within valid range |
2105 | | let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); |
2106 | | let year_2000_millis = (year_2000 - epoch).num_milliseconds(); |
2107 | | |
2108 | | // Test edge case: subtracting zero should always work for valid dates |
2109 | | let zero_interval = IntervalMonthDayNano::new(0, 0, 0); |
2110 | | assert!( |
2111 | | Date64Type::subtract_month_day_nano_opt(year_2000_millis, zero_interval).is_some(), |
2112 | | "Subtracting zero interval should always succeed for valid dates" |
2113 | | ); |
2114 | | |
2115 | | // Test with a large year that is still within chrono's safe range |
2116 | | let near_max_date = NaiveDate::from_ymd_opt(5000, 11, 1).unwrap(); |
2117 | | let near_max_millis = (near_max_date - epoch).num_milliseconds(); |
2118 | | |
2119 | | // Adding 1 month and 30 days (subtracting negative) should succeed |
2120 | | let interval_safe_neg = IntervalMonthDayNano::new(-1, -30, 0); |
2121 | | assert!( |
2122 | | Date64Type::subtract_month_day_nano_opt(near_max_millis, interval_safe_neg).is_some(), |
2123 | | "Adding 1 month 30 days to large year should succeed" |
2124 | | ); |
2125 | | |
2126 | | // Test with extreme input values that would cause overflow |
2127 | | assert!( |
2128 | | Date64Type::subtract_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(1, 0, 0)) |
2129 | | .is_none(), |
2130 | | "Subtracting interval from i64::MIN should fail" |
2131 | | ); |
2132 | | |
2133 | | let interval_normal = IntervalMonthDayNano::new(2, 10, 123_456_789_000); |
2134 | | assert!( |
2135 | | Date64Type::subtract_month_day_nano_opt(year_2000_millis, interval_normal).is_some(), |
2136 | | "Subtracting 2 months, 10 days, and nanos from year 2000 should succeed" |
2137 | | ); |
2138 | | |
2139 | | // Test with extreme input values that would cause overflow |
2140 | | assert!( |
2141 | | Date64Type::subtract_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(1, 0, 0)) |
2142 | | .is_none(), |
2143 | | "Subtracting interval from i64::MIN should fail" |
2144 | | ); |
2145 | | assert!( |
2146 | | Date64Type::subtract_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(-1, 0, 0)) |
2147 | | .is_none(), |
2148 | | "Adding interval to i64::MAX should fail" |
2149 | | ); |
2150 | | |
2151 | | // Test nanosecond precision (should not affect boundary since it's < 1ms) |
2152 | | let nano_interval = IntervalMonthDayNano::new(0, 0, 999_999_999); |
2153 | | assert!( |
2154 | | Date64Type::subtract_month_day_nano_opt(year_2000_millis, nano_interval).is_some(), |
2155 | | "Subtracting nanoseconds within valid range should succeed" |
2156 | | ); |
2157 | | |
2158 | | // Test large nanosecond values that convert to milliseconds |
2159 | | let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); // 1 day in nanos |
2160 | | assert!( |
2161 | | Date64Type::subtract_month_day_nano_opt(year_2000_millis, large_nano_interval) |
2162 | | .is_some(), |
2163 | | "Subtracting 1 day worth of nanoseconds should succeed" |
2164 | | ); |
2165 | | |
2166 | | // Test edge case: subtracting zero should always work for valid dates |
2167 | | let zero_interval = IntervalMonthDayNano::new(0, 0, 0); |
2168 | | let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap(); |
2169 | | let valid_millis = (valid_date - epoch).num_milliseconds(); |
2170 | | assert!( |
2171 | | Date64Type::subtract_month_day_nano_opt(valid_millis, zero_interval).is_some(), |
2172 | | "Subtracting zero interval should always succeed for valid dates" |
2173 | | ); |
2174 | | } |
2175 | | } |