Skip to content

Commit ba6dbb8

Browse files
authored
Add support for more fused boolean operations (#5298)
* feat: don't check schema and batch * Add support for more fused boolean operations * Add support for more fused boolean operations * remove git error message in code. * format code * fix the Clippy error.
1 parent 8345991 commit ba6dbb8

File tree

3 files changed

+92
-1
lines changed

3 files changed

+92
-1
lines changed

arrow-arith/src/bitwise.rs

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,20 @@ where
116116
Ok(unary(array, |value| !value))
117117
}
118118

119+
/// Perform `left & !right` operation on two arrays. If either left or right value is null
120+
/// then the result is also null.
121+
pub fn bitwise_and_not<T>(
122+
left: &PrimitiveArray<T>,
123+
right: &PrimitiveArray<T>,
124+
) -> Result<PrimitiveArray<T>, ArrowError>
125+
where
126+
T: ArrowNumericType,
127+
T::Native: BitAnd<Output = T::Native>,
128+
T::Native: Not<Output = T::Native>,
129+
{
130+
bitwise_op(left, right, |a, b| a & !b)
131+
}
132+
119133
/// Perform bitwise `and` every value in an array with the scalar. If any value in the array is null then the
120134
/// result is also null.
121135
pub fn bitwise_and_scalar<T>(
@@ -298,6 +312,31 @@ mod tests {
298312
assert_eq!(expected, result);
299313
}
300314

315+
#[test]
316+
fn test_bitwise_and_not_array() {
317+
// unsigned value
318+
let left = UInt64Array::from(vec![Some(8), Some(2), None, Some(4)]);
319+
let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
320+
let expected = UInt64Array::from(vec![Some(8), Some(2), None, Some(0)]);
321+
let result = bitwise_and_not(&left, &right).unwrap();
322+
assert_eq!(expected, result);
323+
assert_eq!(
324+
bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
325+
result
326+
);
327+
328+
// signed value
329+
let left = Int32Array::from(vec![Some(2), Some(1), None, Some(3)]);
330+
let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
331+
let expected = Int32Array::from(vec![Some(2), Some(0), None, Some(2)]);
332+
let result = bitwise_and_not(&left, &right).unwrap();
333+
assert_eq!(expected, result);
334+
assert_eq!(
335+
bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
336+
result
337+
);
338+
}
339+
301340
#[test]
302341
fn test_bitwise_or_array_scalar() {
303342
// unsigned value

arrow-arith/src/boolean.rs

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
2525
use arrow_array::*;
2626
use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper};
27-
use arrow_buffer::{BooleanBuffer, NullBuffer};
27+
use arrow_buffer::{buffer_bin_and_not, BooleanBuffer, NullBuffer};
2828
use arrow_schema::ArrowError;
2929

3030
/// Logical 'and' boolean values with Kleene logic
@@ -272,6 +272,27 @@ pub fn or(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, Arr
272272
binary_boolean_kernel(left, right, |a, b| a | b)
273273
}
274274

275+
/// Performs `AND_NOT` operation on two arrays. If either left or right value is null then the
276+
/// result is also null.
277+
/// # Error
278+
/// This function errors when the arrays have different lengths.
279+
/// # Example
280+
/// ```rust
281+
/// # use arrow_array::BooleanArray;
282+
/// # use arrow_arith::boolean::{and, not, and_not};
283+
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
284+
/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
285+
/// let andn_ab = and_not(&a, &b).unwrap();
286+
/// assert_eq!(andn_ab, BooleanArray::from(vec![Some(false), Some(false), None]));
287+
/// // It's equal to and(left, not(right))
288+
/// assert_eq!(andn_ab, and(&a, &not(&b).unwrap()).unwrap());
289+
pub fn and_not(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
290+
binary_boolean_kernel(left, right, |a, b| {
291+
let buffer = buffer_bin_and_not(a.inner(), b.offset(), b.inner(), a.offset(), a.len());
292+
BooleanBuffer::new(buffer, left.offset(), left.len())
293+
})
294+
}
295+
275296
/// Performs unary `NOT` operation on an arrays. If value is null then the result is also
276297
/// null.
277298
/// # Error
@@ -356,6 +377,18 @@ mod tests {
356377
assert_eq!(c, expected);
357378
}
358379

380+
#[test]
381+
fn test_bool_array_and_not() {
382+
let a = BooleanArray::from(vec![false, false, true, true]);
383+
let b = BooleanArray::from(vec![false, true, false, true]);
384+
let c = and_not(&a, &b).unwrap();
385+
386+
let expected = BooleanArray::from(vec![false, false, true, false]);
387+
388+
assert_eq!(c, expected);
389+
assert_eq!(c, and(&a, &not(&b).unwrap()).unwrap());
390+
}
391+
359392
#[test]
360393
fn test_bool_array_or_nulls() {
361394
let a = BooleanArray::from(vec![

arrow-buffer/src/buffer/ops.rs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,25 @@ pub fn buffer_bin_xor(
182182
)
183183
}
184184

185+
/// Apply a bitwise and_not to two inputs and return the result as a Buffer.
186+
/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
187+
pub fn buffer_bin_and_not(
188+
left: &Buffer,
189+
left_offset_in_bits: usize,
190+
right: &Buffer,
191+
right_offset_in_bits: usize,
192+
len_in_bits: usize,
193+
) -> Buffer {
194+
bitwise_bin_op_helper(
195+
left,
196+
left_offset_in_bits,
197+
right,
198+
right_offset_in_bits,
199+
len_in_bits,
200+
|a, b| a & !b,
201+
)
202+
}
203+
185204
/// Apply a bitwise not to one input and return the result as a Buffer.
186205
/// The input is treated as a bitmap, meaning that offset and length are specified in number of bits.
187206
pub fn buffer_unary_not(left: &Buffer, offset_in_bits: usize, len_in_bits: usize) -> Buffer {

0 commit comments

Comments
 (0)