Jefffrey commented on code in PR #10813:
URL: https://github.com/apache/arrow-rs/pull/10813#discussion_r3892015318


##########
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:
   given this unit test:
   
   ```rust
   #[test]
   fn test123() {
       let array = BooleanArray::from(vec![true, false, true]);
       let indices = Int32Array::from(vec![0, 1, 10]);
       let array = take(&array, &indices, None).unwrap();
       dbg!(array);
   }
   ```
   
   running it, it succeeds normally:
   
   ```sh
   running 1 test
   [arrow-select/src/take.rs:3329:9] array = BooleanArray
   [
     true,
     false,
     false,
   ]
   test take::tests::test123 ... ok
   ```
   
   but under miri it fails, because that last index isnt bounds checked against 
the input array:
   
   ```sh
   running 1 test
   test take::tests::test123 ... error: Undefined Behavior: memory access 
failed: attempting to access 1 byte, but got alloc274465+0x1 which is at or 
beyond the end of the allocation of size 1 byte
      --> arrow-select/src/take.rs:491:25
       |
   491 |     let byte = unsafe { *src.add(bit_idx >> 3) }; // byte containing 
bit `bit_idx`
       |                         ^^^^^^^^^^^^^^^^^^^^^^ Undefined Behavior 
occurred here
       |
       = help: this indicates a bug in the program: it performed an invalid 
operation, and caused Undefined Behavior
       = help: see 
https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html 
for further information
   help: alloc274465 was allocated here:
      --> arrow-buffer/src/buffer/mutable.rs:215:39
       = note: this is on thread `take::tests::test123`
       = note: stack backtrace:
               0: take::pack_bit
                   at arrow-select/src/take.rs:491:25: 491:47
               1: take::take_bits::<arrow_array::types::UInt32Type, true>
                   at arrow-select/src/take.rs:566:38: 566:69
               2: take::take_boolean::<arrow_array::types::UInt32Type, true>
                   at arrow-select/src/take.rs:688:27: 688:65
               3: take::take_impl::<arrow_array::types::UInt32Type, true>
                   at arrow-select/src/take.rs:223:25: 223:68
               4: take::take
                   at arrow-select/src/take.rs:101:13: 101:51
               5: take::tests::test123
                   at arrow-select/src/take.rs:3328:21: 3328:49
               6: take::tests::test123::{closure#0}
                   at arrow-select/src/take.rs:3325:17: 3325:17
   
   note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` 
for a verbose backtrace
   
   error: aborting due to 1 previous error
   
   error: test failed, to rerun pass `-p arrow-select --lib`
   ```



-- 
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]

Reply via email to