Improve

arrow-select/src/take.rs

@@ -424,22 +424,25 @@ fn take_bits<I: ArrowPrimitiveType>(
     indices: &PrimitiveArray<I>,
 ) -> BooleanBuffer {
     let len = indices.len();
-    let mut output_buffer = MutableBuffer::new_null(len);
-    let output_slice = output_buffer.as_slice_mut();
 
     match indices.nulls().filter(|n| n.null_count() > 0) {
-        Some(nulls) => nulls.valid_indices().for_each(|idx| {
-            if values.value(indices.value(idx).as_usize()) {
-                bit_util::set_bit(output_slice, idx);
-            }
-        }),
+        Some(nulls) => {
+            let mut output_buffer = MutableBuffer::new_null(len);
+            let output_slice: &mut [u8] = output_buffer.as_slice_mut();
+
+            nulls.valid_indices().for_each(|idx| {
+                if values.value(indices.value(idx).as_usize()) {
+                    bit_util::set_bit(output_slice, idx);
+                }
+            });
+            return BooleanBuffer::new(output_buffer.into(), 0, indices.len());
+        }
         None => {
             return BooleanBuffer::collect_bool(len, |idx: usize| {
                 values.value(unsafe { indices.value_unchecked(idx).as_usize() })
             });
         }
     }
-    BooleanBuffer::new(output_buffer.into(), 0, indices.len())
 }
 /// `take` implementation for boolean arrays
 fn take_boolean<IndexType: ArrowPrimitiveType>(