Add BooleanBuffer::try_into_builder and BooleanBuffer::into_builder_or_clone

arrow-array/src/array/boolean_array.rs

@@ -593,8 +593,19 @@ impl BooleanArray {
             return self;
         };
 
-        let mut builder = BooleanBufferBuilder::new(len);
-        builder.append_buffer(&self.values.slice(0, end));
+        let mut builder = match self.values.try_into_builder() {
+            Ok(mut builder) => {
+                // reuse allocation, just resize buffer
+                builder.truncate(end);
+                builder
+            }
+            Err(values) => {
+                // copy only retained prefix
+                let mut builder = BooleanBufferBuilder::new(len);
+                builder.append_buffer(&values.slice(0, end));
+                builder
+            }
+        };
         builder.append_n(len - end, false);
         BooleanArray::new(builder.finish(), self.nulls)
     }
@@ -1654,6 +1665,24 @@ mod tests {
         assert_eq!(r.true_count(), 0);
     }
 
+    #[test]
+    fn test_take_n_true_unshared_reuses_allocation() {
+        let a = BooleanArray::from(vec![true, false, true, true]);
+        let ptr = a.values().values().as_ptr();
+        let r = a.take_n_true(1);
+        assert_eq!(r, BooleanArray::from(vec![true, false, false, false]));
+        // the values buffer was not shared so its allocation is reused
+        assert_eq!(r.values().values().as_ptr(), ptr);
+    }
+
+    #[test]
+    fn test_take_n_true_sliced() {
+        let a = BooleanArray::from(vec![false, true, true, true, false, true]);
+        let r = a.slice(1, 5).take_n_true(2);
+        // values: [true, true, true, false, true], keep the first 2 trues
+        assert_eq!(r, BooleanArray::from(vec![true, true, false, false, false]));
+    }
+
     #[test]
     fn test_take_n_true_preserves_nulls_and_skips_them() {
         // Non-null trues: positions 0, 3, 5. Null at 2 must not count toward `n`.

arrow-buffer/src/buffer/boolean.rs

@@ -555,6 +555,84 @@ impl BooleanBuffer {
         self.buffer
     }
 
+    /// Try to convert self into a [`BooleanBufferBuilder`] without copying.
+    ///
+    /// Reuses the underlying [`Buffer`] allocation if
+    /// 1. it is not shared (no other references to it exist) and
+    /// 2. the bit offset is zero
+    ///
+    /// Returns `Err(self)` if the allocation cannot be reused. See
+    /// [`Self::into_builder_or_clone`] for a variant that copies the bits in
+    /// this case instead.
+    ///
+    /// # Example
+    /// ```
+    /// # use arrow_buffer::BooleanBuffer;
+    /// let buffer = BooleanBuffer::from(vec![true, false, true]);
+    /// // The buffer is not shared, so the builder reuses its allocation
+    /// let mut builder = buffer.try_into_builder().expect("buffer was not shared");
+    /// builder.append(false);
+    /// assert_eq!(builder.finish(), BooleanBuffer::from(vec![true, false, true, false]));
+    ///
+    /// // Conversion fails if the buffer is shared
+    /// let buffer = BooleanBuffer::from(vec![true, false, true]);
+    /// let shared = buffer.clone();
+    /// let buffer = buffer.try_into_builder().expect_err("buffer was shared");
+    /// # assert_eq!(buffer, shared);
+    /// ```
+    pub fn try_into_builder(self) -> Result<BooleanBufferBuilder, Self> {
+        if self.bit_offset != 0 {
+            return Err(self);
+        }
+
+        let Self {
+            buffer,
+            bit_offset,
+            bit_len,
+        } = self;
+
+        match buffer.into_mutable() {
+            Ok(mutable_buffer) => Ok(BooleanBufferBuilder::new_from_buffer(
+                mutable_buffer,
+                bit_len,
+            )),
+            Err(buffer) => Err(Self {
+                buffer,
+                bit_offset,
+                bit_len,
+            }),
+        }
+    }
+
+    /// Converts self into a [`BooleanBufferBuilder`], reusing the existing
+    /// allocation if possible, and copying the bits otherwise.
+    ///
+    /// This is a convenience wrapper around [`Self::try_into_builder`] that
+    /// always succeeds. If the allocation cannot be reused (for example, the
+    /// underlying [`Buffer`] is shared, or the bit offset is non zero), the
+    /// bits are copied into a new allocation.
+    ///
+    /// # Example
+    /// ```
+    /// # use arrow_buffer::BooleanBuffer;
+    /// let buffer = BooleanBuffer::from(vec![true, false, true]);
+    /// // holding a second reference forces a copy
+    /// let shared = buffer.clone();
+    /// let mut builder = buffer.into_builder_or_clone();
+    /// builder.append(false);
+    /// assert_eq!(builder.finish(), BooleanBuffer::from(vec![true, false, true, false]));
+    /// // the original buffer is unchanged
+    /// assert_eq!(shared, BooleanBuffer::from(vec![true, false, true]));
+    /// ```
+    pub fn into_builder_or_clone(self) -> BooleanBufferBuilder {
+        self.try_into_builder().unwrap_or_else(|buffer| {
+            //  copy needed
+            let mut builder = BooleanBufferBuilder::new(buffer.len());
+            builder.append_buffer(&buffer);
+            builder
+        })
+    }
+
     /// Claim memory used by this buffer in the provided memory pool.
     ///
     /// See [`Buffer::claim`] for details.
@@ -826,6 +904,74 @@ mod tests {
         assert!(!boolean_buf.is_empty())
     }
 
+    #[test]
+    fn test_try_into_builder_reuses_allocation() {
+        let buffer = BooleanBuffer::from(vec![true, false, true]);
+        let ptr = buffer.values().as_ptr();
+        let mut builder = buffer.try_into_builder().unwrap();
+        builder.append(true);
+        let buffer = builder.finish();
+        assert_eq!(buffer, BooleanBuffer::from(vec![true, false, true, true]));
+        // same allocation: no copy was made
+        assert_eq!(buffer.values().as_ptr(), ptr);
+    }
+
+    #[test]
+    fn test_try_into_builder_shared() {
+        let buffer = BooleanBuffer::from(vec![true, false, true]);
+        let shared = buffer.clone();
+        let buffer = buffer.try_into_builder().unwrap_err();
+        assert_eq!(buffer, shared);
+    }
+
+    #[test]
+    fn test_try_into_builder_non_zero_offset() {
+        // an unshared buffer with a non zero bit offset cannot be converted
+        let buffer = BooleanBuffer::new(Buffer::from(vec![0b0000_0101u8]), 1, 2);
+        assert!(buffer.try_into_builder().is_err());
+    }
+
+    #[test]
+    fn test_try_into_builder_clears_padding() {
+        // bits between the logical length and the end of the last byte are
+        // cleared so subsequent appends see zeroed padding
+        let buffer = BooleanBuffer::new(Buffer::from(vec![0xFFu8]), 0, 3);
+        let mut builder = buffer.try_into_builder().unwrap();
+        builder.append(false);
+        builder.append(true);
+        assert_eq!(builder.finish().values(), &[0b0001_0111]);
+    }
+
+    #[test]
+    fn test_into_builder_or_clone() {
+        // unshared: reuses the allocation
+        let buffer = BooleanBuffer::from(vec![true, false]);
+        let ptr = buffer.values().as_ptr();
+        let mut builder = buffer.into_builder_or_clone();
+        builder.append(true);
+        let result = builder.finish();
+        assert_eq!(result, BooleanBuffer::from(vec![true, false, true]));
+        assert_eq!(result.values().as_ptr(), ptr);
+
+        // shared: copies, leaving the other reference intact
+        let shared = result.clone();
+        let mut builder = result.into_builder_or_clone();
+        builder.append(false);
+        let copied = builder.finish();
+        assert_eq!(copied, BooleanBuffer::from(vec![true, false, true, false]));
+        assert_eq!(shared, BooleanBuffer::from(vec![true, false, true]));
+
+        // non zero offset: copies, producing an offset 0 builder
+        let sliced = shared.slice(1, 2);
+        drop(shared);
+        let mut builder = sliced.into_builder_or_clone();
+        builder.append(true);
+        assert_eq!(
+            builder.finish(),
+            BooleanBuffer::from(vec![false, true, true])
+        );
+    }
+
     #[test]
     fn test_boolean_data_equality() {
         let boolean_buf1 = BooleanBuffer::new(Buffer::from(&[0, 1, 4, 3, 5]), 0, 32);

parquet/src/arrow/arrow_reader/mod.rs

@@ -1386,8 +1386,7 @@ impl FilterMaskAccumulator {
         *self = match std::mem::take(self) {
             Self::Empty => Self::Single(mask),
             Self::Single(first) => {
-                let mut combined = BooleanBufferBuilder::new(first.len() + mask.len());
-                combined.append_buffer(&first);
+                let mut combined = first.into_builder_or_clone();
                 combined.append_buffer(&mask);
                 Self::Combined(combined)
             }