This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 6e25fe3a67 fix(arrow-buffer): add a fallible collect_bool (#10984)
6e25fe3a67 is described below
commit 6e25fe3a67584f7286da327176dec8da0d40b5e9
Author: pawan <[email protected]>
AuthorDate: Mon Sep 7 06:18:32 2026 +0530
fix(arrow-buffer): add a fallible collect_bool (#10984)
# Which issue does this PR close?
- Closes #10973.
# Rationale for this change
`collect_bool` reserves ceil(len / 64) u64 words before it calls the
closure once, so a `len` that comes from a row or bit count in user
controlled data aborts the process rather than failing. the reservation
went through `Vec::with_capacity`, which has no way to report that.
this file already has the shape for it. `with_capacity` is
`try_with_capacity(..).unwrap_or_else(|e| panic!("{e}"))` at line 154,
and there are seven other fallible twins beside it:
`try_from_len_zeroed`, `try_reserve`, `try_repeat_slice_n_times`,
`try_resize`, `try_shrink_to_fit`, `try_extend_from_slice` and
`try_extend_zeros`. `collect_bool` is the allocating entry point that
does not have one.
# What changes are included in this PR?
- add `try_collect_bool` returning `Result<Self, MutableBufferError>`
- `collect_bool` calls it and unwraps, the same way `with_capacity`
calls `try_with_capacity`
- the reservation goes through `Vec::try_reserve`, and a failure maps
onto the existing `AllocationError` variant carrying the layout it tried
to take
no behaviour change for any `len` that already worked. `collect_bool`
still panics, one level down, and its doc comment now says so.
i went with a fallible entry point rather than capping the reservation,
since a cap only moves where the abort happens rather than letting a
caller handle it. happy to do the cap instead if you would rather not
grow the api surface.
# Are these changes tested?
yes, two tests.
`try_collect_bool_reports_a_len_it_cannot_reserve` uses the 2^60 from
the issue and asserts it comes back as `AllocationError` and that the
closure ran zero times. the reproducer in the issue aborts on that same
input with `memory allocation of 144115188075855872 bytes failed`.
`try_collect_bool_matches_collect_bool_for_sizes_that_fit` checks the
two agree byte for byte at 0, 1, 63, 64, 65 and 1000 bits, since 64 is
the word boundary and 8 is the truncation boundary.
`arrow-buffer` is 351 passed, `arrow-array` is 721 passed, and fmt and
clippy with `-D warnings` are clean. i ran arrow-array as well because
`collect_bool` has around twenty call sites in it and this changes how
the function is built.
# Are there any user-facing changes?
`try_collect_bool` is new and additive. `collect_bool` keeps its
signature and still panics on a `len` it cannot reserve.
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
arrow-buffer/src/buffer/mutable.rs | 28 +++++++++++++++++++++++++---
1 file changed, 25 insertions(+), 3 deletions(-)
diff --git a/arrow-buffer/src/buffer/mutable.rs
b/arrow-buffer/src/buffer/mutable.rs
index 7ae7d1337f..610d9a738b 100644
--- a/arrow-buffer/src/buffer/mutable.rs
+++ b/arrow-buffer/src/buffer/mutable.rs
@@ -737,9 +737,31 @@ impl MutableBuffer {
///
/// This is similar to `from_trusted_len_iter_bool`, however, can be
significantly faster
/// as it eliminates the conditional `Iterator::next`
+ ///
+ /// # Panics
+ ///
+ /// Panics if the backing storage for `len` bits cannot be allocated. Use
+ /// [`MutableBuffer::try_collect_bool`] for a fallible version.
+ #[inline]
+ pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
+ Self::try_collect_bool(len, f).unwrap_or_else(|e| panic!("{e}"))
+ }
+
+ /// Fallible version of [`MutableBuffer::collect_bool`].
+ ///
+ /// `len` is a bit count, so the reservation is `ceil(len / 64)` u64
slots. This function
+ /// returns an error if that much memory cannot be reserved up front.
#[inline]
- pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, mut f: F) -> Self
{
- let mut buffer: Vec<u64> = Vec::with_capacity(bit_util::ceil(len, 64));
+ pub fn try_collect_bool<F: FnMut(usize) -> bool>(
+ len: usize,
+ mut f: F,
+ ) -> Result<Self, MutableBufferError> {
+ let words = bit_util::ceil(len, 64);
+ let layout = Layout::array::<u64>(words).map_err(|_|
MutableBufferError::LayoutError)?;
+ let mut buffer: Vec<u64> = Vec::new();
+ buffer
+ .try_reserve(words)
+ .map_err(|_| MutableBufferError::AllocationError(layout))?;
let chunks = len / 64;
let remainder = len % 64;
@@ -765,7 +787,7 @@ impl MutableBuffer {
let mut buffer: MutableBuffer = buffer.into();
buffer.truncate(bit_util::ceil(len, 8));
- buffer
+ Ok(buffer)
}
/// Extends this buffer with boolean values.
