From 91636169b0fb01adcf7254170b18c2cf3bb47b86 Mon Sep 17 00:00:00 2001 From: yjh <465402634@qq.com> Date: Fri, 8 Oct 2021 14:27:33 +0800 Subject: [PATCH] Add Wrapping version ops for primitiveArray (#496) --- src/compute/arithmetics/basic/add.rs | 62 +++++++++++++++++++++++++++- src/compute/arithmetics/basic/mul.rs | 62 +++++++++++++++++++++++++++- src/compute/arithmetics/basic/sub.rs | 62 +++++++++++++++++++++++++++- src/compute/arithmetics/mod.rs | 21 ++++++++++ 4 files changed, 204 insertions(+), 3 deletions(-) diff --git a/src/compute/arithmetics/basic/add.rs b/src/compute/arithmetics/basic/add.rs index 152b903ee9a..73bf1420a4a 100644 --- a/src/compute/arithmetics/basic/add.rs +++ b/src/compute/arithmetics/basic/add.rs @@ -1,9 +1,10 @@ //! Definition of basic add operations with primitive arrays use std::ops::Add; -use num_traits::{ops::overflowing::OverflowingAdd, CheckedAdd, SaturatingAdd, Zero}; +use num_traits::{ops::overflowing::OverflowingAdd, CheckedAdd, SaturatingAdd, WrappingAdd, Zero}; use crate::compute::arithmetics::basic::check_same_type; +use crate::compute::arithmetics::ArrayWrappingAdd; use crate::{ array::{Array, PrimitiveArray}, bitmap::Bitmap, @@ -42,6 +43,34 @@ where binary(lhs, rhs, lhs.data_type().clone(), |a, b| a + b) } +/// Wrapping addition of two [`PrimitiveArray`]s. +/// It wraps around at the boundary of the type if the result overflows. +/// +/// # Examples +/// ``` +/// use arrow2::compute::arithmetics::basic::wrapping_add; +/// use arrow2::array::PrimitiveArray; +/// +/// let a = PrimitiveArray::from([Some(-100i8), Some(100i8), Some(100i8)]); +/// let b = PrimitiveArray::from([Some(0i8), Some(100i8), Some(0i8)]); +/// let result = wrapping_add(&a, &b).unwrap(); +/// let expected = PrimitiveArray::from([Some(-100i8), Some(-56i8), Some(100i8)]); +/// assert_eq!(result, expected); +/// ``` +pub fn wrapping_add( + lhs: &PrimitiveArray, + rhs: &PrimitiveArray, +) -> Result> +where + T: NativeType + WrappingAdd, +{ + check_same_type(lhs, rhs)?; + + let op = move |a: T, b: T| a.wrapping_add(&b); + + binary(lhs, rhs, lhs.data_type().clone(), op) +} + /// Checked addition of two primitive arrays. If the result from the sum /// overflows, the validity for that index is changed to None /// @@ -138,6 +167,17 @@ where } } +impl ArrayWrappingAdd> for PrimitiveArray +where + T: NativeType + WrappingAdd + NotI128, +{ + type Output = Self; + + fn wrapping_add(&self, rhs: &PrimitiveArray) -> Result { + wrapping_add(self, rhs) + } +} + // Implementation of ArrayCheckedAdd trait for PrimitiveArrays impl ArrayCheckedAdd> for PrimitiveArray where @@ -195,6 +235,26 @@ where unary(lhs, |a| a + rhs, lhs.data_type().clone()) } +/// Wrapping addition of a scalar T to a [`PrimitiveArray`] of type T. +/// It do nothing if the result overflows. +/// +/// # Examples +/// ``` +/// use arrow2::compute::arithmetics::basic::wrapping_add_scalar; +/// use arrow2::array::Int8Array; +/// +/// let a = Int8Array::from(&[None, Some(100)]); +/// let result = wrapping_add_scalar(&a, &100i8); +/// let expected = Int8Array::from(&[None, Some(-56)]); +/// assert_eq!(result, expected); +/// ``` +pub fn wrapping_add_scalar(lhs: &PrimitiveArray, rhs: &T) -> PrimitiveArray +where + T: NativeType + WrappingAdd, +{ + unary(lhs, |a| a.wrapping_add(rhs), lhs.data_type().clone()) +} + /// Checked addition of a scalar T to a primitive array of type T. If the /// result from the sum overflows then the validity index for that value is /// changed to None diff --git a/src/compute/arithmetics/basic/mul.rs b/src/compute/arithmetics/basic/mul.rs index 9b6e27ee1d4..101442947d3 100644 --- a/src/compute/arithmetics/basic/mul.rs +++ b/src/compute/arithmetics/basic/mul.rs @@ -1,9 +1,10 @@ //! Definition of basic mul operations with primitive arrays use std::ops::Mul; -use num_traits::{ops::overflowing::OverflowingMul, CheckedMul, SaturatingMul, Zero}; +use num_traits::{ops::overflowing::OverflowingMul, CheckedMul, SaturatingMul, WrappingMul, Zero}; use crate::compute::arithmetics::basic::check_same_type; +use crate::compute::arithmetics::ArrayWrappingMul; use crate::{ array::{Array, PrimitiveArray}, bitmap::Bitmap, @@ -42,6 +43,34 @@ where binary(lhs, rhs, lhs.data_type().clone(), |a, b| a * b) } +/// Wrapping multiplication of two [`PrimitiveArray`]s. +/// It wraps around at the boundary of the type if the result overflows. +/// +/// # Examples +/// ``` +/// use arrow2::compute::arithmetics::basic::wrapping_mul; +/// use arrow2::array::PrimitiveArray; +/// +/// let a = PrimitiveArray::from([Some(100i8), Some(0x10i8), Some(100i8)]); +/// let b = PrimitiveArray::from([Some(0i8), Some(0x10i8), Some(0i8)]); +/// let result = wrapping_mul(&a, &b).unwrap(); +/// let expected = PrimitiveArray::from([Some(0), Some(0), Some(0)]); +/// assert_eq!(result, expected); +/// ``` +pub fn wrapping_mul( + lhs: &PrimitiveArray, + rhs: &PrimitiveArray, +) -> Result> +where + T: NativeType + WrappingMul, +{ + check_same_type(lhs, rhs)?; + + let op = move |a: T, b: T| a.wrapping_mul(&b); + + binary(lhs, rhs, lhs.data_type().clone(), op) +} + /// Checked multiplication of two primitive arrays. If the result from the /// multiplications overflows, the validity for that index is changed /// returned. @@ -139,6 +168,17 @@ where } } +impl ArrayWrappingMul> for PrimitiveArray +where + T: NativeType + WrappingMul + NotI128, +{ + type Output = Self; + + fn wrapping_mul(&self, rhs: &PrimitiveArray) -> Result { + wrapping_mul(self, rhs) + } +} + // Implementation of ArrayCheckedMul trait for PrimitiveArrays impl ArrayCheckedMul> for PrimitiveArray where @@ -195,6 +235,26 @@ where unary(lhs, |a| a * rhs, lhs.data_type().clone()) } +/// Wrapping multiplication of a scalar T to a [`PrimitiveArray`] of type T. +/// It do nothing if the result overflows. +/// +/// # Examples +/// ``` +/// use arrow2::compute::arithmetics::basic::wrapping_mul_scalar; +/// use arrow2::array::Int8Array; +/// +/// let a = Int8Array::from(&[None, Some(0x10)]); +/// let result = wrapping_mul_scalar(&a, &0x10); +/// let expected = Int8Array::from(&[None, Some(0)]); +/// assert_eq!(result, expected); +/// ``` +pub fn wrapping_mul_scalar(lhs: &PrimitiveArray, rhs: &T) -> PrimitiveArray +where + T: NativeType + WrappingMul, +{ + unary(lhs, |a| a.wrapping_mul(rhs), lhs.data_type().clone()) +} + /// Checked multiplication of a scalar T to a primitive array of type T. If the /// result from the multiplication overflows, then the validity for that index is /// changed to None diff --git a/src/compute/arithmetics/basic/sub.rs b/src/compute/arithmetics/basic/sub.rs index d69f2b04f4c..eaefed35f1b 100644 --- a/src/compute/arithmetics/basic/sub.rs +++ b/src/compute/arithmetics/basic/sub.rs @@ -1,9 +1,10 @@ //! Definition of basic sub operations with primitive arrays use std::ops::Sub; -use num_traits::{ops::overflowing::OverflowingSub, CheckedSub, SaturatingSub, Zero}; +use num_traits::{ops::overflowing::OverflowingSub, CheckedSub, SaturatingSub, WrappingSub, Zero}; use crate::compute::arithmetics::basic::check_same_type; +use crate::compute::arithmetics::ArrayWrappingSub; use crate::{ array::{Array, PrimitiveArray}, bitmap::Bitmap, @@ -42,6 +43,34 @@ where binary(lhs, rhs, lhs.data_type().clone(), |a, b| a - b) } +/// Wrapping subtraction of two [`PrimitiveArray`]s. +/// It wraps around at the boundary of the type if the result overflows. +/// +/// # Examples +/// ``` +/// use arrow2::compute::arithmetics::basic::wrapping_sub; +/// use arrow2::array::PrimitiveArray; +/// +/// let a = PrimitiveArray::from([Some(-100i8), Some(-100i8), Some(100i8)]); +/// let b = PrimitiveArray::from([Some(0i8), Some(100i8), Some(0i8)]); +/// let result = wrapping_sub(&a, &b).unwrap(); +/// let expected = PrimitiveArray::from([Some(-100i8), Some(56i8), Some(100i8)]); +/// assert_eq!(result, expected); +/// ``` +pub fn wrapping_sub( + lhs: &PrimitiveArray, + rhs: &PrimitiveArray, +) -> Result> +where + T: NativeType + WrappingSub, +{ + check_same_type(lhs, rhs)?; + + let op = move |a: T, b: T| a.wrapping_sub(&b); + + binary(lhs, rhs, lhs.data_type().clone(), op) +} + /// Checked subtraction of two primitive arrays. If the result from the /// subtraction overflow, the validity for that index is changed /// @@ -138,6 +167,17 @@ where } } +impl ArrayWrappingSub> for PrimitiveArray +where + T: NativeType + WrappingSub + NotI128, +{ + type Output = Self; + + fn wrapping_sub(&self, rhs: &PrimitiveArray) -> Result { + wrapping_sub(self, rhs) + } +} + // Implementation of ArrayCheckedSub trait for PrimitiveArrays impl ArrayCheckedSub> for PrimitiveArray where @@ -195,6 +235,26 @@ where unary(lhs, |a| a - rhs, lhs.data_type().clone()) } +/// Wrapping subtraction of a scalar T to a [`PrimitiveArray`] of type T. +/// It do nothing if the result overflows. +/// +/// # Examples +/// ``` +/// use arrow2::compute::arithmetics::basic::wrapping_sub_scalar; +/// use arrow2::array::Int8Array; +/// +/// let a = Int8Array::from(&[None, Some(-100)]); +/// let result = wrapping_sub_scalar(&a, &100i8); +/// let expected = Int8Array::from(&[None, Some(56)]); +/// assert_eq!(result, expected); +/// ``` +pub fn wrapping_sub_scalar(lhs: &PrimitiveArray, rhs: &T) -> PrimitiveArray +where + T: NativeType + WrappingSub, +{ + unary(lhs, |a| a.wrapping_sub(rhs), lhs.data_type().clone()) +} + /// Checked subtraction of a scalar T to a primitive array of type T. If the /// result from the subtraction overflows, then the validity for that index /// is changed to None diff --git a/src/compute/arithmetics/mod.rs b/src/compute/arithmetics/mod.rs index 99378716cab..e53fd96a191 100644 --- a/src/compute/arithmetics/mod.rs +++ b/src/compute/arithmetics/mod.rs @@ -309,6 +309,13 @@ pub trait ArrayAdd { fn add(&self, rhs: &Rhs) -> Result; } +/// Defines wrapping addition operation for primitive arrays +pub trait ArrayWrappingAdd { + type Output; + + fn wrapping_add(&self, rhs: &Rhs) -> Result; +} + /// Defines checked addition operation for primitive arrays pub trait ArrayCheckedAdd { type Output; @@ -337,6 +344,13 @@ pub trait ArraySub { fn sub(&self, rhs: &Rhs) -> Result; } +/// Defines wrapping subtraction operation for primitive arrays +pub trait ArrayWrappingSub { + type Output; + + fn wrapping_sub(&self, rhs: &Rhs) -> Result; +} + /// Defines checked subtraction operation for primitive arrays pub trait ArrayCheckedSub { type Output; @@ -365,6 +379,13 @@ pub trait ArrayMul { fn mul(&self, rhs: &Rhs) -> Result; } +/// Defines wrapping multiplication operation for primitive arrays +pub trait ArrayWrappingMul { + type Output; + + fn wrapping_mul(&self, rhs: &Rhs) -> Result; +} + /// Defines checked multiplication operation for primitive arrays pub trait ArrayCheckedMul { type Output;