Filter DecimalArray as PrimitiveArray ~80% Faster (#2637) (#2870)

* Filter DecimalArray as PrimitiveArray (#2637) * Add decimal filter benches * Format
apache · Oct 15, 2022 · ede36d7 · ede36d7
1 parent f055f51
commit ede36d7
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/arrow/benches/filter_kernels.rs b/arrow/benches/filter_kernels.rs
@@ -26,6 +26,7 @@ use arrow::array::*;
 use arrow::compute::filter;
 use arrow::datatypes::{Field, Float32Type, Int32Type, Schema, UInt8Type};
 
+use arrow_array::types::Decimal128Type;
 use criterion::{criterion_group, criterion_main, Criterion};
 
 fn bench_filter(data_array: &dyn Array, filter_array: &BooleanArray) {
@@ -143,6 +144,29 @@ fn add_benchmark(c: &mut Criterion) {
         b.iter(|| bench_built_filter(&sparse_filter, &data_array))
     });
 
+    let data_array = create_primitive_array::<Decimal128Type>(size, 0.0);
+    c.bench_function("filter decimal128 (kept 1/2)", |b| {
+        b.iter(|| bench_filter(&data_array, &filter_array))
+    });
+    c.bench_function("filter decimal128 high selectivity (kept 1023/1024)", |b| {
+        b.iter(|| bench_filter(&data_array, &dense_filter_array))
+    });
+    c.bench_function("filter decimal128 low selectivity (kept 1/1024)", |b| {
+        b.iter(|| bench_filter(&data_array, &sparse_filter_array))
+    });
+
+    c.bench_function("filter context decimal128 (kept 1/2)", |b| {
+        b.iter(|| bench_built_filter(&filter, &data_array))
+    });
+    c.bench_function(
+        "filter context decimal128 high selectivity (kept 1023/1024)",
+        |b| b.iter(|| bench_built_filter(&dense_filter, &data_array)),
+    );
+    c.bench_function(
+        "filter context decimal128 low selectivity (kept 1/1024)",
+        |b| b.iter(|| bench_built_filter(&sparse_filter, &data_array)),
+    );
+
     let data_array = create_string_array::<i32>(size, 0.5);
     c.bench_function("filter context string (kept 1/2)", |b| {
         b.iter(|| bench_built_filter(&filter, &data_array))

diff --git a/arrow/src/compute/kernels/filter.rs b/arrow/src/compute/kernels/filter.rs
@@ -338,6 +338,16 @@ fn filter_array(values: &dyn Array, predicate: &FilterPredicate) -> Result<Array
         // actually filter
         _ => downcast_primitive_array! {
             values => Ok(Arc::new(filter_primitive(values, predicate))),
+            DataType::Decimal128(p, s) => {
+                let values = values.as_any().downcast_ref::<Decimal128Array>().unwrap();
+                let filtered = filter_primitive(values, predicate);
+                Ok(Arc::new(filtered.with_precision_and_scale(*p, *s).unwrap()))
+            }
+            DataType::Decimal256(p, s) => {
+                let values = values.as_any().downcast_ref::<Decimal256Array>().unwrap();
+                let filtered = filter_primitive(values, predicate);
+                Ok(Arc::new(filtered.with_precision_and_scale(*p, *s).unwrap()))
+            }
             DataType::Boolean => {
                 let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
                 Ok(Arc::new(filter_boolean(values, predicate)))