1#[forbid(unsafe_code)]
15pub mod basic;
16#[cfg(feature = "compute_arithmetics_decimal")]
17pub mod decimal;
18pub mod time;
19
20use crate::{
21 array::{Array, DictionaryArray, PrimitiveArray},
22 bitmap::Bitmap,
23 datatypes::{DataType, IntervalUnit, TimeUnit},
24 scalar::{PrimitiveScalar, Scalar},
25 types::NativeType,
26};
27
28fn binary_dyn<T: NativeType, F: Fn(&PrimitiveArray<T>, &PrimitiveArray<T>) -> PrimitiveArray<T>>(
29 lhs: &dyn Array,
30 rhs: &dyn Array,
31 op: F,
32) -> Box<dyn Array> {
33 let lhs = lhs.as_any().downcast_ref().unwrap();
34 let rhs = rhs.as_any().downcast_ref().unwrap();
35 op(lhs, rhs).boxed()
36}
37
38macro_rules! arith {
41 ($lhs:expr, $rhs:expr, $op:tt $(, decimal = $op_decimal:tt )? $(, duration = $op_duration:tt )? $(, interval = $op_interval:tt )? $(, timestamp = $op_timestamp:tt )?) => {{
42 let lhs = $lhs;
43 let rhs = $rhs;
44 use DataType::*;
45 match (lhs.data_type(), rhs.data_type()) {
46 (Int8, Int8) => binary_dyn::<i8, _>(lhs, rhs, basic::$op),
47 (Int16, Int16) => binary_dyn::<i16, _>(lhs, rhs, basic::$op),
48 (Int32, Int32) => binary_dyn::<i32, _>(lhs, rhs, basic::$op),
49 (Int64, Int64) | (Duration(_), Duration(_)) => {
50 binary_dyn::<i64, _>(lhs, rhs, basic::$op)
51 }
52 (UInt8, UInt8) => binary_dyn::<u8, _>(lhs, rhs, basic::$op),
53 (UInt16, UInt16) => binary_dyn::<u16, _>(lhs, rhs, basic::$op),
54 (UInt32, UInt32) => binary_dyn::<u32, _>(lhs, rhs, basic::$op),
55 (UInt64, UInt64) => binary_dyn::<u64, _>(lhs, rhs, basic::$op),
56 (Float32, Float32) => binary_dyn::<f32, _>(lhs, rhs, basic::$op),
57 (Float64, Float64) => binary_dyn::<f64, _>(lhs, rhs, basic::$op),
58 $ (
59 (Decimal(_, _), Decimal(_, _)) => {
60 let lhs = lhs.as_any().downcast_ref().unwrap();
61 let rhs = rhs.as_any().downcast_ref().unwrap();
62 Box::new(decimal::$op_decimal(lhs, rhs)) as Box<dyn Array>
63 }
64 )?
65 $ (
66 (Time32(TimeUnit::Second), Duration(_))
67 | (Time32(TimeUnit::Millisecond), Duration(_))
68 | (Date32, Duration(_)) => {
69 let lhs = lhs.as_any().downcast_ref().unwrap();
70 let rhs = rhs.as_any().downcast_ref().unwrap();
71 Box::new(time::$op_duration::<i32>(lhs, rhs)) as Box<dyn Array>
72 }
73 (Time64(TimeUnit::Microsecond), Duration(_))
74 | (Time64(TimeUnit::Nanosecond), Duration(_))
75 | (Date64, Duration(_))
76 | (Timestamp(_, _), Duration(_)) => {
77 let lhs = lhs.as_any().downcast_ref().unwrap();
78 let rhs = rhs.as_any().downcast_ref().unwrap();
79 Box::new(time::$op_duration::<i64>(lhs, rhs)) as Box<dyn Array>
80 }
81 )?
82 $ (
83 (Timestamp(_, _), Interval(IntervalUnit::MonthDayNano)) => {
84 let lhs = lhs.as_any().downcast_ref().unwrap();
85 let rhs = rhs.as_any().downcast_ref().unwrap();
86 time::$op_interval(lhs, rhs).map(|x| Box::new(x) as Box<dyn Array>).unwrap()
87 }
88 )?
89 $ (
90 (Timestamp(_, None), Timestamp(_, None)) => {
91 let lhs = lhs.as_any().downcast_ref().unwrap();
92 let rhs = rhs.as_any().downcast_ref().unwrap();
93 time::$op_timestamp(lhs, rhs).map(|x| Box::new(x) as Box<dyn Array>).unwrap()
94 }
95 )?
96 _ => todo!(
97 "Addition of {:?} with {:?} is not supported",
98 lhs.data_type(),
99 rhs.data_type()
100 ),
101 }
102 }};
103}
104
105fn binary_scalar<T: NativeType, F: Fn(&PrimitiveArray<T>, &T) -> PrimitiveArray<T>>(
106 lhs: &PrimitiveArray<T>,
107 rhs: &PrimitiveScalar<T>,
108 op: F,
109) -> PrimitiveArray<T> {
110 let rhs = if let Some(rhs) = *rhs.value() {
111 rhs
112 } else {
113 return PrimitiveArray::<T>::new_null(lhs.data_type().clone(), lhs.len());
114 };
115 op(lhs, &rhs)
116}
117
118fn binary_scalar_dyn<T: NativeType, F: Fn(&PrimitiveArray<T>, &T) -> PrimitiveArray<T>>(
119 lhs: &dyn Array,
120 rhs: &dyn Scalar,
121 op: F,
122) -> Box<dyn Array> {
123 let lhs = lhs.as_any().downcast_ref().unwrap();
124 let rhs = rhs.as_any().downcast_ref().unwrap();
125 binary_scalar(lhs, rhs, op).boxed()
126}
127
128macro_rules! arith_scalar {
131 ($lhs:expr, $rhs:expr, $op:tt $(, decimal = $op_decimal:tt )? $(, duration = $op_duration:tt )? $(, interval = $op_interval:tt )? $(, timestamp = $op_timestamp:tt )?) => {{
132 let lhs = $lhs;
133 let rhs = $rhs;
134 use DataType::*;
135 match (lhs.data_type(), rhs.data_type()) {
136 (Int8, Int8) => binary_scalar_dyn::<i8, _>(lhs, rhs, basic::$op),
137 (Int16, Int16) => binary_scalar_dyn::<i16, _>(lhs, rhs, basic::$op),
138 (Int32, Int32) => binary_scalar_dyn::<i32, _>(lhs, rhs, basic::$op),
139 (Int64, Int64) | (Duration(_), Duration(_)) => {
140 binary_scalar_dyn::<i64, _>(lhs, rhs, basic::$op)
141 }
142 (UInt8, UInt8) => binary_scalar_dyn::<u8, _>(lhs, rhs, basic::$op),
143 (UInt16, UInt16) => binary_scalar_dyn::<u16, _>(lhs, rhs, basic::$op),
144 (UInt32, UInt32) => binary_scalar_dyn::<u32, _>(lhs, rhs, basic::$op),
145 (UInt64, UInt64) => binary_scalar_dyn::<u64, _>(lhs, rhs, basic::$op),
146 (Float32, Float32) => binary_scalar_dyn::<f32, _>(lhs, rhs, basic::$op),
147 (Float64, Float64) => binary_scalar_dyn::<f64, _>(lhs, rhs, basic::$op),
148 $ (
149 (Decimal(_, _), Decimal(_, _)) => {
150 let lhs = lhs.as_any().downcast_ref().unwrap();
151 let rhs = rhs.as_any().downcast_ref().unwrap();
152 decimal::$op_decimal(lhs, rhs).boxed()
153 }
154 )?
155 $ (
156 (Time32(TimeUnit::Second), Duration(_))
157 | (Time32(TimeUnit::Millisecond), Duration(_))
158 | (Date32, Duration(_)) => {
159 let lhs = lhs.as_any().downcast_ref().unwrap();
160 let rhs = rhs.as_any().downcast_ref().unwrap();
161 time::$op_duration::<i32>(lhs, rhs).boxed()
162 }
163 (Time64(TimeUnit::Microsecond), Duration(_))
164 | (Time64(TimeUnit::Nanosecond), Duration(_))
165 | (Date64, Duration(_))
166 | (Timestamp(_, _), Duration(_)) => {
167 let lhs = lhs.as_any().downcast_ref().unwrap();
168 let rhs = rhs.as_any().downcast_ref().unwrap();
169 time::$op_duration::<i64>(lhs, rhs).boxed()
170 }
171 )?
172 $ (
173 (Timestamp(_, _), Interval(IntervalUnit::MonthDayNano)) => {
174 let lhs = lhs.as_any().downcast_ref().unwrap();
175 let rhs = rhs.as_any().downcast_ref().unwrap();
176 time::$op_interval(lhs, rhs).unwrap().boxed()
177 }
178 )?
179 $ (
180 (Timestamp(_, None), Timestamp(_, None)) => {
181 let lhs = lhs.as_any().downcast_ref().unwrap();
182 let rhs = rhs.as_any().downcast_ref().unwrap();
183 time::$op_timestamp(lhs, rhs).unwrap().boxed()
184 }
185 )?
186 _ => todo!(
187 "Addition of {:?} with {:?} is not supported",
188 lhs.data_type(),
189 rhs.data_type()
190 ),
191 }
192 }};
193}
194
195pub fn add(lhs: &dyn Array, rhs: &dyn Array) -> Box<dyn Array> {
202 arith!(
203 lhs,
204 rhs,
205 add,
206 duration = add_duration,
207 interval = add_interval
208 )
209}
210
211pub fn add_scalar(lhs: &dyn Array, rhs: &dyn Scalar) -> Box<dyn Array> {
218 arith_scalar!(
219 lhs,
220 rhs,
221 add_scalar,
222 duration = add_duration_scalar,
223 interval = add_interval_scalar
224 )
225}
226
227pub fn can_add(lhs: &DataType, rhs: &DataType) -> bool {
229 use DataType::*;
230 matches!(
231 (lhs, rhs),
232 (Int8, Int8)
233 | (Int16, Int16)
234 | (Int32, Int32)
235 | (Int64, Int64)
236 | (UInt8, UInt8)
237 | (UInt16, UInt16)
238 | (UInt32, UInt32)
239 | (UInt64, UInt64)
240 | (Float64, Float64)
241 | (Float32, Float32)
242 | (Duration(_), Duration(_))
243 | (Decimal(_, _), Decimal(_, _))
244 | (Date32, Duration(_))
245 | (Date64, Duration(_))
246 | (Time32(TimeUnit::Millisecond), Duration(_))
247 | (Time32(TimeUnit::Second), Duration(_))
248 | (Time64(TimeUnit::Microsecond), Duration(_))
249 | (Time64(TimeUnit::Nanosecond), Duration(_))
250 | (Timestamp(_, _), Duration(_))
251 | (Timestamp(_, _), Interval(IntervalUnit::MonthDayNano))
252 )
253}
254
255pub fn sub(lhs: &dyn Array, rhs: &dyn Array) -> Box<dyn Array> {
262 arith!(
263 lhs,
264 rhs,
265 sub,
266 decimal = sub,
267 duration = subtract_duration,
268 timestamp = subtract_timestamps
269 )
270}
271
272pub fn sub_scalar(lhs: &dyn Array, rhs: &dyn Scalar) -> Box<dyn Array> {
279 arith_scalar!(
280 lhs,
281 rhs,
282 sub_scalar,
283 duration = sub_duration_scalar,
284 timestamp = sub_timestamps_scalar
285 )
286}
287
288pub fn can_sub(lhs: &DataType, rhs: &DataType) -> bool {
290 use DataType::*;
291 matches!(
292 (lhs, rhs),
293 (Int8, Int8)
294 | (Int16, Int16)
295 | (Int32, Int32)
296 | (Int64, Int64)
297 | (UInt8, UInt8)
298 | (UInt16, UInt16)
299 | (UInt32, UInt32)
300 | (UInt64, UInt64)
301 | (Float64, Float64)
302 | (Float32, Float32)
303 | (Duration(_), Duration(_))
304 | (Decimal(_, _), Decimal(_, _))
305 | (Date32, Duration(_))
306 | (Date64, Duration(_))
307 | (Time32(TimeUnit::Millisecond), Duration(_))
308 | (Time32(TimeUnit::Second), Duration(_))
309 | (Time64(TimeUnit::Microsecond), Duration(_))
310 | (Time64(TimeUnit::Nanosecond), Duration(_))
311 | (Timestamp(_, _), Duration(_))
312 | (Timestamp(_, None), Timestamp(_, None))
313 )
314}
315
316pub fn mul(lhs: &dyn Array, rhs: &dyn Array) -> Box<dyn Array> {
322 arith!(lhs, rhs, mul, decimal = mul)
323}
324
325pub fn mul_scalar(lhs: &dyn Array, rhs: &dyn Scalar) -> Box<dyn Array> {
330 arith_scalar!(lhs, rhs, mul_scalar, decimal = mul_scalar)
331}
332
333pub fn can_mul(lhs: &DataType, rhs: &DataType) -> bool {
335 use DataType::*;
336 matches!(
337 (lhs, rhs),
338 (Int8, Int8)
339 | (Int16, Int16)
340 | (Int32, Int32)
341 | (Int64, Int64)
342 | (UInt8, UInt8)
343 | (UInt16, UInt16)
344 | (UInt32, UInt32)
345 | (UInt64, UInt64)
346 | (Float64, Float64)
347 | (Float32, Float32)
348 | (Decimal(_, _), Decimal(_, _))
349 )
350}
351
352pub fn div(lhs: &dyn Array, rhs: &dyn Array) -> Box<dyn Array> {
358 arith!(lhs, rhs, div, decimal = div)
359}
360
361pub fn div_scalar(lhs: &dyn Array, rhs: &dyn Scalar) -> Box<dyn Array> {
366 arith_scalar!(lhs, rhs, div_scalar, decimal = div_scalar)
367}
368
369pub fn can_div(lhs: &DataType, rhs: &DataType) -> bool {
371 can_mul(lhs, rhs)
372}
373
374pub fn rem(lhs: &dyn Array, rhs: &dyn Array) -> Box<dyn Array> {
380 arith!(lhs, rhs, rem)
381}
382
383pub fn can_rem(lhs: &DataType, rhs: &DataType) -> bool {
385 use DataType::*;
386 matches!(
387 (lhs, rhs),
388 (Int8, Int8)
389 | (Int16, Int16)
390 | (Int32, Int32)
391 | (Int64, Int64)
392 | (UInt8, UInt8)
393 | (UInt16, UInt16)
394 | (UInt32, UInt32)
395 | (UInt64, UInt64)
396 | (Float64, Float64)
397 | (Float32, Float32)
398 )
399}
400
401macro_rules! with_match_negatable {(
402 $key_type:expr, | $_:tt $T:ident | $($body:tt)*
403) => ({
404 macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
405 use crate::datatypes::PrimitiveType::*;
406 use crate::types::{days_ms, months_days_ns, i256};
407 match $key_type {
408 Int8 => __with_ty__! { i8 },
409 Int16 => __with_ty__! { i16 },
410 Int32 => __with_ty__! { i32 },
411 Int64 => __with_ty__! { i64 },
412 Int128 => __with_ty__! { i128 },
413 Int256 => __with_ty__! { i256 },
414 DaysMs => __with_ty__! { days_ms },
415 MonthDayNano => __with_ty__! { months_days_ns },
416 UInt8 | UInt16 | UInt32 | UInt64 | Float16 => todo!(),
417 Float32 => __with_ty__! { f32 },
418 Float64 => __with_ty__! { f64 },
419 }
420})}
421
422pub fn neg(array: &dyn Array) -> Box<dyn Array> {
428 use crate::datatypes::PhysicalType::*;
429 match array.data_type().to_physical_type() {
430 Primitive(primitive) => with_match_negatable!(primitive, |$T| {
431 let array = array.as_any().downcast_ref().unwrap();
432
433 let result = basic::negate::<$T>(array);
434 Box::new(result) as Box<dyn Array>
435 }),
436 Dictionary(key) => match_integer_type!(key, |$T| {
437 let array = array.as_any().downcast_ref::<DictionaryArray<$T>>().unwrap();
438
439 let values = neg(array.values().as_ref());
440
441 unsafe{
443 DictionaryArray::<$T>::try_new_unchecked(array.data_type().clone(), array.keys().clone(), values).unwrap().boxed()
444 }
445 }),
446 _ => todo!(),
447 }
448}
449
450pub fn can_neg(data_type: &DataType) -> bool {
452 if let DataType::Dictionary(_, values, _) = data_type.to_logical_type() {
453 return can_neg(values.as_ref());
454 }
455
456 use crate::datatypes::PhysicalType::*;
457 use crate::datatypes::PrimitiveType::*;
458 matches!(
459 data_type.to_physical_type(),
460 Primitive(Int8)
461 | Primitive(Int16)
462 | Primitive(Int32)
463 | Primitive(Int64)
464 | Primitive(Float64)
465 | Primitive(Float32)
466 | Primitive(DaysMs)
467 | Primitive(MonthDayNano)
468 )
469}
470
471pub trait ArrayAdd<Rhs>: Sized {
473 fn add(&self, rhs: &Rhs) -> Self;
475}
476
477pub trait ArrayWrappingAdd<Rhs>: Sized {
479 fn wrapping_add(&self, rhs: &Rhs) -> Self;
481}
482
483pub trait ArrayCheckedAdd<Rhs>: Sized {
485 fn checked_add(&self, rhs: &Rhs) -> Self;
487}
488
489pub trait ArraySaturatingAdd<Rhs>: Sized {
491 fn saturating_add(&self, rhs: &Rhs) -> Self;
493}
494
495pub trait ArrayOverflowingAdd<Rhs>: Sized {
497 fn overflowing_add(&self, rhs: &Rhs) -> (Self, Bitmap);
499}
500
501pub trait ArraySub<Rhs>: Sized {
503 fn sub(&self, rhs: &Rhs) -> Self;
505}
506
507pub trait ArrayWrappingSub<Rhs>: Sized {
509 fn wrapping_sub(&self, rhs: &Rhs) -> Self;
511}
512
513pub trait ArrayCheckedSub<Rhs>: Sized {
515 fn checked_sub(&self, rhs: &Rhs) -> Self;
517}
518
519pub trait ArraySaturatingSub<Rhs>: Sized {
521 fn saturating_sub(&self, rhs: &Rhs) -> Self;
523}
524
525pub trait ArrayOverflowingSub<Rhs>: Sized {
527 fn overflowing_sub(&self, rhs: &Rhs) -> (Self, Bitmap);
529}
530
531pub trait ArrayMul<Rhs>: Sized {
533 fn mul(&self, rhs: &Rhs) -> Self;
535}
536
537pub trait ArrayWrappingMul<Rhs>: Sized {
539 fn wrapping_mul(&self, rhs: &Rhs) -> Self;
541}
542
543pub trait ArrayCheckedMul<Rhs>: Sized {
545 fn checked_mul(&self, rhs: &Rhs) -> Self;
547}
548
549pub trait ArraySaturatingMul<Rhs>: Sized {
551 fn saturating_mul(&self, rhs: &Rhs) -> Self;
553}
554
555pub trait ArrayOverflowingMul<Rhs>: Sized {
557 fn overflowing_mul(&self, rhs: &Rhs) -> (Self, Bitmap);
559}
560
561pub trait ArrayDiv<Rhs>: Sized {
563 fn div(&self, rhs: &Rhs) -> Self;
565}
566
567pub trait ArrayCheckedDiv<Rhs>: Sized {
569 fn checked_div(&self, rhs: &Rhs) -> Self;
571}
572
573pub trait ArrayRem<Rhs>: Sized {
575 fn rem(&self, rhs: &Rhs) -> Self;
577}
578
579pub trait ArrayCheckedRem<Rhs>: Sized {
581 fn checked_rem(&self, rhs: &Rhs) -> Self;
583}