-
Notifications
You must be signed in to change notification settings - Fork 658
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.
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 and operation in the kernel #2703
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,109 @@ | ||
// Licensed to the Apache Software Foundation (ASF) under one | ||
// or more contributor license agreements. See the NOTICE file | ||
// distributed with this work for additional information | ||
// regarding copyright ownership. The ASF licenses this file | ||
// to you under the Apache License, Version 2.0 (the | ||
// "License"); you may not use this file except in compliance | ||
// with the License. You may obtain a copy of the License at | ||
// | ||
// http://www.apache.org/licenses/LICENSE-2.0 | ||
// | ||
// Unless required by applicable law or agreed to in writing, | ||
// software distributed under the License is distributed on an | ||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
// KIND, either express or implied. See the License for the | ||
// specific language governing permissions and limitations | ||
// under the License. | ||
|
||
use crate::array::PrimitiveArray; | ||
use crate::compute::{binary, unary}; | ||
use crate::datatypes::ArrowNumericType; | ||
use crate::error::{ArrowError, Result}; | ||
use std::ops::BitAnd; | ||
|
||
// The helper function for bitwise operation with two array | ||
fn bitwise_op<T, F>( | ||
left: &PrimitiveArray<T>, | ||
right: &PrimitiveArray<T>, | ||
op: F, | ||
) -> Result<PrimitiveArray<T>> | ||
where | ||
T: ArrowNumericType, | ||
F: Fn(T::Native, T::Native) -> T::Native, | ||
{ | ||
if left.len() != right.len() { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We also check the length in Maybe we could let There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think it is fine, the compiler should eliminate it, and even if it doesn't the cost will be amortised over the much more expensive code below There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
If the binary return the |
||
return Err(ArrowError::ComputeError( | ||
"Cannot perform bitwise operation on arrays of different length".to_string(), | ||
)); | ||
} | ||
Ok(binary(left, right, op)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could we directly operate on the value buffer, There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You would need to take into account both endianess, and the size of the types, but theoretically yes you could do this. I suspect it would perform worse, however, as you'll have a variable length type now There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
@HaoYang670 we can track a issue for this optimization, and can try it after finish the bitwise kernel There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agreed @liukun4515. And we should add some benchmarks to see whether we ccoul get performance improvement by doing this. |
||
} | ||
|
||
/// Perform `left & right` operation on two arrays. If either left or right value is null | ||
/// then the result is also null. | ||
pub fn bitwise_and<T>( | ||
left: &PrimitiveArray<T>, | ||
right: &PrimitiveArray<T>, | ||
) -> Result<PrimitiveArray<T>> | ||
where | ||
T: ArrowNumericType, | ||
T::Native: BitAnd<Output = T::Native>, | ||
{ | ||
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 | ||
/// result is also null. | ||
pub fn bitwise_and_scalar<T>( | ||
array: &PrimitiveArray<T>, | ||
scalar: T::Native, | ||
) -> Result<PrimitiveArray<T>> | ||
where | ||
T: ArrowNumericType, | ||
T::Native: BitAnd<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::error::Result; | ||
|
||
#[test] | ||
fn test_bitwise_and_array() -> Result<()> { | ||
// unsigned value | ||
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]); | ||
let right = UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12)]); | ||
let expected = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]); | ||
let result = bitwise_and(&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(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); | ||
Ok(()) | ||
} | ||
|
||
#[test] | ||
fn test_bitwise_and_array_scalar() -> Result<()> { | ||
// unsigned value | ||
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]); | ||
let scalar = 7; | ||
let expected = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could we get a test where the result is changed 😄 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done |
||
let result = bitwise_and_scalar(&left, scalar)?; | ||
assert_eq!(expected, result); | ||
|
||
// signed value | ||
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]); | ||
let scalar = 7; | ||
let expected = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]); | ||
let result = bitwise_and_scalar(&left, scalar)?; | ||
assert_eq!(expected, result); | ||
Ok(()) | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What is the difference between this function and
math_op
(https://github.com/apache/arrow-rs/blob/master/arrow/src/compute/kernels/arithmetic.rs#L60)?I find the 2 have same types and function body. Could we directly reuse
math_op
, such as?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
two difference about the
math_op
T
, and the output data type is alsoT
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We can use the
math_op
to replace the code body, but I concern the error message which can point out what operation is error