From 4b15b7e400b3cd7e5806958612c812e81d342355 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 4 Aug 2022 07:57:03 +0200 Subject: [PATCH] Speedup take_bits (#2307) --- arrow/src/compute/kernels/take.rs | 40 ++++++++++++++++++++++--------- 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/arrow/src/compute/kernels/take.rs b/arrow/src/compute/kernels/take.rs index ab99acd2c04..05832e4f527 100644 --- a/arrow/src/compute/kernels/take.rs +++ b/arrow/src/compute/kernels/take.rs @@ -614,23 +614,41 @@ where let mut output_buffer = MutableBuffer::new_null(len); let output_slice = output_buffer.as_slice_mut(); - indices - .iter() - .enumerate() - .try_for_each::<_, Result<()>>(|(i, index)| { - if let Some(index) = index { - let index = ToPrimitive::to_usize(&index).ok_or_else(|| { + let indices_has_nulls = indices.null_count() > 0; + + if indices_has_nulls { + indices + .iter() + .enumerate() + .try_for_each::<_, Result<()>>(|(i, index)| { + if let Some(index) = index { + let index = ToPrimitive::to_usize(&index).ok_or_else(|| { + ArrowError::ComputeError("Cast to usize failed".to_string()) + })?; + + if bit_util::get_bit(values_slice, values_offset + index) { + bit_util::set_bit(output_slice, i); + } + } + + Ok(()) + })?; + } else { + indices + .values() + .iter() + .enumerate() + .try_for_each::<_, Result<()>>(|(i, index)| { + let index = ToPrimitive::to_usize(index).ok_or_else(|| { ArrowError::ComputeError("Cast to usize failed".to_string()) })?; if bit_util::get_bit(values_slice, values_offset + index) { bit_util::set_bit(output_slice, i); } - } - - Ok(()) - })?; - + Ok(()) + })?; + } Ok(output_buffer.into()) }