diff --git a/arrow/src/compute/kernels/bitwise.rs b/arrow/src/compute/kernels/bitwise.rs new file mode 100644 index 000000000000..18b9f4bb760c --- /dev/null +++ b/arrow/src/compute/kernels/bitwise.rs @@ -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( + left: &PrimitiveArray, + right: &PrimitiveArray, + op: F, +) -> Result> +where + T: ArrowNumericType, + F: Fn(T::Native, T::Native) -> T::Native, +{ + if left.len() != right.len() { + return Err(ArrowError::ComputeError( + "Cannot perform bitwise operation on arrays of different length".to_string(), + )); + } + Ok(binary(left, right, op)) +} + +/// 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( + left: &PrimitiveArray, + right: &PrimitiveArray, +) -> Result> +where + T: ArrowNumericType, + T::Native: BitAnd, +{ + 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( + array: &PrimitiveArray, + scalar: T::Native, +) -> Result> +where + T: ArrowNumericType, + T::Native: BitAnd, +{ + 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(15), Some(2), None, Some(4)]); + let scalar = 7; + let expected = UInt64Array::from(vec![Some(7), Some(2), None, Some(4)]); + 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 = 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(()) + } +} diff --git a/arrow/src/compute/kernels/mod.rs b/arrow/src/compute/kernels/mod.rs index c615d3a55e1a..99cdcf460ce1 100644 --- a/arrow/src/compute/kernels/mod.rs +++ b/arrow/src/compute/kernels/mod.rs @@ -20,6 +20,7 @@ pub mod aggregate; pub mod arithmetic; pub mod arity; +pub mod bitwise; pub mod boolean; pub mod cast; pub mod cast_utils;