ARROW-10540: [Rust] Extended filter kernel to all types and improved …

…performance

This PR improves the filter kernel:

* made the filter benchmarks more realistic
* performance improved by 1.2-4x for all multi-filter operations
* performance decreased by 30% for a single-filter operation with 50% taken (and 2x faster for dense)
* filter now supports all types supported by `MutableArrayData` (in particular nested lists, `struct`, etc.)
* removed 400 LOC

There are two novel ideas here:

1. it minimizes the number of memcopies when building the filtered array, both for single filter and multi-filter operations.

2. for single filter operations, it leverages an iterator to create the new array on the fly. For multi filter operations, it persists the iterator's result in a vector and iterates over it per array.

This PR also improves the performance of `MutableArrayData` by avoiding some bound checks via `unsafe` (properly documented).

Summary of the benchmarks:

|  benchmark | variation (%) |
|-------------- | -------------- |
| filter u8 | 29.5 |
| filter u8 low selectivity | 7.3 |
| filter context u8 w NULLs | -17.5 |
| filter context u8 w NULLs high selectivity | -21.9 |
| filter context f32 high selectivity | -22.0 |
| filter context f32 | -26.8 |
| filter context string high selectivity | -27.5 |
| filter context string | -31.4 |
| filter context u8 | -40.3 |
| filter u8 high selectivity | -47.3 |
| filter context string low selectivity | -55.3 |
| filter context u8 w NULLs low selectivity | -57.7 |
| filter context f32 low selectivity | -64.8 |
| filter context u8 low selectivity | -66.0 |
| filter context u8 high selectivity | -77.2 |

Code used to benchmark:
```bash
git checkout 54da437
cargo bench --bench filter_kernels
git checkout mutable_filter2
cargo bench --bench filter_kernels
```

Benchmark result:
```
   Compiling arrow v3.0.0-SNAPSHOT (/Users/jorgecarleitao/projects/arrow/rust/arrow)
    Finished bench [optimized] target(s) in 1m 01s
     Running /Users/jorgecarleitao/projects/arrow/rust/target/release/deps/filter_kernels-5208f9a404de52c9
Gnuplot not found, using plotters backend
filter u8               time:   [512.54 us 513.43 us 514.37 us]
                        change: [+29.070% +29.548% +30.003%] (p = 0.00 < 0.05)
                        Performance has regressed.
Found 5 outliers among 100 measurements (5.00%)
  3 (3.00%) high mild
  2 (2.00%) high severe

filter u8 high selectivity
                        time:   [11.494 us 11.513 us 11.532 us]
                        change: [-47.846% -47.337% -46.755%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 7 outliers among 100 measurements (7.00%)
  2 (2.00%) high mild
  5 (5.00%) high severe

filter u8 low selectivity
                        time:   [7.0342 us 7.0520 us 7.0693 us]
                        change: [+6.5543% +7.3409% +8.1080%] (p = 0.00 < 0.05)
                        Performance has regressed.
Found 5 outliers among 100 measurements (5.00%)
  1 (1.00%) high mild
  4 (4.00%) high severe

filter context u8       time:   [233.81 us 234.31 us 234.93 us]
                        change: [-40.715% -40.329% -39.886%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 8 outliers among 100 measurements (8.00%)
  3 (3.00%) high mild
  5 (5.00%) high severe

filter context u8 high selectivity
                        time:   [4.5943 us 4.6100 us 4.6276 us]
                        change: [-77.449% -77.231% -77.022%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 18 outliers among 100 measurements (18.00%)
  10 (10.00%) high mild
  8 (8.00%) high severe

filter context u8 low selectivity
                        time:   [1.7582 us 1.7664 us 1.7742 us]
                        change: [-66.250% -65.989% -65.669%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 2 outliers among 100 measurements (2.00%)
  2 (2.00%) high severe

filter context u8 w NULLs
                        time:   [476.99 us 477.71 us 478.44 us]
                        change: [-17.852% -17.457% -17.000%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 6 outliers among 100 measurements (6.00%)
  3 (3.00%) high mild
  3 (3.00%) high severe

filter context u8 w NULLs high selectivity
                        time:   [296.46 us 297.03 us 297.67 us]
                        change: [-22.297% -21.871% -21.393%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 7 outliers among 100 measurements (7.00%)
  3 (3.00%) high mild
  4 (4.00%) high severe

filter context u8 w NULLs low selectivity
                        time:   [2.5988 us 2.6124 us 2.6268 us]
                        change: [-58.065% -57.668% -57.237%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
  4 (4.00%) high mild
  1 (1.00%) high severe

filter context f32      time:   [470.69 us 472.39 us 474.73 us]
                        change: [-29.574% -26.769% -24.242%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 14 outliers among 100 measurements (14.00%)
  9 (9.00%) high mild
  5 (5.00%) high severe

filter context f32 high selectivity
                        time:   [307.16 us 307.58 us 308.03 us]
                        change: [-22.472% -22.039% -21.532%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 6 outliers among 100 measurements (6.00%)
  2 (2.00%) high mild
  4 (4.00%) high severe

filter context f32 low selectivity
                        time:   [2.4266 us 2.4323 us 2.4384 us]
                        change: [-65.024% -64.764% -64.517%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 7 outliers among 100 measurements (7.00%)
  5 (5.00%) high mild
  2 (2.00%) high severe

filter context string   time:   [645.82 us 647.32 us 649.04 us]
                        change: [-31.810% -31.427% -31.046%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 11 outliers among 100 measurements (11.00%)
  6 (6.00%) high mild
  5 (5.00%) high severe

Benchmarking filter context string high selectivity: Warming up for 3.0000 s
Warning: Unable to complete 100 samples in 5.0s. You may wish to increase target time to 5.2s, enable flat sampling, or reduce sample count to 60.
filter context string high selectivity
                        time:   [999.11 us 1.0008 ms 1.0027 ms]
                        change: [-28.133% -27.524% -26.930%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 11 outliers among 100 measurements (11.00%)
  7 (7.00%) high mild
  4 (4.00%) high severe

filter context string low selectivity
                        time:   [3.6441 us 3.6623 us 3.6799 us]
                        change: [-55.650% -55.329% -55.013%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 8 outliers among 100 measurements (8.00%)
  6 (6.00%) low mild
  2 (2.00%) high severe
```

Closes #8960 from jorgecarleitao/mutable_filter2

Authored-by: Jorge C. Leitao <jorgecarleitao@gmail.com>
Signed-off-by: Andrew Lamb <andrew@nerdnetworks.org>

rust/arrow/Cargo.toml

@@ -137,3 +137,7 @@ harness = false
 [[bench]]
 name = "concatenate_kernel"
 harness = false
+
+[[bench]]
+name = "mutable_array"
+harness = false

rust/arrow/benches/filter_kernels.rs

@@ -14,128 +14,144 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
+extern crate arrow;
+
+use arrow::{compute::Filter, util::test_util::seedable_rng};
+use rand::{
+    distributions::{Alphanumeric, Standard},
+    prelude::Distribution,
+    Rng,
+};
 
 use arrow::array::*;
-use arrow::compute::{filter, FilterContext};
+use arrow::compute::{build_filter, filter};
 use arrow::datatypes::ArrowNumericType;
+use arrow::datatypes::{Float32Type, UInt8Type};
+
 use criterion::{criterion_group, criterion_main, Criterion};
 
-fn create_primitive_array<T, F>(size: usize, value_fn: F) -> PrimitiveArray<T>
+fn create_primitive_array<T>(size: usize, null_density: f32) -> PrimitiveArray<T>
 where
     T: ArrowNumericType,
-    F: Fn(usize) -> T::Native,
+    Standard: Distribution<T::Native>,
 {
+    // use random numbers to avoid spurious compiler optimizations wrt to branching
+    let mut rng = seedable_rng();
     let mut builder = PrimitiveArray::<T>::builder(size);
-    for i in 0..size {
-        builder.append_value(value_fn(i)).unwrap();
+
+    for _ in 0..size {
+        if rng.gen::<f32>() < null_density {
+            builder.append_null().unwrap();
+        } else {
+            builder.append_value(rng.gen()).unwrap();
+        }
     }
     builder.finish()
 }
 
-fn create_u8_array_with_nulls(size: usize) -> UInt8Array {
-    let mut builder = UInt8Builder::new(size);
-    for i in 0..size {
-        if i % 2 == 0 {
-            builder.append_value(1).unwrap();
-        } else {
+fn create_string_array(size: usize, null_density: f32) -> StringArray {
+    // use random numbers to avoid spurious compiler optimizations wrt to branching
+    let mut rng = seedable_rng();
+    let mut builder = StringBuilder::new(size);
+
+    for _ in 0..size {
+        if rng.gen::<f32>() < null_density {
             builder.append_null().unwrap();
+        } else {
+            let value = (&mut rng)
+                .sample_iter(&Alphanumeric)
+                .take(10)
+                .collect::<String>();
+            builder.append_value(&value).unwrap();
         }
     }
     builder.finish()
 }
 
-fn create_bool_array<F>(size: usize, value_fn: F) -> BooleanArray
-where
-    F: Fn(usize) -> bool,
-{
+fn create_bool_array(size: usize, trues_density: f32) -> BooleanArray {
+    let mut rng = seedable_rng();
     let mut builder = BooleanBuilder::new(size);
-    for i in 0..size {
-        builder.append_value(value_fn(i)).unwrap();
+    for _ in 0..size {
+        let value = rng.gen::<f32>() < trues_density;
+        builder.append_value(value).unwrap();
     }
     builder.finish()
 }
 
-fn bench_filter_u8(data_array: &UInt8Array, filter_array: &BooleanArray) {
-    filter(
-        criterion::black_box(data_array),
-        criterion::black_box(filter_array),
-    )
-    .unwrap();
-}
-
-// fn bench_filter_f32(data_array: &Float32Array, filter_array: &BooleanArray) {
-//     filter(criterion::black_box(data_array), criterion::black_box(filter_array)).unwrap();
-// }
-
-fn bench_filter_context_u8(data_array: &UInt8Array, filter_context: &FilterContext) {
-    filter_context
-        .filter(criterion::black_box(data_array))
-        .unwrap();
+fn bench_filter(data_array: &dyn Array, filter_array: &BooleanArray) {
+    criterion::black_box(filter(data_array, filter_array).unwrap());
 }
 
-fn bench_filter_context_f32(data_array: &Float32Array, filter_context: &FilterContext) {
-    filter_context
-        .filter(criterion::black_box(data_array))
-        .unwrap();
+fn bench_built_filter<'a>(filter: &Filter<'a>, data: &impl Array) {
+    criterion::black_box(filter(&data.data()));
 }
 
 fn add_benchmark(c: &mut Criterion) {
     let size = 65536;
-    let filter_array = create_bool_array(size, |i| matches!(i % 2, 0));
-    let sparse_filter_array = create_bool_array(size, |i| matches!(i % 8000, 0));
-    let dense_filter_array = create_bool_array(size, |i| !matches!(i % 8000, 0));
+    let filter_array = create_bool_array(size, 0.5);
+    let dense_filter_array = create_bool_array(size, 1.0 - 1.0 / 1024.0);
+    let sparse_filter_array = create_bool_array(size, 1.0 / 1024.0);
 
-    let filter_context = FilterContext::new(&filter_array).unwrap();
-    let sparse_filter_context = FilterContext::new(&sparse_filter_array).unwrap();
-    let dense_filter_context = FilterContext::new(&dense_filter_array).unwrap();
+    let filter = build_filter(&filter_array).unwrap();
+    let dense_filter = build_filter(&dense_filter_array).unwrap();
+    let sparse_filter = build_filter(&sparse_filter_array).unwrap();
 
-    let data_array = create_primitive_array(size, |i| match i % 2 {
-        0 => 1,
-        _ => 0,
-    });
-    c.bench_function("filter u8 low selectivity", |b| {
-        b.iter(|| bench_filter_u8(&data_array, &filter_array))
+    let data_array = create_primitive_array::<UInt8Type>(size, 0.0);
+
+    c.bench_function("filter u8", |b| {
+        b.iter(|| bench_filter(&data_array, &filter_array))
     });
     c.bench_function("filter u8 high selectivity", |b| {
-        b.iter(|| bench_filter_u8(&data_array, &sparse_filter_array))
+        b.iter(|| bench_filter(&data_array, &dense_filter_array))
     });
-    c.bench_function("filter u8 very low selectivity", |b| {
-        b.iter(|| bench_filter_u8(&data_array, &dense_filter_array))
+    c.bench_function("filter u8 low selectivity", |b| {
+        b.iter(|| bench_filter(&data_array, &sparse_filter_array))
     });
 
-    c.bench_function("filter context u8 low selectivity", |b| {
-        b.iter(|| bench_filter_context_u8(&data_array, &filter_context))
+    c.bench_function("filter context u8", |b| {
+        b.iter(|| bench_built_filter(&filter, &data_array))
     });
     c.bench_function("filter context u8 high selectivity", |b| {
-        b.iter(|| bench_filter_context_u8(&data_array, &sparse_filter_context))
+        b.iter(|| bench_built_filter(&dense_filter, &data_array))
     });
-    c.bench_function("filter context u8 very low selectivity", |b| {
-        b.iter(|| bench_filter_context_u8(&data_array, &dense_filter_context))
+    c.bench_function("filter context u8 low selectivity", |b| {
+        b.iter(|| bench_built_filter(&sparse_filter, &data_array))
     });
 
-    let data_array = create_u8_array_with_nulls(size);
-    c.bench_function("filter context u8 w NULLs low selectivity", |b| {
-        b.iter(|| bench_filter_context_u8(&data_array, &filter_context))
+    let data_array = create_primitive_array::<UInt8Type>(size, 0.5);
+    c.bench_function("filter context u8 w NULLs", |b| {
+        b.iter(|| bench_built_filter(&filter, &data_array))
     });
     c.bench_function("filter context u8 w NULLs high selectivity", |b| {
-        b.iter(|| bench_filter_context_u8(&data_array, &sparse_filter_context))
+        b.iter(|| bench_built_filter(&dense_filter, &data_array))
     });
-    c.bench_function("filter context u8 w NULLs very low selectivity", |b| {
-        b.iter(|| bench_filter_context_u8(&data_array, &dense_filter_context))
+    c.bench_function("filter context u8 w NULLs low selectivity", |b| {
+        b.iter(|| bench_built_filter(&sparse_filter, &data_array))
     });
 
-    let data_array = create_primitive_array(size, |i| match i % 2 {
-        0 => 1.0,
-        _ => 0.0,
+    let data_array = create_primitive_array::<Float32Type>(size, 0.5);
+    c.bench_function("filter f32", |b| {
+        b.iter(|| bench_filter(&data_array, &filter_array))
     });
-    c.bench_function("filter context f32 low selectivity", |b| {
-        b.iter(|| bench_filter_context_f32(&data_array, &filter_context))
+    c.bench_function("filter context f32", |b| {
+        b.iter(|| bench_built_filter(&filter, &data_array))
     });
     c.bench_function("filter context f32 high selectivity", |b| {
-        b.iter(|| bench_filter_context_f32(&data_array, &sparse_filter_context))
+        b.iter(|| bench_built_filter(&dense_filter, &data_array))
+    });
+    c.bench_function("filter context f32 low selectivity", |b| {
+        b.iter(|| bench_built_filter(&sparse_filter, &data_array))
+    });
+
+    let data_array = create_string_array(size, 0.5);
+    c.bench_function("filter context string", |b| {
+        b.iter(|| bench_built_filter(&filter, &data_array))
+    });
+    c.bench_function("filter context string high selectivity", |b| {
+        b.iter(|| bench_built_filter(&dense_filter, &data_array))
     });
-    c.bench_function("filter context f32 very low selectivity", |b| {
-        b.iter(|| bench_filter_context_f32(&data_array, &dense_filter_context))
+    c.bench_function("filter context string low selectivity", |b| {
+        b.iter(|| bench_built_filter(&sparse_filter, &data_array))
     });
 }
 

rust/arrow/benches/mutable_array.rs

@@ -0,0 +1,79 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#[macro_use]
+extern crate criterion;
+use criterion::Criterion;
+
+use rand::distributions::Alphanumeric;
+use rand::Rng;
+
+use std::sync::Arc;
+
+extern crate arrow;
+
+use arrow::array::*;
+use arrow::util::test_util::seedable_rng;
+
+fn create_strings(size: usize, null_density: f32) -> ArrayRef {
+    let rng = &mut seedable_rng();
+
+    let mut builder = StringBuilder::new(size);
+    for _ in 0..size {
+        let x = rng.gen::<f32>();
+        if x < null_density {
+            let value = rng.sample_iter(&Alphanumeric).take(4).collect::<String>();
+            builder.append_value(&value).unwrap();
+        } else {
+            builder.append_null().unwrap()
+        }
+    }
+    Arc::new(builder.finish())
+}
+
+fn create_slices(size: usize) -> Vec<(usize, usize)> {
+    let rng = &mut seedable_rng();
+
+    (0..size)
+        .map(|_| {
+            let start = rng.gen_range(0, size / 2);
+            let end = rng.gen_range(start + 1, size);
+            (start, end)
+        })
+        .collect()
+}
+
+fn bench(v1: &ArrayRef, slices: &[(usize, usize)]) {
+    let mut mutable = MutableArrayData::new(vec![v1.data_ref()], false, 5);
+    for (start, end) in slices {
+        mutable.extend(0, *start, *end)
+    }
+    mutable.freeze();
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    let v1 = create_strings(1024, 0.0);
+    let v2 = create_slices(1024);
+    c.bench_function("mutable str 1024", |b| b.iter(|| bench(&v1, &v2)));
+
+    let v1 = create_strings(1024, 0.5);
+    let v2 = create_slices(1024);
+    c.bench_function("mutable str nulls 1024", |b| b.iter(|| bench(&v1, &v2)));
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

rust/arrow/src/array/transform/boolean.rs

@@ -26,7 +26,7 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend {
     let values = array.buffers()[0].data();
     Box::new(
         move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
-            let buffer = &mut mutable.buffers[0];
+            let buffer = &mut mutable.buffer1;
             reserve_for_bits(buffer, mutable.len + len);
             set_bits(
                 &mut buffer.data_mut(),
@@ -40,6 +40,6 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend {
 }
 
 pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) {
-    let buffer = &mut mutable.buffers[0];
+    let buffer = &mut mutable.buffer1;
     reserve_for_bits(buffer, mutable.len + len);
 }

rust/arrow/src/array/transform/fixed_binary.rs

@@ -30,15 +30,15 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend {
         // fast case where we can copy regions without null issues
         Box::new(
             move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
-                let buffer = &mut mutable.buffers[0];
+                let buffer = &mut mutable.buffer1;
                 buffer.extend_from_slice(&values[start * size..(start + len) * size]);
             },
         )
     } else {
         Box::new(
             move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
                 // nulls present: append item by item, ignoring null entries
-                let values_buffer = &mut mutable.buffers[0];
+                let values_buffer = &mut mutable.buffer1;
 
                 (start..start + len).for_each(|i| {
                     if array.is_valid(i) {
@@ -60,6 +60,6 @@ pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) {
         _ => unreachable!(),
     };
 
-    let values_buffer = &mut mutable.buffers[0];
+    let values_buffer = &mut mutable.buffer1;
     values_buffer.extend(len * size);
 }