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

feat: support bitwise shift left/right #4148

Merged
merged 3 commits into from
Apr 28, 2023
Merged

feat: support bitwise shift left/right #4148

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ca30337
feat: support bitwise shift left/right
Weijun-H Apr 27, 2023
3caf1cc
ignore truncation
Weijun-H Apr 27, 2023
6f7a533
add test when shifting by more than the number of bits
Weijun-H Apr 28, 2023
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
56 changes: 56 additions & 0 deletions arrow-arith/src/bitwise.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is great -- when I looked into doing this in DataFusion I noticed these traits weren't in std -- 100% for finding them in num

use std::ops::{BitAnd, BitOr, BitXor, Not};

// The helper function for bitwise operation with two array
Expand Down Expand Up @@ -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<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)
})
}

#[cfg(test)]
mod tests {
use super::*;
Expand All @@ -143,6 +177,28 @@ mod tests {
Ok(())
}

#[test]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it worth adding a test of what happens if you shift by more than the number of bits?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I add a test for it. It should be 0

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are you sure? https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=0bc4c3229b2ffd98a64eaf71573d0199

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes sense, as wrapping_shr would shift the position would be the remainder of the type. So wrapping_shr(3_i32, 32) == 3 and wrapping_shr(3_i32, 33) == 1

Copy link
Member Author

@Weijun-H Weijun-H Apr 28, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shift operations (<<, >>) on a value of width N can be passed a shift value >= N. It is unclear what behaviour should result from this, so the shift value is unconditionally masked to be modulo N to ensure that the argument is always in range.

I found the reason here https://github.com/rust-lang/rfcs/blob/master/text/0560-integer-overflow.md

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
Expand Down