You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2022/10/29 04:08:59 UTC

[arrow-rs] branch master updated: Add Decimal Arithmetic (#2881)

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new e6eb8341c Add Decimal Arithmetic (#2881)
e6eb8341c is described below

commit e6eb8341cb20460847ba12949ad7a937a62d76b7
Author: Raphael Taylor-Davies <17...@users.noreply.github.com>
AuthorDate: Sat Oct 29 17:08:54 2022 +1300

    Add Decimal Arithmetic (#2881)
    
    * Add Decimal Arithmetic
    
    * Derive operations
    
    * Add neg
    
    * Clippy
---
 arrow-buffer/src/bigint.rs              | 47 +++++++++++++++++
 arrow/src/compute/kernels/arithmetic.rs | 63 ++++++++++++++++++++++
 arrow/src/datatypes/native.rs           | 12 +++--
 arrow/src/datatypes/numeric.rs          | 93 +++++++++++++++++++++++++++++++++
 4 files changed, 211 insertions(+), 4 deletions(-)

diff --git a/arrow-buffer/src/bigint.rs b/arrow-buffer/src/bigint.rs
index 892c6c99d..463c63729 100644
--- a/arrow-buffer/src/bigint.rs
+++ b/arrow-buffer/src/bigint.rs
@@ -189,6 +189,18 @@ impl i256 {
         (self != Self::MIN).then(|| self.wrapping_abs())
     }
 
+    /// Negates this i256
+    #[inline]
+    pub fn wrapping_neg(self) -> Self {
+        Self::from_parts(!self.low, !self.high).wrapping_add(i256::ONE)
+    }
+
+    /// Negates this i256 returning `None` if `Self == Self::MIN`
+    #[inline]
+    pub fn checked_neg(self) -> Option<Self> {
+        (self != Self::MIN).then(|| self.wrapping_neg())
+    }
+
     /// Performs wrapping addition
     #[inline]
     pub fn wrapping_add(self, other: Self) -> Self {
@@ -396,6 +408,30 @@ fn mulx(a: u128, b: u128) -> (u128, u128) {
     (low, high)
 }
 
+macro_rules! derive_op {
+    ($t:ident, $op:ident, $wrapping:ident, $checked:ident) => {
+        impl std::ops::$t for i256 {
+            type Output = i256;
+
+            #[cfg(debug_assertions)]
+            fn $op(self, rhs: Self) -> Self::Output {
+                self.$checked(rhs).expect("i256 overflow")
+            }
+
+            #[cfg(not(debug_assertions))]
+            fn $op(self, rhs: Self) -> Self::Output {
+                self.$wrapping(rhs)
+            }
+        }
+    };
+}
+
+derive_op!(Add, add, wrapping_add, checked_add);
+derive_op!(Sub, sub, wrapping_sub, checked_sub);
+derive_op!(Mul, mul, wrapping_mul, checked_mul);
+derive_op!(Div, div, wrapping_div, checked_div);
+derive_op!(Rem, rem, wrapping_rem, checked_rem);
+
 macro_rules! define_as_primitive {
     ($native_ty:ty) => {
         impl AsPrimitive<i256> for $native_ty {
@@ -416,6 +452,7 @@ mod tests {
     use super::*;
     use num::{BigInt, FromPrimitive, Signed, ToPrimitive};
     use rand::{thread_rng, Rng};
+    use std::ops::Neg;
 
     #[test]
     fn test_signed_cmp() {
@@ -466,6 +503,16 @@ mod tests {
         assert_eq!(ir.wrapping_abs(), abs);
         assert_eq!(ir.checked_abs().is_none(), overflow);
 
+        // Negation
+        let (neg, overflow) = i256::from_bigint_with_overflow(bl.clone().neg());
+        assert_eq!(il.wrapping_neg(), neg);
+        assert_eq!(il.checked_neg().is_none(), overflow);
+
+        // Negation
+        let (neg, overflow) = i256::from_bigint_with_overflow(br.clone().neg());
+        assert_eq!(ir.wrapping_neg(), neg);
+        assert_eq!(ir.checked_neg().is_none(), overflow);
+
         // Addition
         let actual = il.wrapping_add(ir);
         let (expected, overflow) =
diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs
index e7fcb50cc..d12a0c196 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -1625,6 +1625,7 @@ mod tests {
     use super::*;
     use crate::array::Int32Array;
     use crate::datatypes::{Date64Type, Int32Type, Int8Type};
+    use arrow_buffer::i256;
     use chrono::NaiveDate;
     use half::f16;
 
@@ -2898,6 +2899,68 @@ mod tests {
         overflow.expect_err("overflow should be detected");
     }
 
+    #[test]
+    fn test_decimal128() {
+        let a = Decimal128Array::from_iter_values([1, 2, 4, 5]);
+        let b = Decimal128Array::from_iter_values([7, -3, 6, 3]);
+        let e = Decimal128Array::from_iter_values([8, -1, 10, 8]);
+        let r = add(&a, &b).unwrap();
+        assert_eq!(e, r);
+
+        let e = Decimal128Array::from_iter_values([-6, 5, -2, 2]);
+        let r = subtract(&a, &b).unwrap();
+        assert_eq!(e, r);
+
+        let e = Decimal128Array::from_iter_values([7, -6, 24, 15]);
+        let r = multiply(&a, &b).unwrap();
+        assert_eq!(e, r);
+
+        let a = Decimal128Array::from_iter_values([23, 56, 32, 55]);
+        let b = Decimal128Array::from_iter_values([1, -2, 4, 5]);
+        let e = Decimal128Array::from_iter_values([23, -28, 8, 11]);
+        let r = divide(&a, &b).unwrap();
+        assert_eq!(e, r);
+    }
+
+    #[test]
+    fn test_decimal256() {
+        let a = Decimal256Array::from_iter_values(
+            [1, 2, 4, 5].into_iter().map(i256::from_i128),
+        );
+        let b = Decimal256Array::from_iter_values(
+            [7, -3, 6, 3].into_iter().map(i256::from_i128),
+        );
+        let e = Decimal256Array::from_iter_values(
+            [8, -1, 10, 8].into_iter().map(i256::from_i128),
+        );
+        let r = add(&a, &b).unwrap();
+        assert_eq!(e, r);
+
+        let e = Decimal256Array::from_iter_values(
+            [-6, 5, -2, 2].into_iter().map(i256::from_i128),
+        );
+        let r = subtract(&a, &b).unwrap();
+        assert_eq!(e, r);
+
+        let e = Decimal256Array::from_iter_values(
+            [7, -6, 24, 15].into_iter().map(i256::from_i128),
+        );
+        let r = multiply(&a, &b).unwrap();
+        assert_eq!(e, r);
+
+        let a = Decimal256Array::from_iter_values(
+            [23, 56, 32, 55].into_iter().map(i256::from_i128),
+        );
+        let b = Decimal256Array::from_iter_values(
+            [1, -2, 4, 5].into_iter().map(i256::from_i128),
+        );
+        let e = Decimal256Array::from_iter_values(
+            [23, -28, 8, 11].into_iter().map(i256::from_i128),
+        );
+        let r = divide(&a, &b).unwrap();
+        assert_eq!(e, r);
+    }
+
     #[test]
     #[cfg(feature = "dyn_arith_dict")]
     fn test_dictionary_div_dyn_wrapping_overflow() {
diff --git a/arrow/src/datatypes/native.rs b/arrow/src/datatypes/native.rs
index 2643025f1..bbdec14b4 100644
--- a/arrow/src/datatypes/native.rs
+++ b/arrow/src/datatypes/native.rs
@@ -17,7 +17,7 @@
 
 use crate::error::{ArrowError, Result};
 pub use arrow_array::ArrowPrimitiveType;
-pub use arrow_buffer::{ArrowNativeType, ToByteSlice};
+pub use arrow_buffer::{i256, ArrowNativeType, ToByteSlice};
 use half::f16;
 
 /// Trait for [`ArrowNativeType`] that adds checked and unchecked arithmetic operations,
@@ -85,9 +85,12 @@ pub trait ArrowNativeTypeOp: ArrowNativeType {
 
 macro_rules! native_type_op {
     ($t:tt) => {
+        native_type_op!($t, 0, 1);
+    };
+    ($t:tt, $zero:expr, $one: expr) => {
         impl ArrowNativeTypeOp for $t {
-            const ZERO: Self = 0;
-            const ONE: Self = 1;
+            const ZERO: Self = $zero;
+            const ONE: Self = $one;
 
             fn add_checked(self, rhs: Self) -> Result<Self> {
                 self.checked_add(rhs).ok_or_else(|| {
@@ -173,7 +176,7 @@ macro_rules! native_type_op {
             }
 
             fn is_zero(self) -> bool {
-                self == 0
+                self == Self::ZERO
             }
 
             fn is_eq(self, rhs: Self) -> bool {
@@ -212,6 +215,7 @@ native_type_op!(u8);
 native_type_op!(u16);
 native_type_op!(u32);
 native_type_op!(u64);
+native_type_op!(i256, i256::ZERO, i256::ONE);
 
 macro_rules! native_type_float_op {
     ($t:tt, $zero:expr, $one:expr) => {
diff --git a/arrow/src/datatypes/numeric.rs b/arrow/src/datatypes/numeric.rs
index e74764d4c..61fd05d52 100644
--- a/arrow/src/datatypes/numeric.rs
+++ b/arrow/src/datatypes/numeric.rs
@@ -365,6 +365,7 @@ make_numeric_type!(DurationSecondType, i64, i64x8, m64x8);
 make_numeric_type!(DurationMillisecondType, i64, i64x8, m64x8);
 make_numeric_type!(DurationMicrosecondType, i64, i64x8, m64x8);
 make_numeric_type!(DurationNanosecondType, i64, i64x8, m64x8);
+make_numeric_type!(Decimal128Type, i128, i128x4, m128x4);
 
 #[cfg(not(feature = "simd"))]
 impl ArrowNumericType for Float16Type {}
@@ -462,6 +463,98 @@ impl ArrowNumericType for Float16Type {
     }
 }
 
+#[cfg(not(feature = "simd"))]
+impl ArrowNumericType for Decimal256Type {}
+
+#[cfg(feature = "simd")]
+impl ArrowNumericType for Decimal256Type {
+    type Simd = i256;
+    type SimdMask = bool;
+
+    fn lanes() -> usize {
+        1
+    }
+
+    fn init(value: Self::Native) -> Self::Simd {
+        value
+    }
+
+    fn load(slice: &[Self::Native]) -> Self::Simd {
+        slice[0]
+    }
+
+    fn mask_init(value: bool) -> Self::SimdMask {
+        value
+    }
+
+    fn mask_from_u64(mask: u64) -> Self::SimdMask {
+        mask != 0
+    }
+
+    fn mask_to_u64(mask: &Self::SimdMask) -> u64 {
+        *mask as u64
+    }
+
+    fn mask_get(mask: &Self::SimdMask, _idx: usize) -> bool {
+        *mask
+    }
+
+    fn mask_set(_mask: Self::SimdMask, _idx: usize, value: bool) -> Self::SimdMask {
+        value
+    }
+
+    fn mask_select(mask: Self::SimdMask, a: Self::Simd, b: Self::Simd) -> Self::Simd {
+        match mask {
+            true => a,
+            false => b,
+        }
+    }
+
+    fn mask_any(mask: Self::SimdMask) -> bool {
+        mask
+    }
+
+    fn bin_op<F: Fn(Self::Simd, Self::Simd) -> Self::Simd>(
+        left: Self::Simd,
+        right: Self::Simd,
+        op: F,
+    ) -> Self::Simd {
+        op(left, right)
+    }
+
+    fn eq(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
+        left.eq(&right)
+    }
+
+    fn ne(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
+        left.ne(&right)
+    }
+
+    fn lt(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
+        left.lt(&right)
+    }
+
+    fn le(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
+        left.le(&right)
+    }
+
+    fn gt(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
+        left.gt(&right)
+    }
+
+    fn ge(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
+        left.ge(&right)
+    }
+
+    fn write(simd_result: Self::Simd, slice: &mut [Self::Native]) {
+        slice[0] = simd_result
+    }
+
+    fn unary_op<F: Fn(Self::Simd) -> Self::Simd>(a: Self::Simd, op: F) -> Self::Simd {
+        op(a)
+    }
+}
+
 #[cfg(feature = "simd")]
 pub trait ArrowFloatNumericType: ArrowNumericType {
     fn pow(base: Self::Simd, raise: Self::Simd) -> Self::Simd;