Add BooleanArray::from_unary and BooleanArray::from_binary

arrow-array/src/array/boolean_array.rs

@@ -20,8 +20,9 @@ use crate::iterator::BooleanIter;
 use crate::raw_pointer::RawPtrBox;
 use crate::{print_long_array, Array, ArrayAccessor};
 use arrow_buffer::{bit_util, Buffer, MutableBuffer};
+use arrow_data::bit_mask::combine_option_bitmap;
 use arrow_data::ArrayData;
-use arrow_schema::DataType;
+use arrow_schema::{ArrowError, DataType};
 use std::any::Any;
 
 /// Array of bools
@@ -173,6 +174,92 @@ impl BooleanArray {
     ) -> impl Iterator<Item = Option<bool>> + 'a {
         indexes.map(|opt_index| opt_index.map(|index| self.value_unchecked(index)))
     }
+
+    /// Create a [`BooleanArray`] by evaluating the operation for
+    /// each element of the provided array
+    ///
+    /// ```
+    /// # use arrow_array::{BooleanArray, Int32Array};
+    ///
+    /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
+    /// let r = BooleanArray::from_unary(&array, |x| x > 2);
+    /// assert_eq!(&r, &BooleanArray::from(vec![false, false, true, true, true]));
+    /// ```
+    pub fn from_unary<T: ArrayAccessor, F>(left: T, mut op: F) -> Self
+    where
+        F: FnMut(T::Item) -> bool,
+    {
+        let null_bit_buffer = left
+            .data()
+            .null_buffer()
+            .map(|b| b.bit_slice(left.offset(), left.len()));
+
+        let buffer = MutableBuffer::collect_bool(left.len(), |i| unsafe {
+            // SAFETY: i in range 0..len
+            op(left.value_unchecked(i))
+        });
+
+        let data = unsafe {
+            ArrayData::new_unchecked(
+                DataType::Boolean,
+                left.len(),
+                None,
+                null_bit_buffer,
+                0,
+                vec![Buffer::from(buffer)],
+                vec![],
+            )
+        };
+        Self::from(data)
+    }
+
+    /// Create a [`BooleanArray`] by evaluating the binary operation for
+    /// each element of the provided arrays
+    ///
+    /// ```
+    /// # use arrow_array::{BooleanArray, Int32Array};
+    ///
+    /// let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
+    /// let b = Int32Array::from(vec![1, 2, 0, 2, 5]);
+    /// let r = BooleanArray::from_binary(&a, &b, |a, b| a == b).unwrap();
+    /// assert_eq!(&r, &BooleanArray::from(vec![true, true, false, false, true]));
+    /// ```
+    pub fn from_binary<T: ArrayAccessor, S: ArrayAccessor, F>(
+        left: T,
+        right: S,
+        mut op: F,
+    ) -> Result<Self, ArrowError>
+    where
+        F: FnMut(T::Item, S::Item) -> bool,
+    {
+        if left.len() != right.len() {
+            return Err(ArrowError::ComputeError(
+                "Cannot perform binary operation on arrays of different length"
+                    .to_string(),
+            ));
+        }
+
+        let null_bit_buffer =
+            combine_option_bitmap(&[left.data_ref(), right.data_ref()], left.len());
+
+        let buffer = MutableBuffer::collect_bool(left.len(), |i| unsafe {
+            // SAFETY: i in range 0..len
+            op(left.value_unchecked(i), right.value_unchecked(i))
+        });
+
+        let data = unsafe {
+            ArrayData::new_unchecked(
+                DataType::Boolean,
+                left.len(),
+                None,
+                null_bit_buffer,
+                0,
+                vec![Buffer::from(buffer)],
+                vec![],
+            )
+        };
+        Ok(Self::from(data))
+    }
 }
 
 impl Array for BooleanArray {

arrow-data/src/bit_mask.rs

@@ -17,8 +17,11 @@
 
 //! Utils for working with packed bit masks
 
+use crate::ArrayData;
 use arrow_buffer::bit_chunk_iterator::BitChunks;
 use arrow_buffer::bit_util::{ceil, get_bit, set_bit};
+use arrow_buffer::buffer::buffer_bin_and;
+use arrow_buffer::Buffer;
 
 /// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]` to be equal to the
 /// bits in `data` in the range `[offset_read..offset_read+len]`
@@ -62,9 +65,42 @@ pub fn set_bits(
     null_count as usize
 }
 
+/// Combines the null bitmaps of multiple arrays using a bitwise `and` operation.
+///
+/// This function is useful when implementing operations on higher level arrays.
+pub fn combine_option_bitmap(
+    arrays: &[&ArrayData],
+    len_in_bits: usize,
+) -> Option<Buffer> {
+    let (buffer, offset) = arrays
+        .iter()
+        .map(|array| (array.null_buffer().cloned(), array.offset()))
+        .reduce(|acc, buffer_and_offset| match (acc, buffer_and_offset) {
+            ((None, _), (None, _)) => (None, 0),
+            ((Some(buffer), offset), (None, _)) | ((None, _), (Some(buffer), offset)) => {
+                (Some(buffer), offset)
+            }
+            ((Some(buffer_left), offset_left), (Some(buffer_right), offset_right)) => (
+                Some(buffer_bin_and(
+                    &buffer_left,
+                    offset_left,
+                    &buffer_right,
+                    offset_right,
+                    len_in_bits,
+                )),
+                0,
+            ),
+        })?;
+
+    Some(buffer?.bit_slice(offset, len_in_bits))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
+    use arrow_buffer::buffer::buffer_bin_or;
+    use arrow_schema::DataType;
+    use std::sync::Arc;
 
     #[test]
     fn test_set_bits_aligned() {
@@ -187,4 +223,169 @@ mod tests {
         assert_eq!(destination, expected_data);
         assert_eq!(result, expected_null_count);
     }
+
+    /// Compares the null bitmaps of two arrays using a bitwise `or` operation.
+    ///
+    /// This function is useful when implementing operations on higher level arrays.
+    pub(super) fn compare_option_bitmap(
+        left_data: &ArrayData,
+        right_data: &ArrayData,
+        len_in_bits: usize,
+    ) -> Result<Option<Buffer>, ArrowError> {
+        let left_offset_in_bits = left_data.offset();
+        let right_offset_in_bits = right_data.offset();
+
+        let left = left_data.null_buffer();
+        let right = right_data.null_buffer();
+
+        match left {
+            None => match right {
+                None => Ok(None),
+                Some(r) => Ok(Some(r.bit_slice(right_offset_in_bits, len_in_bits))),
+            },
+            Some(l) => match right {
+                None => Ok(Some(l.bit_slice(left_offset_in_bits, len_in_bits))),
+
+                Some(r) => Ok(Some(buffer_bin_or(
+                    l,
+                    left_offset_in_bits,
+                    r,
+                    right_offset_in_bits,
+                    len_in_bits,
+                ))),
+            },
+        }
+    }
+
+    fn make_data_with_null_bit_buffer(
+        len: usize,
+        offset: usize,
+        null_bit_buffer: Option<Buffer>,
+    ) -> Arc<ArrayData> {
+        let buffer = Buffer::from(&vec![11; len + offset]);
+
+        Arc::new(
+            ArrayData::try_new(
+                DataType::UInt8,
+                len,
+                null_bit_buffer,
+                offset,
+                vec![buffer],
+                vec![],
+            )
+            .unwrap(),
+        )
+    }
+
+    #[test]
+    fn test_combine_option_bitmap() {
+        let none_bitmap = make_data_with_null_bit_buffer(8, 0, None);
+        let some_bitmap =
+            make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b01001010])));
+        let inverse_bitmap =
+            make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b10110101])));
+        let some_other_bitmap =
+            make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b11010111])));
+        assert_eq!(None, combine_option_bitmap(&[], 8));
+        assert_eq!(
+            Some(Buffer::from([0b01001010])),
+            combine_option_bitmap(&[&some_bitmap], 8)
+        );
+        assert_eq!(
+            None,
+            combine_option_bitmap(&[&none_bitmap, &none_bitmap], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b01001010])),
+            combine_option_bitmap(&[&some_bitmap, &none_bitmap], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b11010111])),
+            combine_option_bitmap(&[&none_bitmap, &some_other_bitmap], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b01001010])),
+            combine_option_bitmap(&[&some_bitmap, &some_bitmap], 8,)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b0])),
+            combine_option_bitmap(&[&some_bitmap, &inverse_bitmap], 8,)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b01000010])),
+            combine_option_bitmap(&[&some_bitmap, &some_other_bitmap, &none_bitmap], 8,)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b00001001])),
+            combine_option_bitmap(
+                &[
+                    &some_bitmap.slice(3, 5),
+                    &inverse_bitmap.slice(2, 5),
+                    &some_other_bitmap.slice(1, 5)
+                ],
+                5,
+            )
+        );
+    }
+
+    #[test]
+    fn test_combine_option_bitmap_with_offsets() {
+        let none_bitmap = make_data_with_null_bit_buffer(8, 0, None);
+        let bitmap0 =
+            make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b10101010])));
+        let bitmap1 =
+            make_data_with_null_bit_buffer(8, 1, Some(Buffer::from([0b01010100, 0b1])));
+        let bitmap2 =
+            make_data_with_null_bit_buffer(8, 2, Some(Buffer::from([0b10101000, 0b10])));
+        assert_eq!(
+            Some(Buffer::from([0b10101010])),
+            combine_option_bitmap(&[&bitmap1], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b10101010])),
+            combine_option_bitmap(&[&bitmap2], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b10101010])),
+            combine_option_bitmap(&[&bitmap1, &none_bitmap], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b10101010])),
+            combine_option_bitmap(&[&none_bitmap, &bitmap2], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b10101010])),
+            combine_option_bitmap(&[&bitmap0, &bitmap1], 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b10101010])),
+            combine_option_bitmap(&[&bitmap1, &bitmap2], 8)
+        );
+    }
+
+    #[test]
+    fn test_compare_option_bitmap() {
+        let none_bitmap = make_data_with_null_bit_buffer(8, 0, None);
+        let some_bitmap =
+            make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b01001010])));
+        let inverse_bitmap =
+            make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b10110101])));
+        assert_eq!(None, compare_option_bitmap(&none_bitmap, &none_bitmap, 8));
+        assert_eq!(
+            Some(Buffer::from([0b01001010])),
+            compare_option_bitmap(&some_bitmap, &none_bitmap, 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b01001010])),
+            compare_option_bitmap(&none_bitmap, &some_bitmap, 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b01001010])),
+            compare_option_bitmap(&some_bitmap, &some_bitmap, 8)
+        );
+        assert_eq!(
+            Some(Buffer::from([0b11111111])),
+            compare_option_bitmap(&some_bitmap, &inverse_bitmap, 8)
+        );
+    }
 }

arrow/src/compute/kernels/arithmetic.rs

@@ -310,10 +310,10 @@ where
     }
 
     // Create the combined `Bitmap`
-    let null_bit_buffer = crate::compute::util::combine_option_bitmap(
+    let null_bit_buffer = arrow_data::bit_mask::combine_option_bitmap(
         &[left.data_ref(), right.data_ref()],
         left.len(),
-    )?;
+    );
 
     let lanes = T::lanes();
     let buffer_size = left.len() * std::mem::size_of::<T::Native>();
@@ -652,10 +652,10 @@ where
         )));
     }
 
-    let null_bit_buffer = crate::compute::util::combine_option_bitmap(
+    let null_bit_buffer = arrow_data::bit_mask::combine_option_bitmap(
         &[left.data_ref(), right.data_ref()],
         left.len(),
-    )?;
+    );
 
     // Safety justification: Since the inputs are valid Arrow arrays, all values are
     // valid indexes into the dictionary (which is verified during construction)