Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support bitwise op in kernel: or,xor,not #2716

Merged
merged 1 commit into from
Sep 13, 2022
Merged

Support bitwise op in kernel: or,xor,not #2716

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
167 changes: 162 additions & 5 deletions arrow/src/compute/kernels/bitwise.rs
View file
Open in desktop
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(())
}
}