Skip to content

Commit a9294d7

Browse files
authored
Optimize take_bits / take_boolean: up to -25% (#6622)
* Optimize take * Wip * Revert "Wip" This reverts commit 3846bfc. * Revert "Revert "Wip"" This reverts commit b6a59fc. * Improve * Rev * Rev * Rev * Clippy * Fix * Style * Bring back take bytes optimization * Revert "Bring back take bytes optimization" This reverts commit 0449080. * Revert take_bytes change * Style
1 parent 1103939 commit a9294d7

File tree

1 file changed

+16
-13
lines changed

1 file changed

+16
-13
lines changed

arrow-select/src/take.rs

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)