From 18193b30d69e07b68ce0efd09c749992b4f6144d Mon Sep 17 00:00:00 2001 From: Simon Vandel Sillesen Date: Sat, 20 Jul 2024 10:59:47 +0000 Subject: [PATCH 1/2] Optimize `max_boolean` Operate on bit chunks instead of individual booleans, which can lead to massive speedups while not regressing the short-circuiting behavior of the existing implementation. `cargo bench --bench aggregate_kernels -- "bool/max"` shows throughput improvements between 50% to 23390% on my machine. --- arrow-arith/src/aggregate.rs | 33 +++++++++++++++++---- arrow/benches/aggregate_kernels.rs | 47 ++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 5 deletions(-) diff --git a/arrow-arith/src/aggregate.rs b/arrow-arith/src/aggregate.rs index 0e4d31eee7a1..a218ee28922e 100644 --- a/arrow-arith/src/aggregate.rs +++ b/arrow-arith/src/aggregate.rs @@ -371,11 +371,26 @@ pub fn max_boolean(array: &BooleanArray) -> Option { } // Note the max bool is true (1), so short circuit as soon as we see it - array - .iter() - .find(|&b| b == Some(true)) - .flatten() - .or(Some(false)) + match array.nulls() { + None => array + .values() + .bit_chunks() + .iter_padded() + // We found a true if any bit is set + .map(|x| x != 0) + .find(|b| *b) + .or(Some(false)), + Some(nulls) => { + let validity_chunks = nulls.inner().bit_chunks().iter_padded(); + let value_chunks = array.values().bit_chunks().iter_padded(); + value_chunks + .zip(validity_chunks) + // We found a true if the value bit is 1, AND the validity bit is 1 for any bits in the chunk + .map(|(value_bits, validity_bits)| (value_bits & validity_bits) != 0) + .find(|b| *b) + .or(Some(false)) + } + } } /// Helper to compute min/max of [`ArrayAccessor`]. @@ -1318,6 +1333,14 @@ mod tests { let a = BooleanArray::from(vec![Some(false), Some(true), None, Some(false), None]); assert_eq!(Some(false), min_boolean(&a)); assert_eq!(Some(true), max_boolean(&a)); + + let a = BooleanArray::from(vec![Some(true), None]); + assert_eq!(Some(true), min_boolean(&a)); + assert_eq!(Some(true), max_boolean(&a)); + + let a = BooleanArray::from(vec![Some(false), None]); + assert_eq!(Some(false), min_boolean(&a)); + assert_eq!(Some(false), max_boolean(&a)); } #[test] diff --git a/arrow/benches/aggregate_kernels.rs b/arrow/benches/aggregate_kernels.rs index 1e7b9f894f2a..71cd24278424 100644 --- a/arrow/benches/aggregate_kernels.rs +++ b/arrow/benches/aggregate_kernels.rs @@ -66,6 +66,53 @@ fn add_benchmark(c: &mut Criterion) { .bench_function("min nullable", |b| b.iter(|| min_string(&nullable_strings))) .bench_function("max nullable", |b| b.iter(|| max_string(&nullable_strings))); } + + { + let nonnull_bools_mixed = create_boolean_array(BATCH_SIZE, 0.0, 0.5); + let nonnull_bools_all_false = create_boolean_array(BATCH_SIZE, 0.0, 0.0); + let nonnull_bools_all_true = create_boolean_array(BATCH_SIZE, 0.0, 1.0); + let nullable_bool_mixed = create_boolean_array(BATCH_SIZE, 0.5, 0.5); + let nullable_bool_all_false = create_boolean_array(BATCH_SIZE, 0.5, 0.0); + let nullable_bool_all_true = create_boolean_array(BATCH_SIZE, 0.5, 1.0); + c.benchmark_group("bool") + .throughput(Throughput::Elements(BATCH_SIZE as u64)) + .bench_function("min nonnull mixed", |b| { + b.iter(|| min_boolean(&nonnull_bools_mixed)) + }) + .bench_function("max nonnull mixed", |b| { + b.iter(|| max_boolean(&nonnull_bools_mixed)) + }) + .bench_function("min nonnull false", |b| { + b.iter(|| min_boolean(&nonnull_bools_all_false)) + }) + .bench_function("max nonnull false", |b| { + b.iter(|| max_boolean(&nonnull_bools_all_false)) + }) + .bench_function("min nonnull true", |b| { + b.iter(|| min_boolean(&nonnull_bools_all_true)) + }) + .bench_function("max nonnull true", |b| { + b.iter(|| max_boolean(&nonnull_bools_all_true)) + }) + .bench_function("min nullable mixed", |b| { + b.iter(|| min_boolean(&nullable_bool_mixed)) + }) + .bench_function("max nullable mixed", |b| { + b.iter(|| max_boolean(&nullable_bool_mixed)) + }) + .bench_function("min nullable false", |b| { + b.iter(|| min_boolean(&nullable_bool_all_false)) + }) + .bench_function("max nullable false", |b| { + b.iter(|| max_boolean(&nullable_bool_all_false)) + }) + .bench_function("min nullable true", |b| { + b.iter(|| min_boolean(&nullable_bool_all_true)) + }) + .bench_function("max nullable true", |b| { + b.iter(|| max_boolean(&nullable_bool_all_true)) + }); + } } criterion_group!(benches, add_benchmark); From 5c40f28ec4716e65b9802e8d06941d2f093e3b8c Mon Sep 17 00:00:00 2001 From: Simon Vandel Sillesen Date: Sun, 21 Jul 2024 17:41:13 +0000 Subject: [PATCH 2/2] add tests exercising u64 chunk code --- arrow-arith/src/aggregate.rs | 87 ++++++++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/arrow-arith/src/aggregate.rs b/arrow-arith/src/aggregate.rs index a218ee28922e..ae4a5b008a83 100644 --- a/arrow-arith/src/aggregate.rs +++ b/arrow-arith/src/aggregate.rs @@ -763,6 +763,7 @@ where mod tests { use super::*; use arrow_array::types::*; + use builder::BooleanBuilder; use std::sync::Arc; #[test] @@ -1362,6 +1363,92 @@ mod tests { assert_eq!(Some(true), max_boolean(&a)); } + #[test] + fn test_boolean_min_max_64_true_64_false() { + let mut no_nulls = BooleanBuilder::new(); + no_nulls.append_slice(&[true; 64]); + no_nulls.append_slice(&[false; 64]); + let no_nulls = no_nulls.finish(); + + assert_eq!(Some(false), min_boolean(&no_nulls)); + assert_eq!(Some(true), max_boolean(&no_nulls)); + + let mut with_nulls = BooleanBuilder::new(); + with_nulls.append_slice(&[true; 31]); + with_nulls.append_null(); + with_nulls.append_slice(&[true; 32]); + with_nulls.append_slice(&[false; 1]); + with_nulls.append_nulls(63); + let with_nulls = with_nulls.finish(); + + assert_eq!(Some(false), min_boolean(&with_nulls)); + assert_eq!(Some(true), max_boolean(&with_nulls)); + } + + #[test] + fn test_boolean_min_max_64_false_64_true() { + let mut no_nulls = BooleanBuilder::new(); + no_nulls.append_slice(&[false; 64]); + no_nulls.append_slice(&[true; 64]); + let no_nulls = no_nulls.finish(); + + assert_eq!(Some(false), min_boolean(&no_nulls)); + assert_eq!(Some(true), max_boolean(&no_nulls)); + + let mut with_nulls = BooleanBuilder::new(); + with_nulls.append_slice(&[false; 31]); + with_nulls.append_null(); + with_nulls.append_slice(&[false; 32]); + with_nulls.append_slice(&[true; 1]); + with_nulls.append_nulls(63); + let with_nulls = with_nulls.finish(); + + assert_eq!(Some(false), min_boolean(&with_nulls)); + assert_eq!(Some(true), max_boolean(&with_nulls)); + } + + #[test] + fn test_boolean_min_max_96_true() { + let mut no_nulls = BooleanBuilder::new(); + no_nulls.append_slice(&[true; 96]); + let no_nulls = no_nulls.finish(); + + assert_eq!(Some(true), min_boolean(&no_nulls)); + assert_eq!(Some(true), max_boolean(&no_nulls)); + + let mut with_nulls = BooleanBuilder::new(); + with_nulls.append_slice(&[true; 31]); + with_nulls.append_null(); + with_nulls.append_slice(&[true; 32]); + with_nulls.append_slice(&[true; 31]); + with_nulls.append_null(); + let with_nulls = with_nulls.finish(); + + assert_eq!(Some(true), min_boolean(&with_nulls)); + assert_eq!(Some(true), max_boolean(&with_nulls)); + } + + #[test] + fn test_boolean_min_max_96_false() { + let mut no_nulls = BooleanBuilder::new(); + no_nulls.append_slice(&[false; 96]); + let no_nulls = no_nulls.finish(); + + assert_eq!(Some(false), min_boolean(&no_nulls)); + assert_eq!(Some(false), max_boolean(&no_nulls)); + + let mut with_nulls = BooleanBuilder::new(); + with_nulls.append_slice(&[false; 31]); + with_nulls.append_null(); + with_nulls.append_slice(&[false; 32]); + with_nulls.append_slice(&[false; 31]); + with_nulls.append_null(); + let with_nulls = with_nulls.finish(); + + assert_eq!(Some(false), min_boolean(&with_nulls)); + assert_eq!(Some(false), max_boolean(&with_nulls)); + } + #[test] fn test_sum_dyn() { let values = Int8Array::from_iter_values([10_i8, 11, 12, 13, 14, 15, 16, 17]);