diff --git a/arrow-arith/src/bitwise.rs b/arrow-arith/src/bitwise.rs index 08cc246b351a..f9f456bf95fc 100644 --- a/arrow-arith/src/bitwise.rs +++ b/arrow-arith/src/bitwise.rs @@ -17,7 +17,9 @@ use crate::arity::{binary, unary}; use arrow_array::*; +use arrow_buffer::ArrowNativeType; use arrow_schema::ArrowError; +use num::traits::{WrappingShl, WrappingShr}; use std::ops::{BitAnd, BitOr, BitXor, Not}; // The helper function for bitwise operation with two array @@ -121,6 +123,38 @@ where Ok(unary(array, |value| value ^ scalar)) } +/// Perform bitwise 'left << right' operation on two arrays. If either left or right value is null +/// then the result is also null. +pub fn bitwise_shift_left( + left: &PrimitiveArray, + right: &PrimitiveArray, +) -> Result, ArrowError> +where + T: ArrowNumericType, + T::Native: WrappingShl, +{ + bitwise_op(left, right, |a, b| { + let b = b.as_usize(); + a.wrapping_shl(b as u32) + }) +} + +/// Perform bitwise 'left >> right' operation on two arrays. If either left or right value is null +/// then the result is also null. +pub fn bitwise_shift_right( + left: &PrimitiveArray, + right: &PrimitiveArray, +) -> Result, ArrowError> +where + T: ArrowNumericType, + T::Native: WrappingShr, +{ + bitwise_op(left, right, |a, b| { + let b = b.as_usize(); + a.wrapping_shr(b as u32) + }) +} + #[cfg(test)] mod tests { use super::*; @@ -143,6 +177,28 @@ mod tests { Ok(()) } + #[test] + fn test_bitwise_shift_left() { + let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(8)]); + let right = + UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12), Some(u64::MAX)]); + let expected = + UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(0)]); + let result = bitwise_shift_left(&left, &right).unwrap(); + assert_eq!(expected, result); + } + + #[test] + fn test_bitwise_shift_right() { + let left = + UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(3)]); + let right = + UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12), Some(65)]); + let expected = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(1)]); + let result = bitwise_shift_right(&left, &right).unwrap(); + assert_eq!(expected, result); + } + #[test] fn test_bitwise_and_array_scalar() { // unsigned value