Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Added checked_negate and wrapping_negate for PrimitiveArray (#506)
Browse files Browse the repository at this point in the history
  • Loading branch information
yjhmelody authored Oct 9, 2021
1 parent b19dde2 commit e77d4e1
Showing 1 changed file with 41 additions and 2 deletions.
43 changes: 41 additions & 2 deletions src/compute/arithmetics/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,14 +61,14 @@ pub mod time;

use std::ops::{Add, Div, Mul, Neg, Rem, Sub};

use num_traits::{NumCast, Zero};
use num_traits::{CheckedNeg, NumCast, WrappingNeg, Zero};

use crate::datatypes::{DataType, IntervalUnit, TimeUnit};
use crate::error::{ArrowError, Result};
use crate::types::NativeType;
use crate::{array::*, bitmap::Bitmap};

use super::arity::unary;
use super::arity::{unary, unary_checked};

// Macro to evaluate match branch in arithmetic function.
// The macro is used to downcast both arrays to a primitive_array_type. If there
Expand Down Expand Up @@ -302,6 +302,45 @@ where
unary(array, |a| -a, array.data_type().clone())
}

/// Checked negates values from array.
///
/// # Examples
/// ```
/// use arrow2::compute::arithmetics::checked_negate;
/// use arrow2::array::{Array, PrimitiveArray};
///
/// let a = PrimitiveArray::from([None, Some(6), Some(i8::MIN), Some(7)]);
/// let result = checked_negate(&a);
/// let expected = PrimitiveArray::from([None, Some(-6), None, Some(-7)]);
/// assert_eq!(result, expected);
/// assert!(!result.is_valid(2))
/// ```
pub fn checked_negate<T>(array: &PrimitiveArray<T>) -> PrimitiveArray<T>
where
T: NativeType + CheckedNeg,
{
unary_checked(array, |a| a.checked_neg(), array.data_type().clone())
}

/// Wrapping negates values from array.
///
/// # Examples
/// ```
/// use arrow2::compute::arithmetics::wrapping_negate;
/// use arrow2::array::{Array, PrimitiveArray};
///
/// let a = PrimitiveArray::from([None, Some(6), Some(i8::MIN), Some(7)]);
/// let result = wrapping_negate(&a);
/// let expected = PrimitiveArray::from([None, Some(-6), Some(i8::MIN), Some(-7)]);
/// assert_eq!(result, expected);
/// ```
pub fn wrapping_negate<T>(array: &PrimitiveArray<T>) -> PrimitiveArray<T>
where
T: NativeType + WrappingNeg,
{
unary(array, |a| a.wrapping_neg(), array.data_type().clone())
}

/// Defines basic addition operation for primitive arrays
pub trait ArrayAdd<Rhs> {
type Output;
Expand Down

0 comments on commit e77d4e1

Please sign in to comment.