diff --git a/arrow/src/compute/kernels/bitwise.rs b/arrow/src/compute/kernels/bitwise.rs index 18b9f4bb760c..2f3c9e490f4c 100644 --- a/arrow/src/compute/kernels/bitwise.rs +++ b/arrow/src/compute/kernels/bitwise.rs @@ -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( @@ -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( + left: &PrimitiveArray, + right: &PrimitiveArray, +) -> Result> +where + T: ArrowNumericType, + T::Native: BitOr, +{ + 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( + left: &PrimitiveArray, + right: &PrimitiveArray, +) -> Result> +where + T: ArrowNumericType, + T::Native: BitXor, +{ + 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(array: &PrimitiveArray) -> Result> +where + T: ArrowNumericType, + T::Native: Not, +{ + 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( array: &PrimitiveArray, @@ -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( + array: &PrimitiveArray, + scalar: T::Native, +) -> Result> +where + T: ArrowNumericType, + T::Native: BitOr, +{ + 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( + array: &PrimitiveArray, + scalar: T::Native, +) -> Result> +where + T: ArrowNumericType, + T::Native: BitXor, +{ + 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] @@ -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); @@ -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(()) + } }