Speed up bound checking in take (#281)

* WIP improve take performance * WIP * Bound checking speed * Simplify * fmt * Improve formatting
apache · May 11, 2021 · 510f02f · 510f02f
1 parent aba044f
commit 510f02f
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 7 deletions.
diff --git 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
@@ -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() {