From c6510a6c8ce8885abd3f124428077fed1946396a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 24 Oct 2024 16:57:32 +0200 Subject: [PATCH] Improve --- arrow-select/src/take.rs | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index d08f3edf61f..2cda392925f 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -424,22 +424,25 @@ fn take_bits( indices: &PrimitiveArray, ) -> 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(