arrow_arith/
bitwise.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Module contains bitwise operations on arrays
19
20use crate::arity::{binary, unary};
21use arrow_array::*;
22use arrow_buffer::ArrowNativeType;
23use arrow_schema::ArrowError;
24use num::traits::{WrappingShl, WrappingShr};
25use std::ops::{BitAnd, BitOr, BitXor, Not};
26
27/// The helper function for bitwise operation with two array
28fn bitwise_op<T, F>(
29    left: &PrimitiveArray<T>,
30    right: &PrimitiveArray<T>,
31    op: F,
32) -> Result<PrimitiveArray<T>, ArrowError>
33where
34    T: ArrowNumericType,
35    F: Fn(T::Native, T::Native) -> T::Native,
36{
37    binary(left, right, op)
38}
39
40/// Perform `left & right` operation on two arrays. If either left or right value is null
41/// then the result is also null.
42pub fn bitwise_and<T>(
43    left: &PrimitiveArray<T>,
44    right: &PrimitiveArray<T>,
45) -> Result<PrimitiveArray<T>, ArrowError>
46where
47    T: ArrowNumericType,
48    T::Native: BitAnd<Output = T::Native>,
49{
50    bitwise_op(left, right, |a, b| a & b)
51}
52
53/// Perform `left | right` operation on two arrays. If either left or right value is null
54/// then the result is also null.
55pub fn bitwise_or<T>(
56    left: &PrimitiveArray<T>,
57    right: &PrimitiveArray<T>,
58) -> Result<PrimitiveArray<T>, ArrowError>
59where
60    T: ArrowNumericType,
61    T::Native: BitOr<Output = T::Native>,
62{
63    bitwise_op(left, right, |a, b| a | b)
64}
65
66/// Perform `left ^ right` operation on two arrays. If either left or right value is null
67/// then the result is also null.
68pub fn bitwise_xor<T>(
69    left: &PrimitiveArray<T>,
70    right: &PrimitiveArray<T>,
71) -> Result<PrimitiveArray<T>, ArrowError>
72where
73    T: ArrowNumericType,
74    T::Native: BitXor<Output = T::Native>,
75{
76    bitwise_op(left, right, |a, b| a ^ b)
77}
78
79/// Perform bitwise `left << right` operation on two arrays. If either left or right value is null
80/// then the result is also null.
81pub fn bitwise_shift_left<T>(
82    left: &PrimitiveArray<T>,
83    right: &PrimitiveArray<T>,
84) -> Result<PrimitiveArray<T>, ArrowError>
85where
86    T: ArrowNumericType,
87    T::Native: WrappingShl<Output = T::Native>,
88{
89    bitwise_op(left, right, |a, b| {
90        let b = b.as_usize();
91        a.wrapping_shl(b as u32)
92    })
93}
94
95/// Perform bitwise `left >> right` operation on two arrays. If either left or right value is null
96/// then the result is also null.
97pub fn bitwise_shift_right<T>(
98    left: &PrimitiveArray<T>,
99    right: &PrimitiveArray<T>,
100) -> Result<PrimitiveArray<T>, ArrowError>
101where
102    T: ArrowNumericType,
103    T::Native: WrappingShr<Output = T::Native>,
104{
105    bitwise_op(left, right, |a, b| {
106        let b = b.as_usize();
107        a.wrapping_shr(b as u32)
108    })
109}
110
111/// Perform `!array` operation on array. If array value is null
112/// then the result is also null.
113pub fn bitwise_not<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>, ArrowError>
114where
115    T: ArrowNumericType,
116    T::Native: Not<Output = T::Native>,
117{
118    Ok(unary(array, |value| !value))
119}
120
121/// Perform `left & !right` operation on two arrays. If either left or right value is null
122/// then the result is also null.
123pub fn bitwise_and_not<T>(
124    left: &PrimitiveArray<T>,
125    right: &PrimitiveArray<T>,
126) -> Result<PrimitiveArray<T>, ArrowError>
127where
128    T: ArrowNumericType,
129    T::Native: BitAnd<Output = T::Native>,
130    T::Native: Not<Output = T::Native>,
131{
132    bitwise_op(left, right, |a, b| a & !b)
133}
134
135/// Perform bitwise `and` every value in an array with the scalar. If any value in the array is null then the
136/// result is also null.
137pub fn bitwise_and_scalar<T>(
138    array: &PrimitiveArray<T>,
139    scalar: T::Native,
140) -> Result<PrimitiveArray<T>, ArrowError>
141where
142    T: ArrowNumericType,
143    T::Native: BitAnd<Output = T::Native>,
144{
145    Ok(unary(array, |value| value & scalar))
146}
147
148/// Perform bitwise `or` every value in an array with the scalar. If any value in the array is null then the
149/// result is also null.
150pub fn bitwise_or_scalar<T>(
151    array: &PrimitiveArray<T>,
152    scalar: T::Native,
153) -> Result<PrimitiveArray<T>, ArrowError>
154where
155    T: ArrowNumericType,
156    T::Native: BitOr<Output = T::Native>,
157{
158    Ok(unary(array, |value| value | scalar))
159}
160
161/// Perform bitwise `xor` every value in an array with the scalar. If any value in the array is null then the
162/// result is also null.
163pub fn bitwise_xor_scalar<T>(
164    array: &PrimitiveArray<T>,
165    scalar: T::Native,
166) -> Result<PrimitiveArray<T>, ArrowError>
167where
168    T: ArrowNumericType,
169    T::Native: BitXor<Output = T::Native>,
170{
171    Ok(unary(array, |value| value ^ scalar))
172}
173
174/// Perform bitwise `left << right` every value in an array with the scalar. If any value in the array is null then the
175/// result is also null.
176pub fn bitwise_shift_left_scalar<T>(
177    array: &PrimitiveArray<T>,
178    scalar: T::Native,
179) -> Result<PrimitiveArray<T>, ArrowError>
180where
181    T: ArrowNumericType,
182    T::Native: WrappingShl<Output = T::Native>,
183{
184    Ok(unary(array, |value| {
185        let scalar = scalar.as_usize();
186        value.wrapping_shl(scalar as u32)
187    }))
188}
189
190/// Perform bitwise `left >> right` every value in an array with the scalar. If any value in the array is null then the
191/// result is also null.
192pub fn bitwise_shift_right_scalar<T>(
193    array: &PrimitiveArray<T>,
194    scalar: T::Native,
195) -> Result<PrimitiveArray<T>, ArrowError>
196where
197    T: ArrowNumericType,
198    T::Native: WrappingShr<Output = T::Native>,
199{
200    Ok(unary(array, |value| {
201        let scalar = scalar.as_usize();
202        value.wrapping_shr(scalar as u32)
203    }))
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    #[test]
211    fn test_bitwise_and_array() -> Result<(), ArrowError> {
212        // unsigned value
213        let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
214        let right = UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12)]);
215        let expected = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
216        let result = bitwise_and(&left, &right)?;
217        assert_eq!(expected, result);
218
219        // signed value
220        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
221        let right = Int32Array::from(vec![Some(5), Some(-10), Some(8), Some(12)]);
222        let expected = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
223        let result = bitwise_and(&left, &right)?;
224        assert_eq!(expected, result);
225        Ok(())
226    }
227
228    #[test]
229    fn test_bitwise_shift_left() {
230        let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(8)]);
231        let right = UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12), Some(u64::MAX)]);
232        let expected = UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(0)]);
233        let result = bitwise_shift_left(&left, &right).unwrap();
234        assert_eq!(expected, result);
235    }
236
237    #[test]
238    fn test_bitwise_shift_left_scalar() {
239        let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(8)]);
240        let scalar = 2;
241        let expected = UInt64Array::from(vec![Some(4), Some(8), None, Some(16), Some(32)]);
242        let result = bitwise_shift_left_scalar(&left, scalar).unwrap();
243        assert_eq!(expected, result);
244    }
245
246    #[test]
247    fn test_bitwise_shift_right() {
248        let left = UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(3)]);
249        let right = UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12), Some(65)]);
250        let expected = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(1)]);
251        let result = bitwise_shift_right(&left, &right).unwrap();
252        assert_eq!(expected, result);
253    }
254
255    #[test]
256    fn test_bitwise_shift_right_scalar() {
257        let left = UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(3)]);
258        let scalar = 2;
259        let expected = UInt64Array::from(vec![Some(8), Some(512), None, Some(4096), Some(0)]);
260        let result = bitwise_shift_right_scalar(&left, scalar).unwrap();
261        assert_eq!(expected, result);
262    }
263
264    #[test]
265    fn test_bitwise_and_array_scalar() {
266        // unsigned value
267        let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
268        let scalar = 7;
269        let expected = UInt64Array::from(vec![Some(7), Some(2), None, Some(4)]);
270        let result = bitwise_and_scalar(&left, scalar).unwrap();
271        assert_eq!(expected, result);
272
273        // signed value
274        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
275        let scalar = -20;
276        let expected = Int32Array::from(vec![Some(0), Some(0), None, Some(4)]);
277        let result = bitwise_and_scalar(&left, scalar).unwrap();
278        assert_eq!(expected, result);
279    }
280
281    #[test]
282    fn test_bitwise_or_array() {
283        // unsigned value
284        let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
285        let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
286        let expected = UInt64Array::from(vec![Some(7), Some(7), None, Some(13)]);
287        let result = bitwise_or(&left, &right).unwrap();
288        assert_eq!(expected, result);
289
290        // signed value
291        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
292        let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
293        let expected = Int32Array::from(vec![Some(-7), Some(-5), None, Some(13)]);
294        let result = bitwise_or(&left, &right).unwrap();
295        assert_eq!(expected, result);
296    }
297
298    #[test]
299    fn test_bitwise_not_array() {
300        // unsigned value
301        let array = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
302        let expected = UInt64Array::from(vec![
303            Some(18446744073709551614),
304            Some(18446744073709551613),
305            None,
306            Some(18446744073709551611),
307        ]);
308        let result = bitwise_not(&array).unwrap();
309        assert_eq!(expected, result);
310        // signed value
311        let array = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
312        let expected = Int32Array::from(vec![Some(-2), Some(-3), None, Some(-5)]);
313        let result = bitwise_not(&array).unwrap();
314        assert_eq!(expected, result);
315    }
316
317    #[test]
318    fn test_bitwise_and_not_array() {
319        // unsigned value
320        let left = UInt64Array::from(vec![Some(8), Some(2), None, Some(4)]);
321        let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
322        let expected = UInt64Array::from(vec![Some(8), Some(2), None, Some(0)]);
323        let result = bitwise_and_not(&left, &right).unwrap();
324        assert_eq!(expected, result);
325        assert_eq!(
326            bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
327            result
328        );
329
330        // signed value
331        let left = Int32Array::from(vec![Some(2), Some(1), None, Some(3)]);
332        let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
333        let expected = Int32Array::from(vec![Some(2), Some(0), None, Some(2)]);
334        let result = bitwise_and_not(&left, &right).unwrap();
335        assert_eq!(expected, result);
336        assert_eq!(
337            bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
338            result
339        );
340    }
341
342    #[test]
343    fn test_bitwise_or_array_scalar() {
344        // unsigned value
345        let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
346        let scalar = 7;
347        let expected = UInt64Array::from(vec![Some(15), Some(7), None, Some(7)]);
348        let result = bitwise_or_scalar(&left, scalar).unwrap();
349        assert_eq!(expected, result);
350
351        // signed value
352        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
353        let scalar = 20;
354        let expected = Int32Array::from(vec![Some(21), Some(22), None, Some(20)]);
355        let result = bitwise_or_scalar(&left, scalar).unwrap();
356        assert_eq!(expected, result);
357    }
358
359    #[test]
360    fn test_bitwise_xor_array() {
361        // unsigned value
362        let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
363        let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
364        let expected = UInt64Array::from(vec![Some(6), Some(7), None, Some(9)]);
365        let result = bitwise_xor(&left, &right).unwrap();
366        assert_eq!(expected, result);
367
368        // signed value
369        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
370        let right = Int32Array::from(vec![Some(-7), Some(5), Some(8), Some(-13)]);
371        let expected = Int32Array::from(vec![Some(-8), Some(7), None, Some(-9)]);
372        let result = bitwise_xor(&left, &right).unwrap();
373        assert_eq!(expected, result);
374    }
375
376    #[test]
377    fn test_bitwise_xor_array_scalar() {
378        // unsigned value
379        let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
380        let scalar = 7;
381        let expected = UInt64Array::from(vec![Some(8), Some(5), None, Some(3)]);
382        let result = bitwise_xor_scalar(&left, scalar).unwrap();
383        assert_eq!(expected, result);
384
385        // signed value
386        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
387        let scalar = -20;
388        let expected = Int32Array::from(vec![Some(-19), Some(-18), None, Some(-24)]);
389        let result = bitwise_xor_scalar(&left, scalar).unwrap();
390        assert_eq!(expected, result);
391    }
392}
arrow_arith/bitwise.rs

arrow_arith/
bitwise.rs
