Skip to content

Commit

Permalink
Optimize take_bits / take_boolean: up to -25% (#6622)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
Dandandan authored Oct 24, 2024
1 parent 1103939 commit a9294d7
Showing 1 changed file with 16 additions and 13 deletions.
29 changes: 16 additions & 13 deletions arrow-select/src/take.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}),
None => indices.values().iter().enumerate().for_each(|(i, index)| {
if values.value(index.as_usize()) {
bit_util::set_bit(output_slice, i);
}
}),
Some(nulls) => {
let mut output_buffer = MutableBuffer::new_null(len);
let output_slice = 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);
}
});
BooleanBuffer::new(output_buffer.into(), 0, len)
}
None => {
BooleanBuffer::collect_bool(len, |idx: usize| {
// SAFETY: idx<indices.len()
values.value(unsafe { indices.value_unchecked(idx).as_usize() })
})
}
}
BooleanBuffer::new(output_buffer.into(), 0, indices.len())
}

/// `take` implementation for boolean arrays
Expand Down

0 comments on commit a9294d7

Please sign in to comment.