From 510f02f449193bea9df3f423d18ce7a9e4112bdf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20Heres?= <danielheres@gmail.com>
Date: Tue, 11 May 2021 07:35:05 +0200
Subject: [PATCH] Speed up bound checking in `take` (#281)

* WIP improve take performance

* WIP

* Bound checking speed

* Simplify

* fmt

* Improve formatting
---
 arrow/benches/take_kernels.rs     | 19 ++++++++++++++++++-
 arrow/src/compute/kernels/take.rs | 25 +++++++++++++++++++------
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/arrow/benches/take_kernels.rs b/arrow/benches/take_kernels.rs
index 2853eb5d4760..b1d03d70c9a7 100644
--- a/arrow/benches/take_kernels.rs
+++ b/arrow/benches/take_kernels.rs
@@ -23,7 +23,7 @@ use rand::Rng;
 
 extern crate arrow;
 
-use arrow::compute::take;
+use arrow::compute::{take, TakeOptions};
 use arrow::datatypes::*;
 use arrow::util::test_util::seedable_rng;
 use arrow::{array::*, util::bench_util::*};
@@ -46,6 +46,12 @@ fn bench_take(values: &dyn Array, indices: &UInt32Array) {
     criterion::black_box(take(values, &indices, None).unwrap());
 }
 
+fn bench_take_bounds_check(values: &dyn Array, indices: &UInt32Array) {
+    criterion::black_box(
+        take(values, &indices, Some(TakeOptions { check_bounds: true })).unwrap(),
+    );
+}
+
 fn add_benchmark(c: &mut Criterion) {
     let values = create_primitive_array::<Int32Type>(512, 0.0);
     let indices = create_random_index(512, 0.0);
@@ -56,6 +62,17 @@ fn add_benchmark(c: &mut Criterion) {
         b.iter(|| bench_take(&values, &indices))
     });
 
+    let values = create_primitive_array::<Int32Type>(512, 0.0);
+    let indices = create_random_index(512, 0.0);
+    c.bench_function("take check bounds i32 512", |b| {
+        b.iter(|| bench_take_bounds_check(&values, &indices))
+    });
+    let values = create_primitive_array::<Int32Type>(1024, 0.0);
+    let indices = create_random_index(1024, 0.0);
+    c.bench_function("take check bounds i32 1024", |b| {
+        b.iter(|| bench_take_bounds_check(&values, &indices))
+    });
+
     let indices = create_random_index(512, 0.5);
     c.bench_function("take i32 nulls 512", |b| {
         b.iter(|| bench_take(&values, &indices))
diff --git a/arrow/src/compute/kernels/take.rs b/arrow/src/compute/kernels/take.rs
index 0217573dc5d3..d325ce44a185 100644
--- a/arrow/src/compute/kernels/take.rs
+++ b/arrow/src/compute/kernels/take.rs
@@ -100,17 +100,30 @@ where
     let options = options.unwrap_or_default();
     if options.check_bounds {
         let len = values.len();
-        for i in 0..indices.len() {
-            if indices.is_valid(i) {
-                let ix = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
+        if indices.null_count() > 0 {
+            indices.iter().flatten().try_for_each(|index| {
+                let ix = ToPrimitive::to_usize(&index).ok_or_else(|| {
                     ArrowError::ComputeError("Cast to usize failed".to_string())
                 })?;
                 if ix >= len {
                     return Err(ArrowError::ComputeError(
-                    format!("Array index out of bounds, cannot get item at index {} from {} entries", ix, len))
-                );
+                        format!("Array index out of bounds, cannot get item at index {} from {} entries", ix, len))
+                    );
                 }
-            }
+                Ok(())
+            })?;
+        } else {
+            indices.values().iter().try_for_each(|index| {
+                let ix = ToPrimitive::to_usize(index).ok_or_else(|| {
+                    ArrowError::ComputeError("Cast to usize failed".to_string())
+                })?;
+                if ix >= len {
+                    return Err(ArrowError::ComputeError(
+                        format!("Array index out of bounds, cannot get item at index {} from {} entries", ix, len))
+                    );
+                }
+                Ok(())
+            })?
         }
     }
     match values.data_type() {