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

rust/arrow/Cargo.toml

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

rust/arrow/benches/filter_kernels.rs

@@ -14,128 +14,141 @@
 // 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: &UInt8Array, 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 context f32", |b| {
+        b.iter(|| bench_built_filter(&filter, &data_array))
+    });
+    c.bench_function("filter context f32 high selectivity", |b| {
+        b.iter(|| bench_built_filter(&dense_filter, &data_array))
     });
     c.bench_function("filter context f32 low selectivity", |b| {
-        b.iter(|| bench_filter_context_f32(&data_array, &filter_context))
+        b.iter(|| bench_built_filter(&sparse_filter, &data_array))
     });
-    c.bench_function("filter context f32 high selectivity", |b| {
-        b.iter(|| bench_filter_context_f32(&data_array, &sparse_filter_context))
+
+    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: &Vec<(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/array_boolean.rs

@@ -78,8 +78,8 @@ impl BooleanArray {
     /// Returns a `Buffer` holding all the values of this array.
     ///
     /// Note this doesn't take the offset of this array into account.
-    pub fn values(&self) -> Buffer {
-        self.data.buffers()[0].clone()
+    pub fn values(&self) -> &Buffer {
+        &self.data.buffers()[0]
     }
 
     /// Returns the boolean value at index `i`.
@@ -245,7 +245,7 @@ mod tests {
     fn test_boolean_array_from_vec() {
         let buf = Buffer::from([10_u8]);
         let arr = BooleanArray::from(vec![false, true, false, true]);
-        assert_eq!(buf, arr.values());
+        assert_eq!(&buf, arr.values());
         assert_eq!(4, arr.len());
         assert_eq!(0, arr.offset());
         assert_eq!(0, arr.null_count());
@@ -260,7 +260,7 @@ mod tests {
     fn test_boolean_array_from_vec_option() {
         let buf = Buffer::from([10_u8]);
         let arr = BooleanArray::from(vec![Some(false), Some(true), None, Some(true)]);
-        assert_eq!(buf, arr.values());
+        assert_eq!(&buf, arr.values());
         assert_eq!(4, arr.len());
         assert_eq!(0, arr.offset());
         assert_eq!(1, arr.null_count());
@@ -288,7 +288,7 @@ mod tests {
             .add_buffer(buf)
             .build();
         let arr = BooleanArray::from(data);
-        assert_eq!(buf2, arr.values());
+        assert_eq!(&buf2, arr.values());
         assert_eq!(5, arr.len());
         assert_eq!(2, arr.offset());
         assert_eq!(0, arr.null_count());

rust/arrow/src/array/builder.rs

@@ -2618,7 +2618,7 @@ mod tests {
         }
 
         let arr = builder.finish();
-        assert_eq!(buf, arr.values());
+        assert_eq!(&buf, arr.values());
         assert_eq!(10, arr.len());
         assert_eq!(0, arr.offset());
         assert_eq!(0, arr.null_count());

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);
 }