Rich-T-kid commented on code in PR #10813:
URL: https://github.com/apache/arrow-rs/pull/10813#discussion_r3888048452
##########
arrow-select/src/take.rs:
##########
@@ -456,43 +456,229 @@ 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();
+ 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| {
+ let src_idx = if CHECKED {
+ indices.value(valid_idx).as_usize()
+ } else {
+ // SAFETY: valid_idx < len (validity bitmap); caller
guarantees index values are in bounds when CHECKED=false
+ unsafe { indices.value_unchecked(valid_idx) }.as_usize()
+ } + src_offset;
+ // SAFETY: src_idx bounded by take's prior bounds check
+ unsafe { copy_bit_if_set(src_ptr, src_idx, out_ptr, valid_idx)
};
+ });
+ BooleanBuffer::new(Buffer::from(output), 0, len)
+ }
+ None => {
+ // 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 {
+ let src_idx = if CHECKED {
+ indices.value(base + bit).as_usize()
+ } else {
+ // SAFETY: base + bit < len (loop bound); caller
guarantees index values are in bounds when CHECKED=false
+ unsafe { indices.value_unchecked(base + bit)
}.as_usize()
+ } + src_offset;
+ // SAFETY: src_idx bounded by take's prior bounds check
+ byte |= unsafe { pack_bit(src_ptr, src_idx, 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) {
+ let src_idx = if CHECKED {
+ indices.value(base + bit).as_usize()
+ } else {
+ // SAFETY: base + bit < len (loop bound); caller
guarantees index values are in bounds when CHECKED=false
+ unsafe { indices.value_unchecked(base + bit)
}.as_usize()
+ } + src_offset;
+ // SAFETY: src_idx bounded by take's prior bounds check
+ byte |= unsafe { pack_bit(src_ptr, src_idx, bit) };
+ }
+ out_slice[full_bytes] = byte;
+ }
+ BooleanBuffer::new(Buffer::from(output), 0, 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) };
+/// 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) => {
+ // Vec is pre-zeroed; only set bits for valid indices via raw
pointer.
+ 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() {
+ let src_idx = if CHECKED {
+ indices.value(out_pos).as_usize()
+ } else {
+ // SAFETY: out_pos < len (validity bitmap); caller
guarantees index values are in bounds when CHECKED=false
+ unsafe { indices.value_unchecked(out_pos) }.as_usize()
+ };
+ // SAFETY: src_idx < values.len(); check_bounds ensures this
when CHECKED, caller guarantees it otherwise
Review Comment:
```suggestion
// SAFETY: src_idx < values.len();
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
