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;