@@ -424,22 +424,25 @@ fn take_bits<I: ArrowPrimitiveType>(
424424 indices : & PrimitiveArray < I > ,
425425) -> BooleanBuffer {
426426 let len = indices. len ( ) ;
427- let mut output_buffer = MutableBuffer :: new_null ( len) ;
428- let output_slice = output_buffer. as_slice_mut ( ) ;
429427
430428 match indices. nulls ( ) . filter ( |n| n. null_count ( ) > 0 ) {
431- Some ( nulls) => nulls. valid_indices ( ) . for_each ( |idx| {
432- if values. value ( indices. value ( idx) . as_usize ( ) ) {
433- bit_util:: set_bit ( output_slice, idx) ;
434- }
435- } ) ,
436- None => indices. values ( ) . iter ( ) . enumerate ( ) . for_each ( |( i, index) | {
437- if values. value ( index. as_usize ( ) ) {
438- bit_util:: set_bit ( output_slice, i) ;
439- }
440- } ) ,
429+ Some ( nulls) => {
430+ let mut output_buffer = MutableBuffer :: new_null ( len) ;
431+ let output_slice = output_buffer. as_slice_mut ( ) ;
432+ nulls. valid_indices ( ) . for_each ( |idx| {
433+ if values. value ( indices. value ( idx) . as_usize ( ) ) {
434+ bit_util:: set_bit ( output_slice, idx) ;
435+ }
436+ } ) ;
437+ BooleanBuffer :: new ( output_buffer. into ( ) , 0 , len)
438+ }
439+ None => {
440+ BooleanBuffer :: collect_bool ( len, |idx : usize | {
441+ // SAFETY: idx<indices.len()
442+ values. value ( unsafe { indices. value_unchecked ( idx) . as_usize ( ) } )
443+ } )
444+ }
441445 }
442- BooleanBuffer :: new ( output_buffer. into ( ) , 0 , indices. len ( ) )
443446}
444447
445448/// `take` implementation for boolean arrays
0 commit comments