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 c6f031305b perf: optimize take bool & take null buffers (#10813)
c6f031305b is described below
commit c6f031305be81f1ea3f41e88805bd87bb5848d5e
Author: RIchard Baah <[email protected]>
AuthorDate: Tue Sep 1 22:20:24 2026 -0400
perf: optimize take bool & take null buffers (#10813)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- works towards #8879.
# Rationale for this change
take on boolean arrays was slower than necessary: the nullable path
called take_bits twice (once for values, once for validity), and the
non-nullable path used high-level accessors that prevented the compiler
from issuing parallel loads.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
- `take_bits`: rewrites the None-nulls path to build output one byte at
a time with an 8-element inner loop, letting the compiler fully unroll
it and issue 8 source loads in parallel. Branches on the CHECKED const
generic to use value_unchecked when bounds are already guaranteed by the
caller.
- `take_bits_with_validity`: new single-pass helper that gathers value
bits and validity bits together when the source boolean array itself has
nulls, avoiding the previous double traversal.
- `take_boolean`: updated to use take_bits_with_validity in the nullable
case.
- ~~`take_record_batch_unchecked`: new pub unsafe counterpart to
`take_record_batch` that calls `take_impl::<_, false>` directly,
allowing us to use unsafe accessor methods~~
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
yes, I added new several boolean test
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
~~yes new public method. `take_record_batch_unchecked()`~~
none
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
arrow-select/src/take.rs | 349 ++++++++++++++++++++++++++++++++++++++++++++---
1 file changed, 328 insertions(+), 21 deletions(-)
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index 55b3d521ab..14906206e6 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -94,10 +94,10 @@ pub fn take(
let options = options.unwrap_or_default();
downcast_integer_array!(
indices => {
+ let indices = indices.to_indices();
if options.check_bounds {
- check_bounds(values.len(), indices)?;
+ check_bounds(values.len(), &indices)?;
}
- let indices = indices.to_indices();
take_impl::<_, true>(values, &indices)
},
d => Err(ArrowError::InvalidArgumentError(format!("Take only supported
for integers, got {d:?}")))
@@ -509,43 +509,252 @@ fn take_native<T: ArrowNativeType, I:
ArrowPrimitiveType>(
}
}
+/// Read the bit at `src_bit_idx` from `src` and, if it is set, write a `1` to
`dst_bit_idx`
+/// in `dst`. Leaves `dst_bit_idx` unchanged (zero) when the source bit is
unset.
+///
+/// ```text
+/// src = 0b00100000 (bit 5 is set)
+/// copy_bit_if_set(src, 5, dst, 2) → dst bit 2 becomes 1
+/// ```
+///
+/// # Safety
+/// - `src` must be valid for reads up to byte `src_bit_idx / 8`.
+/// - `dst` must be valid for writes up to byte `dst_bit_idx / 8`.
+#[inline(always)]
+unsafe fn copy_bit_if_set(src: *const u8, src_bit_idx: usize, dst: *mut u8,
dst_bit_idx: usize) {
+ unsafe {
+ if bit_util::get_bit_raw(src, src_bit_idx) {
+ bit_util::set_bit_raw(dst, dst_bit_idx);
+ }
+ }
+}
+
+/// Read the bit at `bit_idx` from `src` and return it shifted to `out_pos`,
ready to be
+/// OR'd into an output byte accumulator.
+///
+/// ```text
+/// src = 0b10100000 (bit 5 is set)
+/// pack_bit(src, 5, 2) → 0b00000100 (bit from position 5, placed at
position 2)
+/// ```
+///
+/// # Safety
+/// `src` must be valid for reads up to byte `bit_idx / 8`.
+#[inline(always)]
+unsafe fn pack_bit(src: *const u8, bit_idx: usize, out_pos: usize) -> u8 {
+ let byte = unsafe { *src.add(bit_idx >> 3) }; // byte containing bit
`bit_idx`
+ ((byte >> (bit_idx & 7)) & 1) << out_pos // extract the bit, shift to
output position
+}
+
#[inline(never)]
fn take_bits<I: ArrowPrimitiveType, const CHECKED: bool>(
values: &BooleanBuffer,
indices: &PrimitiveArray<I>,
) -> BooleanBuffer {
let len = indices.len();
-
- match indices.nulls().filter(|n| n.null_count() > 0) {
- Some(nulls) => {
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
- nulls.valid_indices().for_each(|idx| {
- // SAFETY: idx is a valid index in indices.nulls() -->
idx<indices.len()
- if unsafe {
values.value(indices.value_unchecked(idx).as_usize()) } {
- // SAFETY: MutableBuffer was created with space for
indices.len() bit, and idx < indices.len()
- unsafe { bit_util::set_bit_raw(output_slice.as_mut_ptr(),
idx) };
+ let src_offset = values.offset();
+ let src_ptr = values.values().as_ptr();
+ let out_bytes = len.div_ceil(8);
+
+ match indices.nulls().filter(|nulls| nulls.null_count() > 0) {
+ Some(index_nulls) => {
+ let mut output = vec![0u8; out_bytes];
+ let out_ptr = output.as_mut_ptr();
+ index_nulls.valid_indices().for_each(|valid_idx| {
+ // SAFETY: valid_idx < indices.len(), guaranteed by
valid_indices().
+ let index_val = unsafe { indices.value_unchecked(valid_idx)
}.as_usize();
+ if CHECKED {
+ if values.value(index_val) {
+ // SAFETY: valid_idx < indices.len() = len, output
buffer holds len bits.
+ unsafe { bit_util::set_bit_raw(out_ptr, valid_idx) };
+ }
+ } else {
+ // SAFETY: caller guarantees index_val < values.len().
+ unsafe { copy_bit_if_set(src_ptr, index_val + src_offset,
out_ptr, valid_idx) };
}
});
- BooleanBuffer::new(output_buffer.into(), 0, len)
+ BooleanBuffer::new(Buffer::from(output), 0, len)
}
None => {
- BooleanBuffer::collect_bool(len, |idx: usize| {
- // SAFETY: idx<indices.len()
- values.value(unsafe { indices.value_unchecked(idx).as_usize()
})
- })
+ // Build the output byte-by-byte with an 8-element inner loop so
the
+ // compiler can fully unroll it and issue the 8 source loads in
parallel.
+ let mut output = vec![0u8; out_bytes];
+ let out_slice = output.as_mut_slice();
+ let full_bytes = len / 8;
+
+ for (byte_idx, out_byte) in
out_slice.iter_mut().enumerate().take(full_bytes) {
+ let base = byte_idx * 8;
+ let mut byte = 0u8;
+ for bit in 0..8usize {
+ // SAFETY: base + bit < full_bytes * 8 <= len, so base +
bit is a valid
+ // position in the indices array.
+ let index_val = unsafe { indices.value_unchecked(base +
bit) }.as_usize();
+ if CHECKED {
+ byte |= (values.value(index_val) as u8) << bit;
+ } else {
+ // SAFETY: caller guarantees index_val < values.len().
+ byte |= unsafe { pack_bit(src_ptr, index_val +
src_offset, bit) };
+ }
+ }
+ *out_byte = byte;
+ }
+ // Handle remaining bits when len is not a multiple of 8.
+ if full_bytes < out_bytes {
+ let base = full_bytes * 8;
+ let mut byte = 0u8;
+ for bit in 0..(len - base) {
+ // SAFETY: base + bit < len (remainder loop bound), so
base + bit is a
+ // valid position in the indices array.
+ let index_val = unsafe { indices.value_unchecked(base +
bit) }.as_usize();
+ if CHECKED {
+ byte |= (values.value(index_val) as u8) << bit;
+ } else {
+ // SAFETY: caller guarantees index_val < values.len().
+ byte |= unsafe { pack_bit(src_ptr, index_val +
src_offset, bit) };
+ }
+ }
+ out_slice[full_bytes] = byte;
+ }
+ BooleanBuffer::new(Buffer::from(output), 0, len)
+ }
+ }
+}
+
+/// Gather value bits and validity bits from two boolean buffers in a single
pass.
+/// Used when the values array itself has nulls, avoiding two separate
`take_bits` calls.
+#[inline(never)]
+fn take_bits_with_validity<I: ArrowPrimitiveType, const CHECKED: bool>(
+ values: &BooleanBuffer,
+ validity: &BooleanBuffer,
+ indices: &PrimitiveArray<I>,
+) -> (BooleanBuffer, Option<NullBuffer>) {
+ let len = indices.len();
+ let value_bit_offset = values.offset();
+ let validity_bit_offset = validity.offset();
+ let value_data_ptr = values.values().as_ptr();
+ let validity_data_ptr = validity.values().as_ptr();
+ let out_bytes = len.div_ceil(8);
+
+ let mut value_out = vec![0u8; out_bytes];
+ let mut validity_out = vec![0u8; out_bytes];
+
+ match indices.nulls().filter(|nulls| nulls.null_count() > 0) {
+ Some(index_nulls) => {
+ let value_out_ptr = value_out.as_mut_ptr();
+ let validity_out_ptr = validity_out.as_mut_ptr();
+ for out_pos in index_nulls.valid_indices() {
+ // SAFETY: out_pos < indices.len(), guaranteed by
valid_indices().
+ let src_idx = unsafe { indices.value_unchecked(out_pos)
}.as_usize();
+ if CHECKED {
+ if values.value(src_idx) {
+ // SAFETY: out_pos < indices.len() = len, output
buffer holds len bits.
+ unsafe { bit_util::set_bit_raw(value_out_ptr, out_pos)
};
+ }
+ if validity.value(src_idx) {
+ // SAFETY: out_pos < indices.len() = len, output
buffer holds len bits.
+ unsafe { bit_util::set_bit_raw(validity_out_ptr,
out_pos) };
+ }
+ } else {
+ // SAFETY: caller guarantees src_idx < values.len().
+ unsafe {
+ copy_bit_if_set(
+ value_data_ptr,
+ src_idx + value_bit_offset,
+ value_out_ptr,
+ out_pos,
+ );
+ copy_bit_if_set(
+ validity_data_ptr,
+ src_idx + validity_bit_offset,
+ validity_out_ptr,
+ out_pos,
+ );
+ }
+ }
+ }
+ }
+ None => {
+ let value_out_slice = value_out.as_mut_slice();
+ let validity_out_slice = validity_out.as_mut_slice();
+ let full_bytes = len / 8;
+
+ for (byte_idx, (value_out_byte, validity_out_byte)) in
value_out_slice
+ .iter_mut()
+ .zip(validity_out_slice.iter_mut())
+ .enumerate()
+ .take(full_bytes)
+ {
+ let bit_base = byte_idx * 8;
+ let mut packed_values = 0u8;
+ let mut packed_validity = 0u8;
+ for bit_pos in 0..8usize {
+ // SAFETY: bit_base + bit_pos < full_bytes * 8 <= len.
+ let src_idx = unsafe { indices.value_unchecked(bit_base +
bit_pos) }.as_usize();
+ if CHECKED {
+ packed_values |= (values.value(src_idx) as u8) <<
bit_pos;
+ packed_validity |= (validity.value(src_idx) as u8) <<
bit_pos;
+ } else {
+ // SAFETY: caller guarantees src_idx < values.len().
+ packed_values |= unsafe {
+ pack_bit(value_data_ptr, src_idx +
value_bit_offset, bit_pos)
+ };
+ packed_validity |= unsafe {
+ pack_bit(validity_data_ptr, src_idx +
validity_bit_offset, bit_pos)
+ };
+ }
+ }
+ *value_out_byte = packed_values;
+ *validity_out_byte = packed_validity;
+ }
+ // Handle remaining bits when len is not a multiple of 8.
+ if full_bytes < out_bytes {
+ let bit_base = full_bytes * 8;
+ let mut packed_values = 0u8;
+ let mut packed_validity = 0u8;
+ for bit_pos in 0..(len - bit_base) {
+ // SAFETY: bit_base + bit_pos < len (remainder loop bound).
+ let src_idx = unsafe { indices.value_unchecked(bit_base +
bit_pos) }.as_usize();
+ if CHECKED {
+ packed_values |= (values.value(src_idx) as u8) <<
bit_pos;
+ packed_validity |= (validity.value(src_idx) as u8) <<
bit_pos;
+ } else {
+ // SAFETY: caller guarantees src_idx < values.len().
+ packed_values |= unsafe {
+ pack_bit(value_data_ptr, src_idx +
value_bit_offset, bit_pos)
+ };
+ packed_validity |= unsafe {
+ pack_bit(validity_data_ptr, src_idx +
validity_bit_offset, bit_pos)
+ };
+ }
+ }
+ value_out_slice[full_bytes] = packed_values;
+ validity_out_slice[full_bytes] = packed_validity;
+ }
}
}
+
+ let value_buf_out = BooleanBuffer::new(Buffer::from(value_out), 0, len);
+ let validity_buf_out = NullBuffer::from_unsliced_buffer(validity_out, len);
+ (value_buf_out, validity_buf_out)
}
/// `take` implementation for boolean arrays
fn take_boolean<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
- values: &BooleanArray,
+ array: &BooleanArray,
indices: &PrimitiveArray<IndexType>,
) -> BooleanArray {
- let val_buf = take_bits::<_, CHECKED>(values.values(), indices);
- let null_buf = take_nulls::<_, CHECKED>(values.nulls(), indices);
- BooleanArray::new(val_buf, null_buf)
+ let bits = array.values();
+ match array.nulls().filter(|n| n.null_count() > 0) {
+ Some(array_nulls) => {
+ let (val_buf, null_buf) =
+ take_bits_with_validity::<_, CHECKED>(bits,
array_nulls.inner(), indices);
+ BooleanArray::new(val_buf, null_buf)
+ }
+ None => {
+ let val_buf = take_bits::<_, CHECKED>(bits, indices);
+ let null_buf = take_nulls::<_, CHECKED>(None, indices);
+ BooleanArray::new(val_buf, null_buf)
+ }
+ }
}
/// `take` implementation for string arrays
@@ -1946,6 +2155,104 @@ mod tests {
);
}
+ #[test]
+ // Null indices + sliced source boolean array — exercises src_offset in
the sparse take_bits path.
+ fn test_take_bool_nullable_index_sliced_source() {
+ let source = BooleanArray::from(vec![Some(true), Some(false),
Some(true), Some(false)]);
+ let source = source.slice(1, 3); // logical: [false, true, false],
offset=1
+ let source = source;
+
+ let indices = UInt32Array::from(vec![Some(2), None, Some(0)]);
+ let result = take(&source, &indices, None).unwrap();
+ let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let expected = BooleanArray::from(vec![Some(false), None,
Some(false)]);
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ // >8 elements: exercises the 8-at-a-time unrolled byte packing in
take_bits.
+ fn test_take_bool_no_nulls_multi_byte() {
+ let source = BooleanArray::from(vec![
+ true, false, true, true, false, false, true, false, true, true,
+ ]);
+ let indices = UInt32Array::from(vec![0, 2, 4, 6, 8, 1, 3, 5, 7, 9]);
+ let result = take(&source, &indices, None).unwrap();
+ let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let expected = BooleanArray::from(vec![
+ true, true, false, true, true, false, true, false, false, true,
+ ]);
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ // Sliced source: verifies src_offset is applied when no null indices.
+ fn test_take_bool_no_nulls_sliced_source() {
+ let source = BooleanArray::from(vec![true, false, true, false, true]);
+ let source = source.slice(2, 3); // [true, false, true], offset=2
+ let source = source.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let indices = UInt32Array::from(vec![2, 0, 1]);
+ let result = take(source, &indices, None).unwrap();
+ let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let expected = BooleanArray::from(vec![true, true, false]);
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ // Nullable source, >8 elements: exercises the 8-at-a-time unrolled byte
packing in take_bits_with_validity.
+ fn test_take_bool_nullable_values_multi_byte() {
+ let source = BooleanArray::from(vec![
+ Some(true),
+ None,
+ Some(false),
+ Some(true),
+ None,
+ Some(false),
+ Some(true),
+ Some(false),
+ Some(true),
+ None,
+ ]);
+ let indices = UInt32Array::from(vec![0, 2, 4, 6, 8, 1, 3, 5, 7, 9]);
+ let result = take(&source, &indices, None).unwrap();
+ let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let expected = BooleanArray::from(vec![
+ Some(true),
+ Some(false),
+ None,
+ Some(true),
+ Some(true),
+ None,
+ Some(true),
+ Some(false),
+ Some(false),
+ None,
+ ]);
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ // Nullable source, null indices, sliced source: verifies src_offset with
both null paths.
+ fn test_take_bool_nullable_values_sliced_source_null_indices() {
+ let source =
+ BooleanArray::from(vec![Some(true), Some(false), None, Some(true),
Some(false)]);
+ let source = source.slice(1, 4); // [false, null, true, false],
offset=1
+ let source = source.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let indices = UInt32Array::from(vec![Some(3), None, Some(1), Some(0)]);
+ let result = take(source, &indices, None).unwrap();
+ let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let expected = BooleanArray::from(vec![Some(false), None, None,
Some(false)]);
+ assert_eq!(result, &expected);
+ }
+
+ #[test]
+ #[should_panic(expected = "assertion failed: idx < self.bit_len")]
+ fn test_take_bool_oob_no_check_bounds_panics() {
+ let array = BooleanArray::from(vec![true, false, true]);
+ let indices = Int32Array::from(vec![0, 1, 10]);
+ take(&array, &indices, None).unwrap();
+ }
+
fn _test_take_string<'a, K>()
where
K: Array + PartialEq + From<Vec<Option<&'a str>>> + 'static,
