From 772a759ee0d26cecec15f80662da4579326337f3 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 20 Jul 2026 14:56:47 -0400 Subject: [PATCH] Add BooleanBuffer::try_into_builder and into_builder_or_clone Co-Authored-By: Claude Fable 5 --- arrow-array/src/array/boolean_array.rs | 33 +++++- arrow-buffer/src/buffer/boolean.rs | 146 +++++++++++++++++++++++++ parquet/src/arrow/arrow_reader/mod.rs | 3 +- 3 files changed, 178 insertions(+), 4 deletions(-) diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs index 9245f5b33700..f9e2525a9ac5 100644 --- a/arrow-array/src/array/boolean_array.rs +++ b/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`. diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 52f1e3fbb510..211e648128b1 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/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 { + 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); diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 621c1c812afd..68e2c7ec9b12 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/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) }