From 6a0e8e93e6cfeb49cb347035eeb8cecdb3d6a00f Mon Sep 17 00:00:00 2001 From: "Heres, Daniel" Date: Wed, 3 Aug 2022 22:44:24 +0200 Subject: [PATCH] Speedup take_bits --- 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 5bfd257fcf4..41ddb6d0acb 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()) }