Skip to content

Commit

Permalink
support bitwise op: or,xor,not (#2716)
Browse files Browse the repository at this point in the history
  • Loading branch information
@liukun4515
liukun4515 authored Sep 13, 2022
1 parent f56b573 commit 7e47fa6
Showing 1 changed file with 162 additions and 5 deletions.
167 changes: 162 additions & 5 deletions arrow/src/compute/kernels/bitwise.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ use crate::array::PrimitiveArray;
use crate::compute::{binary, unary};
use crate::datatypes::ArrowNumericType;
use crate::error::{ArrowError, Result};
use std::ops::BitAnd;
use std::ops::{BitAnd, BitOr, BitXor, Not};

// The helper function for bitwise operation with two array
fn bitwise_op<T, F>(
Expand Down Expand Up @@ -52,7 +52,43 @@ where
bitwise_op(left, right, |a, b| a & b)
}

/// Perform bitwise and every value in an array with the scalar. If any value in the array is null then the
/// Perform `left | right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_or<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: BitOr<Output = T::Native>,
{
bitwise_op(left, right, |a, b| a | b)
}

/// Perform `left ^ right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_xor<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: BitXor<Output = T::Native>,
{
bitwise_op(left, right, |a, b| a ^ b)
}

/// Perform `!array` operation on array. If array value is null
/// then the result is also null.
pub fn bitwise_not<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: Not<Output = T::Native>,
{
Ok(unary(array, |value| !value))
}

/// Perform bitwise `and` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_and_scalar<T>(
array: &PrimitiveArray<T>,
Expand All @@ -65,10 +101,39 @@ where
Ok(unary(array, |value| value & scalar))
}

/// Perform bitwise `or` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_or_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: BitOr<Output = T::Native>,
{
Ok(unary(array, |value| value | scalar))
}

/// Perform bitwise `xor` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_xor_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: BitXor<Output = T::Native>,
{
Ok(unary(array, |value| value ^ scalar))
}

#[cfg(test)]
mod tests {
use crate::array::{Int32Array, UInt64Array};
use crate::compute::kernels::bitwise::{bitwise_and, bitwise_and_scalar};
use crate::compute::kernels::bitwise::{
bitwise_and, bitwise_and_scalar, bitwise_not, bitwise_or, bitwise_or_scalar,
bitwise_xor, bitwise_xor_scalar,
};
use crate::error::Result;

#[test]
Expand All @@ -82,7 +147,7 @@ mod tests {

// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = Int32Array::from(vec![Some(5), Some(10), Some(8), Some(12)]);
let right = Int32Array::from(vec![Some(5), Some(-10), Some(8), Some(12)]);
let expected = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let result = bitwise_and(&left, &right)?;
assert_eq!(expected, result);
Expand All @@ -100,10 +165,102 @@ mod tests {

// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let scalar = 20;
let scalar = -20;
let expected = Int32Array::from(vec![Some(0), Some(0), None, Some(4)]);
let result = bitwise_and_scalar(&left, scalar)?;
assert_eq!(expected, result);
Ok(())
}

#[test]
fn test_bitwise_or_array() -> Result<()> {
// unsigned value
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
let expected = UInt64Array::from(vec![Some(7), Some(7), None, Some(13)]);
let result = bitwise_or(&left, &right)?;
assert_eq!(expected, result);

// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
let expected = Int32Array::from(vec![Some(-7), Some(-5), None, Some(13)]);
let result = bitwise_or(&left, &right)?;
assert_eq!(expected, result);
Ok(())
}

#[test]
fn test_bitwise_not_array() -> Result<()> {
// unsigned value
let array = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let expected = UInt64Array::from(vec![
Some(18446744073709551614),
Some(18446744073709551613),
None,
Some(18446744073709551611),
]);
let result = bitwise_not(&array)?;
assert_eq!(expected, result);
// signed value
let array = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let expected = Int32Array::from(vec![Some(-2), Some(-3), None, Some(-5)]);
let result = bitwise_not(&array)?;
assert_eq!(expected, result);
Ok(())
}

#[test]
fn test_bitwise_or_array_scalar() -> Result<()> {
// unsigned value
let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
let scalar = 7;
let expected = UInt64Array::from(vec![Some(15), Some(7), None, Some(7)]);
let result = bitwise_or_scalar(&left, scalar)?;
assert_eq!(expected, result);

// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let scalar = 20;
let expected = Int32Array::from(vec![Some(21), Some(22), None, Some(20)]);
let result = bitwise_or_scalar(&left, scalar)?;
assert_eq!(expected, result);
Ok(())
}

#[test]
fn test_bitwise_xor_array() -> Result<()> {
// unsigned value
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
let expected = UInt64Array::from(vec![Some(6), Some(7), None, Some(9)]);
let result = bitwise_xor(&left, &right)?;
assert_eq!(expected, result);

// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = Int32Array::from(vec![Some(-7), Some(5), Some(8), Some(-13)]);
let expected = Int32Array::from(vec![Some(-8), Some(7), None, Some(-9)]);
let result = bitwise_xor(&left, &right)?;
assert_eq!(expected, result);
Ok(())
}

#[test]
fn test_bitwise_xor_array_scalar() -> Result<()> {
// unsigned value
let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
let scalar = 7;
let expected = UInt64Array::from(vec![Some(8), Some(5), None, Some(3)]);
let result = bitwise_xor_scalar(&left, scalar)?;
assert_eq!(expected, result);

// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let scalar = -20;
let expected = Int32Array::from(vec![Some(-19), Some(-18), None, Some(-24)]);
let result = bitwise_xor_scalar(&left, scalar)?;
assert_eq!(expected, result);
Ok(())
}
}

0 comments on commit 7e47fa6

Please sign in to comment.