Speedup take_bits (#2307)

apache · Aug 4, 2022 · 4b15b7e · 4b15b7e
1 parent d56d88e
commit 4b15b7e
Showing 1 changed file with 29 additions and 11 deletions.
diff --git 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())
 }