Add bitand_many, union_many, and relevant tests.

Author: mbutrovich

arrow-buffer/src/buffer/boolean.rs

@@ -625,6 +625,20 @@ impl BooleanBuffer {
     pub fn set_slices(&self) -> BitSliceIterator<'_> {
         BitSliceIterator::new(self.values(), self.bit_offset, self.bit_len)
     }
+
+    /// Bitwise AND of multiple [`BooleanBuffer`]s
+    ///
+    /// # Panics
+    ///
+    /// Panics if `buffers` is empty, or if the buffers have different lengths
+    pub fn bitand_many(buffers: &[&BooleanBuffer]) -> BooleanBuffer {
+        assert!(!buffers.is_empty(), "requires at least one buffer");
+        let mut result = buffers[0].clone();
+        for buf in &buffers[1..] {
+            result &= *buf;
+        }
+        result
+    }
 }
 
 impl Not for &BooleanBuffer {
@@ -1245,4 +1259,29 @@ mod tests {
         let buffer = BooleanBuffer::new_unset(100);
         assert_eq!(buffer.clone().find_nth_set_bit_position(0, 1), 100);
     }
+
+    #[test]
+    fn test_bitand_many_single() {
+        let a = BooleanBuffer::from(&[true, false, true, true] as &[bool]);
+        let result = BooleanBuffer::bitand_many(&[&a]);
+        assert_eq!(result, a);
+    }
+
+    #[test]
+    fn test_bitand_many_two() {
+        let a = BooleanBuffer::from(&[true, false, true, true] as &[bool]);
+        let b = BooleanBuffer::from(&[true, true, false, true] as &[bool]);
+        let result = BooleanBuffer::bitand_many(&[&a, &b]);
+        assert_eq!(result, &a & &b);
+    }
+
+    #[test]
+    fn test_bitand_many_three() {
+        let a = BooleanBuffer::from(&[true, false, true, true] as &[bool]);
+        let b = BooleanBuffer::from(&[true, true, false, true] as &[bool]);
+        let c = BooleanBuffer::from(&[false, true, true, true] as &[bool]);
+        let result = BooleanBuffer::bitand_many(&[&a, &b, &c]);
+        let expected = BooleanBuffer::from(&[false, false, false, true] as &[bool]);
+        assert_eq!(result, expected);
+    }
 }

arrow-buffer/src/buffer/null.rs

@@ -84,6 +84,21 @@ impl NullBuffer {
         }
     }
 
+    /// Computes the union of the nulls in multiple optional [`NullBuffer`]s
+    ///
+    /// See [`union`](Self::union)
+    pub fn union_many(nulls: &[Option<&NullBuffer>]) -> Option<NullBuffer> {
+        // Unwrap to BooleanBuffer because BitAndAssign is not implemented for NullBuffer
+        let buffers: Vec<&BooleanBuffer> = nulls
+            .iter()
+            .filter_map(|nb| nb.map(NullBuffer::inner))
+            .collect();
+        if buffers.is_empty() {
+            return None;
+        }
+        Some(Self::new(BooleanBuffer::bitand_many(&buffers)))
+    }
+
     /// Returns true if all nulls in `other` also exist in self
     pub fn contains(&self, other: &NullBuffer) -> bool {
         if other.null_count == 0 {
@@ -336,4 +351,51 @@ mod tests {
         let result = NullBuffer::from_unsliced_buffer(buf, 0);
         assert!(result.is_none());
     }
+
+    #[test]
+    fn test_union_many_all_none() {
+        let result = NullBuffer::union_many(&[None, None, None]);
+        assert!(result.is_none());
+    }
+
+    #[test]
+    fn test_union_many_single_some() {
+        let a = NullBuffer::from(&[true, false, true, true]);
+        let result = NullBuffer::union_many(&[Some(&a)]);
+        assert_eq!(result, Some(a));
+    }
+
+    #[test]
+    fn test_union_many_two_inputs() {
+        let a = NullBuffer::from(&[true, false, true, true]);
+        let b = NullBuffer::from(&[true, true, false, true]);
+        let result = NullBuffer::union_many(&[Some(&a), Some(&b)]);
+        let expected = NullBuffer::union(Some(&a), Some(&b));
+        assert_eq!(result, expected);
+    }
+
+    #[test]
+    fn test_union_many_three_inputs() {
+        let a = NullBuffer::from(&[true, false, true, true]);
+        let b = NullBuffer::from(&[true, true, false, true]);
+        let c = NullBuffer::from(&[false, true, true, true]);
+        let result = NullBuffer::union_many(&[Some(&a), Some(&b), Some(&c)]);
+        let expected = NullBuffer::from(&[false, false, false, true]);
+        assert_eq!(result, Some(expected));
+    }
+
+    #[test]
+    fn test_union_many_mixed_none() {
+        let a = NullBuffer::from(&[true, false, true, true]);
+        let b = NullBuffer::from(&[false, true, true, true]);
+        let result = NullBuffer::union_many(&[Some(&a), None, Some(&b)]);
+        let expected = NullBuffer::union(Some(&a), Some(&b));
+        assert_eq!(result, expected);
+    }
+
+    #[test]
+    fn test_union_many_empty_slice() {
+        let result = NullBuffer::union_many(&[]);
+        assert!(result.is_none());
+    }
 }