This is an automated email from the ASF dual-hosted git repository.
dheres pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 5de1d5e741 Optimize `max_boolean` by operating on u64 chunks (#6098)
5de1d5e741 is described below
commit 5de1d5e74170902224c715a8ef3bb595abb069b3
Author: Simon Vandel Sillesen <[email protected]>
AuthorDate: Mon Jul 22 17:17:49 2024 +0000
Optimize `max_boolean` by operating on u64 chunks (#6098)
* 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.
* add tests exercising u64 chunk code
---
arrow-arith/src/aggregate.rs | 120 +++++++++++++++++++++++++++++++++++--
arrow/benches/aggregate_kernels.rs | 47 +++++++++++++++
2 files changed, 162 insertions(+), 5 deletions(-)
diff --git a/arrow-arith/src/aggregate.rs b/arrow-arith/src/aggregate.rs
index 0e4d31eee7..ae4a5b008a 100644
--- a/arrow-arith/src/aggregate.rs
+++ b/arrow-arith/src/aggregate.rs
@@ -371,11 +371,26 @@ pub fn max_boolean(array: &BooleanArray) -> Option<bool> {
}
// 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`].
@@ -748,6 +763,7 @@ where
mod tests {
use super::*;
use arrow_array::types::*;
+ use builder::BooleanBuilder;
use std::sync::Arc;
#[test]
@@ -1318,6 +1334,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]
@@ -1339,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]);
diff --git a/arrow/benches/aggregate_kernels.rs
b/arrow/benches/aggregate_kernels.rs
index 1e7b9f894f..71cd242784 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);
