use arrow_array::*;
use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper};
use arrow_buffer::{buffer_bin_and_not, BooleanBuffer, NullBuffer};
use arrow_schema::ArrowError;
pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
let left_values = left.values();
let right_values = right.values();
let buffer = match (left.nulls(), right.nulls()) {
(None, None) => None,
(Some(left_null_buffer), None) => {
Some(bitwise_bin_op_helper(
left_null_buffer.buffer(),
left_null_buffer.offset(),
right_values.inner(),
right_values.offset(),
left.len(),
|a, b| a | !b,
))
}
(None, Some(right_null_buffer)) => {
Some(bitwise_bin_op_helper(
right_null_buffer.buffer(),
right_null_buffer.offset(),
left_values.inner(),
left_values.offset(),
left.len(),
|a, b| a | !b,
))
}
(Some(left_null_buffer), Some(right_null_buffer)) => {
Some(bitwise_quaternary_op_helper(
[
left_null_buffer.buffer(),
left_values.inner(),
right_null_buffer.buffer(),
right_values.inner(),
],
[
left_null_buffer.offset(),
left_values.offset(),
right_null_buffer.offset(),
right_values.offset(),
],
left.len(),
|a, b, c, d| (a | (c & !d)) & (c | (a & !b)),
))
}
};
let nulls = buffer.map(|b| NullBuffer::new(BooleanBuffer::new(b, 0, left.len())));
Ok(BooleanArray::new(left_values & right_values, nulls))
}
pub fn or_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
let left_values = left.values();
let right_values = right.values();
let buffer = match (left.nulls(), right.nulls()) {
(None, None) => None,
(Some(left_nulls), None) => {
Some(bitwise_bin_op_helper(
left_nulls.buffer(),
left_nulls.offset(),
right_values.inner(),
right_values.offset(),
left.len(),
|a, b| a | b,
))
}
(None, Some(right_nulls)) => {
Some(bitwise_bin_op_helper(
right_nulls.buffer(),
right_nulls.offset(),
left_values.inner(),
left_values.offset(),
left.len(),
|a, b| a | b,
))
}
(Some(left_nulls), Some(right_nulls)) => {
Some(bitwise_quaternary_op_helper(
[
left_nulls.buffer(),
left_values.inner(),
right_nulls.buffer(),
right_values.inner(),
],
[
left_nulls.offset(),
left_values.offset(),
right_nulls.offset(),
right_values.offset(),
],
left.len(),
|a, b, c, d| (a | (c & d)) & (c | (a & b)),
))
}
};
let nulls = buffer.map(|b| NullBuffer::new(BooleanBuffer::new(b, 0, left.len())));
Ok(BooleanArray::new(left_values | right_values, nulls))
}
pub(crate) fn binary_boolean_kernel<F>(
left: &BooleanArray,
right: &BooleanArray,
op: F,
) -> Result<BooleanArray, ArrowError>
where
F: Fn(&BooleanBuffer, &BooleanBuffer) -> BooleanBuffer,
{
if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
let nulls = NullBuffer::union(left.nulls(), right.nulls());
let values = op(left.values(), right.values());
Ok(BooleanArray::new(values, nulls))
}
pub fn and(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
binary_boolean_kernel(left, right, |a, b| a & b)
}
pub fn or(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
binary_boolean_kernel(left, right, |a, b| a | b)
}
pub fn and_not(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
binary_boolean_kernel(left, right, |a, b| {
let buffer = buffer_bin_and_not(a.inner(), b.offset(), b.inner(), a.offset(), a.len());
BooleanBuffer::new(buffer, left.offset(), left.len())
})
}
pub fn not(left: &BooleanArray) -> Result<BooleanArray, ArrowError> {
let nulls = left.nulls().cloned();
let values = !left.values();
Ok(BooleanArray::new(values, nulls))
}
pub fn is_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
let values = match input.logical_nulls() {
None => BooleanBuffer::new_unset(input.len()),
Some(nulls) => !nulls.inner(),
};
Ok(BooleanArray::new(values, None))
}
pub fn is_not_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
let values = match input.logical_nulls() {
None => BooleanBuffer::new_set(input.len()),
Some(n) => n.inner().clone(),
};
Ok(BooleanArray::new(values, None))
}
#[cfg(test)]
mod tests {
use arrow_buffer::ScalarBuffer;
use arrow_schema::{DataType, Field, UnionFields};
use super::*;
use std::sync::Arc;
#[test]
fn test_bool_array_and() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![false, true, false, true]);
let c = and(&a, &b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![false, true, false, true]);
let c = or(&a, &b).unwrap();
let expected = BooleanArray::from(vec![false, true, true, true]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_and_not() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![false, true, false, true]);
let c = and_not(&a, &b).unwrap();
let expected = BooleanArray::from(vec![false, false, true, false]);
assert_eq!(c, expected);
assert_eq!(c, and(&a, ¬(&b).unwrap()).unwrap());
}
#[test]
fn test_bool_array_or_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = or(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
None,
None,
None,
Some(false),
Some(true),
None,
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_boolean_array_kleene_no_remainder() {
let n = 1024;
let a = BooleanArray::from(vec![true; n]);
let b = BooleanArray::from(vec![None; n]);
let result = or_kleene(&a, &b).unwrap();
assert_eq!(result, a);
}
#[test]
fn test_bool_array_and_kleene_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = and_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
Some(false),
None,
Some(false),
Some(false),
Some(false),
None,
Some(false),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or_kleene_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = or_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
None,
Some(true),
None,
Some(false),
Some(true),
Some(true),
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or_kleene_right_sided_nulls() {
let a = BooleanArray::from(vec![false, false, false, true, true, true]);
assert!(a.nulls().is_none());
let b = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(false),
None,
]);
assert!(b.nulls().is_some());
let c = or_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or_kleene_left_sided_nulls() {
let a = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(false),
None,
]);
assert!(a.nulls().is_some());
let b = BooleanArray::from(vec![false, false, false, true, true, true]);
assert!(b.nulls().is_none());
let c = or_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_not() {
let a = BooleanArray::from(vec![false, true]);
let c = not(&a).unwrap();
let expected = BooleanArray::from(vec![true, false]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_not_sliced() {
let a = BooleanArray::from(vec![None, Some(true), Some(false), None, Some(true)]);
let a = a.slice(1, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = not(a).unwrap();
let expected = BooleanArray::from(vec![Some(false), Some(true), None, Some(false)]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_and_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = and(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
None,
None,
None,
Some(false),
Some(false),
None,
Some(false),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_and_sliced_same_offset() {
let a = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, false, true, true,
]);
let b = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, true, false, true,
]);
let a = a.slice(8, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let b = b.slice(8, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_sliced_same_offset_mod8() {
let a = BooleanArray::from(vec![
false, false, true, true, false, false, false, false, false, false, false, false,
]);
let b = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, true, false, true,
]);
let a = a.slice(0, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let b = b.slice(8, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_sliced_offset1() {
let a = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, false, true, true,
]);
let b = BooleanArray::from(vec![false, true, false, true]);
let a = a.slice(8, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, &b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_sliced_offset2() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, true, false, true,
]);
let b = b.slice(8, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(&a, b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_nulls_offset() {
let a = BooleanArray::from(vec![None, Some(false), Some(true), None, Some(true)]);
let a = a.slice(1, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let b = BooleanArray::from(vec![
None,
None,
Some(true),
Some(false),
Some(true),
Some(true),
]);
let b = b.slice(2, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, b).unwrap();
let expected = BooleanArray::from(vec![Some(false), Some(false), None, Some(true)]);
assert_eq!(expected, c);
}
#[test]
fn test_nonnull_array_is_null() {
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
let res = is_null(a.as_ref()).unwrap();
let expected = BooleanArray::from(vec![false, false, false, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nonnull_array_with_offset_is_null() {
let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
let a = a.slice(8, 4);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, false, false, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nonnull_array_is_not_null() {
let a = Int32Array::from(vec![1, 2, 3, 4]);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, true, true, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nonnull_array_with_offset_is_not_null() {
let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
let a = a.slice(8, 4);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, true, true, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_is_null() {
let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, true, false, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_with_offset_is_null() {
let a = Int32Array::from(vec![
None,
None,
None,
None,
None,
None,
None,
None,
Some(1),
None,
Some(2),
None,
Some(3),
Some(4),
None,
None,
]);
let a = a.slice(8, 4);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, true, false, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_is_not_null() {
let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, false, true, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_with_offset_is_not_null() {
let a = Int32Array::from(vec![
None,
None,
None,
None,
None,
None,
None,
None,
Some(1),
None,
Some(2),
None,
Some(3),
Some(4),
None,
None,
]);
let a = a.slice(8, 4);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, false, true, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_null_array_is_null() {
let a = NullArray::new(3);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, true, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_null_array_is_not_null() {
let a = NullArray::new(3);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, false, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_dense_union_is_null() {
let int_array = Int32Array::from(vec![Some(1), None]);
let float_array = Float64Array::from(vec![Some(3.2), None]);
let str_array = StringArray::from(vec![Some("a"), None]);
let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
let offsets = [0, 1, 0, 1, 0, 1]
.into_iter()
.collect::<ScalarBuffer<i32>>();
let children = vec![
Arc::new(int_array) as Arc<dyn Array>,
Arc::new(float_array),
Arc::new(str_array),
];
let array = UnionArray::try_new(union_fields(), type_ids, Some(offsets), children).unwrap();
let result = is_null(&array).unwrap();
let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
assert_eq!(expected, &result);
}
#[test]
fn test_sparse_union_is_null() {
let int_array = Int32Array::from(vec![Some(1), None, None, None, None, None]);
let float_array = Float64Array::from(vec![None, None, Some(3.2), None, None, None]);
let str_array = StringArray::from(vec![None, None, None, None, Some("a"), None]);
let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
let children = vec![
Arc::new(int_array) as Arc<dyn Array>,
Arc::new(float_array),
Arc::new(str_array),
];
let array = UnionArray::try_new(union_fields(), type_ids, None, children).unwrap();
let result = is_null(&array).unwrap();
let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
assert_eq!(expected, &result);
}
fn union_fields() -> UnionFields {
[
(0, Arc::new(Field::new("A", DataType::Int32, true))),
(1, Arc::new(Field::new("B", DataType::Float64, true))),
(2, Arc::new(Field::new("C", DataType::Utf8, true))),
]
.into_iter()
.collect()
}
}