Rich-T-kid commented on code in PR #10813:
URL: https://github.com/apache/arrow-rs/pull/10813#discussion_r3891869579
##########
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
Review Comment:
I should rephrase, currently this branch should never panic. Currently the
only way to reach `take_bits` with CHECKED=true is through the public take
method that checks the bounds before any specific take_impl's
https://github.com/apache/arrow-rs/blob/70fa5bcf21924c5a9ea34ee2df846ffae43f2bc6/arrow-select/src/take.rs#L97
>itll panic if we go beyond bounds for indices itself, but the value
returned from indices itself (the src_idx) isnt guaranteed to be in bounds of
the array we're taking on
This is true. for example
```
values = [T,F,T,T]
indices = [2,99]
let src_ind = indices.value(1) // 99
pack_bit(src_ptr,src_idx,bit) // will attempt to read bytes outside of the
array
```
but since this branch isn't reachable since take_impl is called with the
`true` generic in this case. `check_bounds` bounds would have caught this here
https://github.com/apache/arrow-rs/blob/70fa5bcf21924c5a9ea34ee2df846ffae43f2bc6/arrow-select/src/take.rs#L185
the safety guarantee comes from `check_bounds`. will update the comments
agiain
##########
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
Review Comment:
I should rephrase, currently this branch should never panic. Currently the
only way to reach `take_bits` with CHECKED=true is through the public take
method that checks the bounds before any specific take_impl's
https://github.com/apache/arrow-rs/blob/70fa5bcf21924c5a9ea34ee2df846ffae43f2bc6/arrow-select/src/take.rs#L97
>itll panic if we go beyond bounds for indices itself, but the value
returned from indices itself (the src_idx) isnt guaranteed to be in bounds of
the array we're taking on
This is true. for example
```
values = [T,F,T,T]
indices = [2,99]
let src_ind = indices.value(1) // 99
pack_bit(src_ptr,src_idx,bit) // will attempt to read bytes outside of the
array
```
but since this branch isn't reachable since take_impl is called with the
`true` generic in this case. `check_bounds` bounds would have caught this here
https://github.com/apache/arrow-rs/blob/70fa5bcf21924c5a9ea34ee2df846ffae43f2bc6/arrow-select/src/take.rs#L185
the safety guarantee comes from `check_bounds`. will update the comments
--
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]
