fix: Handle leading nulls when computing stats in compressor (#2608)

Author: robert3005

vortex-btrblocks/src/float/stats.rs

@@ -1,5 +1,6 @@
 use std::hash::Hash;
 
+use itertools::Itertools;
 use num_traits::Float;
 use rustc_hash::FxBuildHasher;
 use vortex_array::aliases::hash_map::HashMap;
@@ -9,6 +10,7 @@ use vortex_array::{Array, ToCanonical};
 use vortex_dtype::half::f16;
 use vortex_dtype::{NativePType, PType};
 use vortex_error::{VortexExpect, VortexUnwrap, vortex_panic};
+use vortex_mask::AllOr;
 
 use crate::sample::sample;
 use crate::{CompressorStats, GenerateStatsOptions};
@@ -138,33 +140,57 @@ where
     let value_count = validity.true_count();
     let mut min = T::max_value();
     let mut max = T::min_value();
-    let mut distinct_values_count: u32 = if count_distinct_values { 0 } else { u32::MAX };
+    // Keep a HashMap of T, then convert the keys into PValue afterward since value is
+    // so much more efficient to hash and search for.
     let mut distinct_values = if count_distinct_values {
         HashMap::with_capacity_and_hasher(array.len() / 2, FxBuildHasher)
     } else {
         HashMap::with_hasher(FxBuildHasher)
     };
 
-    // Keep a HashMap of T, then convert the keys into PValue afterward since value is
-    // so much more efficient to hash and search for.
     let mut runs = 1;
-    let mut prev = array.as_slice::<T>()[0];
-
-    for (idx, &value) in array.buffer::<T>().iter().enumerate() {
-        if validity.value(idx) {
-            min = min.min(value);
-            max = max.max(value);
-
-            if count_distinct_values {
-                *distinct_values.entry(value.to_bits()).or_insert(0) += 1;
-                distinct_values_count = distinct_values.len().try_into().vortex_unwrap();
-            } else {
-                distinct_values_count = u32::MAX;
+    let head_idx = validity
+        .first()
+        .vortex_expect("All null masks have been handled before");
+    let buff = array.buffer::<T>();
+    let mut prev = buff[head_idx];
+
+    let first_valid_buff = buff.slice(head_idx..array.len());
+    match validity.boolean_buffer() {
+        AllOr::All => {
+            for value in first_valid_buff {
+                min = min.min(value);
+                max = max.max(value);
+
+                if count_distinct_values {
+                    *distinct_values.entry(value.to_bits()).or_insert(0) += 1;
+                }
+
+                if value != prev {
+                    prev = value;
+                    runs += 1;
+                }
             }
+        }
+        AllOr::None => unreachable!("All invalid arrays have been handled earlier"),
+        AllOr::Some(v) => {
+            for (&value, valid) in first_valid_buff
+                .iter()
+                .zip_eq(v.slice(head_idx, array.len() - head_idx).iter())
+            {
+                if valid {
+                    min = min.min(value);
+                    max = max.max(value);
+
+                    if count_distinct_values {
+                        *distinct_values.entry(value.to_bits()).or_insert(0) += 1;
+                    }
 
-            if value != prev {
-                prev = value;
-                runs += 1;
+                    if value != prev {
+                        prev = value;
+                        runs += 1;
+                    }
+                }
             }
         }
     }
@@ -175,6 +201,11 @@ where
     let value_count = value_count
         .try_into()
         .vortex_expect("null_count must fit in u32");
+    let distinct_values_count = if count_distinct_values {
+        distinct_values.len().try_into().vortex_unwrap()
+    } else {
+        u32::MAX
+    };
 
     FloatStats {
         null_count,
@@ -191,6 +222,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::validity::Validity;
     use vortex_array::{IntoArray, ToCanonical};
     use vortex_buffer::buffer;
 
@@ -209,4 +242,19 @@ mod tests {
         assert_eq!(stats.average_run_length, 1);
         assert_eq!(stats.distinct_values_count, 3);
     }
+
+    #[test]
+    fn test_float_stats_leading_nulls() {
+        let floats = PrimitiveArray::new(
+            buffer![0.0f32, 1.0f32, 2.0f32],
+            Validity::from_iter([false, true, true]),
+        );
+
+        let stats = FloatStats::generate(&floats);
+
+        assert_eq!(stats.value_count, 2);
+        assert_eq!(stats.null_count, 1);
+        assert_eq!(stats.average_run_length, 1);
+        assert_eq!(stats.distinct_values_count, 2);
+    }
 }

vortex-btrblocks/src/integer/stats.rs

@@ -9,6 +9,7 @@ use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ToCanonical};
 use vortex_dtype::{NativePType, match_each_integer_ptype};
 use vortex_error::{VortexExpect, VortexUnwrap};
+use vortex_mask::AllOr;
 use vortex_scalar::PValue;
 
 use crate::sample::sample;
@@ -196,7 +197,12 @@ where
     let value_count = validity.true_count();
 
     // Initialize loop state
-    let head = array.as_slice::<T>()[0];
+    let head_idx = validity
+        .first()
+        .vortex_expect("All null masks have been handled before");
+    let buffer = array.buffer::<T>();
+    let head = buffer[head_idx];
+
     let mut loop_state = LoopState {
         min: head,
         max: head,
@@ -205,43 +211,55 @@ where
         } else {
             HashMap::with_hasher(FxBuildHasher)
         },
-        distinct_values_count: if count_distinct_values { 0 } else { u32::MAX },
         prev: head,
         runs: 1,
     };
 
-    let values = array.buffer::<T>();
-    let mask = validity.to_boolean_buffer();
-
-    let mut offset = 0;
-    for chunk in values.as_slice().chunks(64) {
-        let validity = mask.slice(offset, chunk.len());
-        offset += chunk.len();
-
-        if chunk.len() < 64 {
-            // Final iteration, run naive loop
-            inner_loop_naive(chunk, count_distinct_values, &validity, &mut loop_state);
-            break;
-        }
-
-        let set_bits = validity.count_set_bits();
-
-        match set_bits {
-            // All nulls -> no stats to update
-            0 => continue,
-            // Inner loop for when validity check can be elided
-            64 => inner_loop_nonnull(
-                chunk.try_into().vortex_unwrap(),
+    let sliced = buffer.slice(head_idx..array.len());
+    let mut chunks = sliced.as_slice().array_chunks::<64>();
+    match validity.boolean_buffer() {
+        AllOr::All => {
+            for chunk in &mut chunks {
+                inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state)
+            }
+            let remainder = chunks.remainder();
+            inner_loop_naive(
+                remainder,
                 count_distinct_values,
+                &BooleanBuffer::new_set(remainder.len()),
                 &mut loop_state,
-            ),
-            // Inner loop for when we need to check validity
-            _ => inner_loop_nullable(
-                chunk.try_into().vortex_unwrap(),
+            );
+        }
+        AllOr::None => unreachable!("All invalid arrays have been handled before"),
+        AllOr::Some(v) => {
+            let mask = v.slice(head_idx, array.len() - head_idx);
+            let mut offset = 0;
+            for chunk in &mut chunks {
+                let validity = mask.slice(offset, 64);
+                offset += 64;
+
+                match validity.count_set_bits() {
+                    // All nulls -> no stats to update
+                    0 => continue,
+                    // Inner loop for when validity check can be elided
+                    64 => inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state),
+                    // Inner loop for when we need to check validity
+                    _ => inner_loop_nullable(
+                        chunk,
+                        count_distinct_values,
+                        &validity,
+                        &mut loop_state,
+                    ),
+                }
+            }
+            // Final iteration, run naive loop
+            let remainder = chunks.remainder();
+            inner_loop_naive(
+                remainder,
                 count_distinct_values,
-                &validity,
+                &mask.slice(offset, remainder.len()),
                 &mut loop_state,
-            ),
+            );
         }
     }
 
@@ -257,7 +275,11 @@ where
     };
 
     let runs = loop_state.runs;
-    let distinct_values_count = loop_state.distinct_values_count;
+    let distinct_values_count = if count_distinct_values {
+        loop_state.distinct_values.len().try_into().vortex_unwrap()
+    } else {
+        u32::MAX
+    };
 
     let typed = TypedStats {
         min: loop_state.min,
@@ -289,7 +311,6 @@ struct LoopState<T> {
     max: T,
     prev: T,
     runs: u32,
-    distinct_values_count: u32,
     distinct_values: HashMap<T, u32, FxBuildHasher>,
 }
 
@@ -305,7 +326,6 @@ fn inner_loop_nonnull<T: PrimInt + Hash>(
 
         if count_distinct_values {
             *state.distinct_values.entry(value).or_insert(0) += 1;
-            state.distinct_values_count = state.distinct_values.len().try_into().vortex_unwrap();
         }
 
         if value != state.prev {
@@ -329,8 +349,6 @@ fn inner_loop_nullable<T: PrimInt + Hash>(
 
             if count_distinct_values {
                 *state.distinct_values.entry(value).or_insert(0) += 1;
-                state.distinct_values_count =
-                    state.distinct_values.len().try_into().vortex_unwrap();
             }
 
             if value != state.prev {
@@ -355,8 +373,6 @@ fn inner_loop_naive<T: PrimInt + Hash>(
 
             if count_distinct_values {
                 *state.distinct_values.entry(value).or_insert(0) += 1;
-                state.distinct_values_count =
-                    state.distinct_values.len().try_into().vortex_unwrap();
             }
 
             if value != state.prev {
@@ -376,6 +392,8 @@ mod tests {
     use vortex_array::validity::Validity;
     use vortex_buffer::{Buffer, buffer};
 
+    use crate::CompressorStats;
+    use crate::integer::IntegerStats;
     use crate::integer::stats::typed_int_stats;
 
     #[test]
@@ -413,4 +431,16 @@ mod tests {
         let stats = typed_int_stats::<u8>(&array, true);
         assert_eq!(stats.distinct_values_count, 64);
     }
+
+    #[test]
+    fn test_integer_stats_leading_nulls() {
+        let ints = PrimitiveArray::new(buffer![0, 1, 2], Validity::from_iter([false, true, true]));
+
+        let stats = IntegerStats::generate(&ints);
+
+        assert_eq!(stats.value_count, 2);
+        assert_eq!(stats.null_count, 1);
+        assert_eq!(stats.average_run_length, 1);
+        assert_eq!(stats.distinct_values_count, 2);
+    }
 }

vortex-btrblocks/src/lib.rs

@@ -1,3 +1,5 @@
+#![feature(array_chunks)]
+
 use std::fmt::Debug;
 use std::hash::Hash;