Optimize max_boolean by operating on u64 chunks

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

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