You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2023/05/04 06:58:50 UTC
[arrow-rs] branch master updated: feat: support bitwise shift left/right with scalars (#4159)
This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new c4ac758e1 feat: support bitwise shift left/right with scalars (#4159)
c4ac758e1 is described below
commit c4ac758e13cbd76651d16d06e786bdf7c04f902a
Author: Igor Izvekov <iz...@gmail.com>
AuthorDate: Thu May 4 09:58:43 2023 +0300
feat: support bitwise shift left/right with scalars (#4159)
---
arrow-arith/src/bitwise.rs | 91 ++++++++++++++++++++++++++++++++++++----------
1 file changed, 72 insertions(+), 19 deletions(-)
diff --git a/arrow-arith/src/bitwise.rs b/arrow-arith/src/bitwise.rs
index f9f456bf9..a5dec4638 100644
--- a/arrow-arith/src/bitwise.rs
+++ b/arrow-arith/src/bitwise.rs
@@ -22,7 +22,7 @@ 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
+/// The helper function for bitwise operation with two array
fn bitwise_op<T, F>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
@@ -74,6 +74,38 @@ where
bitwise_op(left, right, |a, b| a ^ b)
}
+/// 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<T>(
+ left: &PrimitiveArray<T>,
+ right: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>, ArrowError>
+where
+ T: ArrowNumericType,
+ T::Native: WrappingShl<Output = T::Native>,
+{
+ 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<T>(
+ left: &PrimitiveArray<T>,
+ right: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>, ArrowError>
+where
+ T: ArrowNumericType,
+ T::Native: WrappingShr<Output = T::Native>,
+{
+ bitwise_op(left, right, |a, b| {
+ let b = b.as_usize();
+ a.wrapping_shr(b as u32)
+ })
+}
+
/// 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>, ArrowError>
@@ -123,36 +155,36 @@ 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<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
+/// Perform bitwise `left << right` 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_shift_left_scalar<T>(
+ array: &PrimitiveArray<T>,
+ scalar: T::Native,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: WrappingShl<Output = T::Native>,
{
- bitwise_op(left, right, |a, b| {
- let b = b.as_usize();
- a.wrapping_shl(b as u32)
- })
+ Ok(unary(array, |value| {
+ let scalar = scalar.as_usize();
+ value.wrapping_shl(scalar 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<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
+/// Perform bitwise `left >> right` 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_shift_right_scalar<T>(
+ array: &PrimitiveArray<T>,
+ scalar: T::Native,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: WrappingShr<Output = T::Native>,
{
- bitwise_op(left, right, |a, b| {
- let b = b.as_usize();
- a.wrapping_shr(b as u32)
- })
+ Ok(unary(array, |value| {
+ let scalar = scalar.as_usize();
+ value.wrapping_shr(scalar as u32)
+ }))
}
#[cfg(test)]
@@ -188,6 +220,16 @@ mod tests {
assert_eq!(expected, result);
}
+ #[test]
+ fn test_bitwise_shift_left_scalar() {
+ let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(8)]);
+ let scalar = 2;
+ let expected =
+ UInt64Array::from(vec![Some(4), Some(8), None, Some(16), Some(32)]);
+ let result = bitwise_shift_left_scalar(&left, scalar).unwrap();
+ assert_eq!(expected, result);
+ }
+
#[test]
fn test_bitwise_shift_right() {
let left =
@@ -199,6 +241,17 @@ mod tests {
assert_eq!(expected, result);
}
+ #[test]
+ fn test_bitwise_shift_right_scalar() {
+ let left =
+ UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(3)]);
+ let scalar = 2;
+ let expected =
+ UInt64Array::from(vec![Some(8), Some(512), None, Some(4096), Some(0)]);
+ let result = bitwise_shift_right_scalar(&left, scalar).unwrap();
+ assert_eq!(expected, result);
+ }
+
#[test]
fn test_bitwise_and_array_scalar() {
// unsigned value