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 0a5979560a Replace OffsetBufferBuilder with Vec in scalar byte-array
zip (#10913)
0a5979560a is described below
commit 0a5979560ac67135134c3618519b9b7384edf4d2
Author: Yash Kucheriya <[email protected]>
AuthorDate: Wed Sep 16 19:55:05 2026 -0700
Replace OffsetBufferBuilder with Vec in scalar byte-array zip (#10913)
# Which issue does this PR close?
Part of #10245.
# Rationale for this change
The byte-array scalar zip path is one of the remaining
`OffsetBufferBuilder` callsites in #10245. Building cumulative offsets
directly in a `Vec` avoids checking each offset through the builder,
which benefits masks containing long runs of the same value.
# What changes are included in this PR?
`BytesScalarImpl::create_output_on_non_nulls` builds offsets with
`Vec<T::Offset>`. The mixed-mask path checks the total byte size and
offset range before allocating the output. An overflowing size returns
`MemoryError`; a size outside the offset type returns
`OffsetOverflowError`. Every generated offset is bounded by that checked
total.
The existing all-true, all-false, and nullable-scalar fast paths are
unchanged.
# Are these changes tested?
- `cargo test -p arrow-select --quiet`: 418 unit tests and 17
documentation tests passed.
- `cargo clippy -p arrow-select --all-targets --all-features -- -D
warnings`: passed.
- `cargo +stable fmt --all -- --check`: passed.
The new regression test uses a 64 KiB binary scalar and a mixed mask
whose output would require 2 GiB. It checks that the call returns
`OffsetOverflowError` before allocating that output, with the large
scalar on either side of the mask. Existing tests cover normal and large
byte arrays, nulls, fragmented masks, and one-sided scalar paths.
## Benchmarks
The project runner measured the previous revision, `4829c8f6d`,
[here](https://github.com/apache/arrow-rs/pull/10913#issuecomment-5467796574).
Those results do not include the review fix.
Local measurements using the existing `arrow/benches/zip_kernels.rs`
benchmark, on macOS arm64 with Rust 1.97.1:
| Case, 8,192 rows | Upstream `cbbb56bba` | Previous PR `4829c8f6d` |
Revised `1ec832648` |
| --- | ---: | ---: | ---: |
| Short strings, 50% true | 40.00 us | 38.78 us | 38.40 us |
| Short strings, true then false | 7.04 us | 2.24 us | 2.24 us |
| Long strings, 50% true | 65.12 us | 64.14 us | 63.88 us |
| Long strings, true then false | 41.90 us | 35.69 us | 36.84 us |
| Short bytes, 50% true | 36.66 us | 34.84 us | 32.96 us |
| Short bytes, true then false | 7.01 us | 2.08 us | 1.99 us |
| Long bytes, 50% true | 54.93 us | 52.28 us | 53.47 us |
| Long bytes, true then false | 32.06 us | 25.53 us | 26.71 us |
These are short, sequential local runs, not a controlled performance
study. The revision remains faster than the upstream baseline in these
cases, but some long-value cases are slower than the previous PR
revision. A project-runner rerun would help confirm the impact of the
added checks.
Command, run for each revision with separate Criterion baselines:
```sh
cargo bench -p arrow --features test_utils --bench zip_kernels -- \
'zip_8192_from_(long|short).+non_nulls_scalars/(50pct_true|true_then_false)$' \
--sample-size 20 --warm-up-time 0.5 --measurement-time 1
```
# Are there any user-facing changes?
Mixed-mask scalar byte-array output that exceeds the offset range now
returns an error before output allocation instead of panicking. No
public API changes.
# Automated assistance
AI assistance was used to prepare the refactor, review fix, regression
test, and validation. The commands and measurements above were run
locally.
---------
Co-authored-by: Yash Kuceriya <[email protected]>
---
arrow-select/src/zip.rs | 69 ++++++++++++++++++++++++++++++++++++++-----------
1 file changed, 54 insertions(+), 15 deletions(-)
diff --git a/arrow-select/src/zip.rs b/arrow-select/src/zip.rs
index c0dcf6fbd9..1abde41d16 100644
--- a/arrow-select/src/zip.rs
+++ b/arrow-select/src/zip.rs
@@ -25,8 +25,8 @@ use arrow_array::types::{
};
use arrow_array::*;
use arrow_buffer::{
- BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer,
OffsetBufferBuilder,
- ScalarBuffer, ToByteSlice,
+ ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer,
OffsetBuffer, ScalarBuffer,
+ ToByteSlice,
};
use arrow_data::transform::MutableArrayData;
use arrow_data::{ArrayData, ByteView};
@@ -595,10 +595,19 @@ impl<T: ByteArrayType> BytesScalarImpl<T> {
}
}
- let total_number_of_bytes =
- true_count * truthy_val.len() + (predicate.len() - true_count) *
falsy_val.len();
+ let total_number_of_bytes = true_count
+ .checked_mul(truthy_val.len())
+ .and_then(|truthy_bytes| {
+ let falsy_bytes = (predicate.len() -
true_count).checked_mul(falsy_val.len())?;
+ truthy_bytes.checked_add(falsy_bytes)
+ })
+ .ok_or_else(|| ArrowError::MemoryError("zip output size
overflow".to_string()))?;
+ T::Offset::from_usize(total_number_of_bytes)
+ .ok_or(ArrowError::OffsetOverflowError(total_number_of_bytes))?;
let mut mutable = MutableBuffer::with_capacity(total_number_of_bytes);
- let mut offset_buffer_builder =
OffsetBufferBuilder::<T::Offset>::new(predicate.len());
+ let mut offsets = Vec::<T::Offset>::with_capacity(predicate.len() + 1);
+ offsets.push(T::Offset::usize_as(0));
+ let mut current_offset: usize = 0;
// keep track of how much is filled
let mut filled = 0;
@@ -606,6 +615,7 @@ impl<T: ByteArrayType> BytesScalarImpl<T> {
let truthy_len = truthy_val.len();
let falsy_len = falsy_val.len();
+ // Each run's offsets are bounded by the checked total output size.
SlicesIterator::from(predicate).try_for_each(|(start, end)| ->
Result<(), ArrowError> {
// the gap needs to be filled with falsy values
if start > filled {
@@ -615,9 +625,12 @@ impl<T: ByteArrayType> BytesScalarImpl<T> {
.try_repeat_slice_n_times(falsy_val, false_repeat_count)
.map_err(|e| ArrowError::MemoryError(e.to_string()))?;
- for _ in 0..false_repeat_count {
- offset_buffer_builder.push_length(falsy_len)
- }
+ let start_offset = current_offset;
+ current_offset += falsy_len * false_repeat_count;
+ offsets.extend(
+ (1..=false_repeat_count)
+ .map(|index| T::Offset::usize_as(start_offset + index
* falsy_len)),
+ );
}
let true_repeat_count = end - start;
@@ -626,9 +639,12 @@ impl<T: ByteArrayType> BytesScalarImpl<T> {
.try_repeat_slice_n_times(truthy_val, true_repeat_count)
.map_err(|e| ArrowError::MemoryError(e.to_string()))?;
- for _ in 0..true_repeat_count {
- offset_buffer_builder.push_length(truthy_len)
- }
+ let start_offset = current_offset;
+ current_offset += truthy_len * true_repeat_count;
+ offsets.extend(
+ (1..=true_repeat_count)
+ .map(|index| T::Offset::usize_as(start_offset + index *
truthy_len)),
+ );
filled = end;
Ok(())
})?;
@@ -640,12 +656,18 @@ impl<T: ByteArrayType> BytesScalarImpl<T> {
.try_repeat_slice_n_times(falsy_val, false_repeat_count)
.map_err(|e| ArrowError::MemoryError(e.to_string()))?;
- for _ in 0..false_repeat_count {
- offset_buffer_builder.push_length(falsy_len)
- }
+ let start_offset = current_offset;
+ current_offset += falsy_len * false_repeat_count;
+ offsets.extend(
+ (1..=false_repeat_count)
+ .map(|index| T::Offset::usize_as(start_offset + index *
falsy_len)),
+ );
}
- Ok((mutable.into(), offset_buffer_builder.finish()))
+ debug_assert_eq!(current_offset, total_number_of_bytes);
+ // SAFETY: offsets start at zero, are monotonically increasing, and
fit in T::Offset.
+ let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
+ Ok((mutable.into(), offsets))
}
}
@@ -1292,6 +1314,23 @@ mod test {
assert_eq!(actual, &expected);
}
+ #[test]
+ fn test_zip_scalar_bytes_offset_overflow() {
+ // Repeating a 64 KiB scalar 32,768 times exceeds i32::MAX.
+ // The size must be rejected before allocating the output buffer.
+ let value = vec![0_u8; 65_536];
+ let large =
Scalar::new(BinaryArray::from_iter_values([value.as_slice()]));
+ let empty =
Scalar::new(BinaryArray::from_iter_values([b"".as_slice()]));
+ let mask = BooleanArray::from_iter((0..65_536).map(|i| Some(i % 2 ==
0)));
+
+ for (truthy, falsy) in [(&large, &empty), (&empty, &large)] {
+ assert!(matches!(
+ zip(&mask, truthy, falsy),
+ Err(ArrowError::OffsetOverflowError(2_147_483_648))
+ ));
+ }
+ }
+
#[test]
fn test_zip_scalar_bytes_only_taking_one_side() {
let mask_len = 5;
