fix[vortex-dict]: faster min_max for dict array (#5288)

Signed-off-by: Alfonso Subiotto Marques <[email protected]>
Signed-off-by: Joe Isaacs <[email protected]>
Signed-off-by: Robert Kruszewski <[email protected]>
Co-authored-by: Alfonso Subiotto Marques <[email protected]>
Co-authored-by: Robert Kruszewski <[email protected]>

Author: 3 people

vortex-array/Cargo.toml

@@ -159,3 +159,7 @@ required-features = ["test-harness"]
 [[bench]]
 name = "dict_mask"
 harness = false
+
+[[bench]]
+name = "dict_unreferenced_mask"
+harness = false

vortex-array/benches/dict_unreferenced_mask.rs

@@ -0,0 +1,111 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![allow(clippy::unwrap_used)]
+
+use divan::Bencher;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use vortex_array::IntoArray;
+use vortex_array::arrays::{DictArray, PrimitiveArray};
+use vortex_array::compute::warm_up_vtables;
+
+fn main() {
+    warm_up_vtables();
+    divan::main();
+}
+
+/// Benchmark with many codes (65K) relative to 1024 values.
+/// This tests performance when the values dictionary is small but many codes reference it.
+#[divan::bench(args = [
+    1024,    // Small dictionary
+    2048,    // Medium dictionary
+    4096,    // Larger dictionary
+])]
+fn bench_many_codes_few_values(bencher: Bencher, num_values: i32) {
+    let mut rng = StdRng::seed_from_u64(0);
+
+    let num_codes = 65_536;
+
+    // Create values array with the specified number of unique values
+    let values = PrimitiveArray::from_iter(0..num_values).into_array();
+
+    // Create codes that randomly reference the values
+    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
+    let codes = PrimitiveArray::from_iter(
+        (0..num_codes).map(|_| rng.random_range(0..num_values as usize) as u32),
+    )
+    .into_array();
+
+    let array = DictArray::try_new(codes, values).unwrap();
+
+    bencher
+        .with_inputs(|| array.clone())
+        .bench_values(|array| array.compute_unreferenced_values_mask().unwrap());
+}
+
+/// Benchmark with many nulls in the codes array.
+/// This tests performance when most codes are null and thus don't reference values.
+#[divan::bench(args = [
+    0.01,   // 1% valid codes
+    0.1,    // 10% valid codes
+    0.5,    // 50% valid codes
+    0.9,    // 90% valid codes
+])]
+fn bench_many_nulls(bencher: Bencher, fraction_valid: f64) {
+    let mut rng = StdRng::seed_from_u64(0);
+
+    let num_codes = 65_536;
+    let num_values = 1024i32;
+
+    // Create values array
+    let values = PrimitiveArray::from_iter(0..num_values).into_array();
+
+    // Create codes with many nulls based on fraction_valid
+    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
+    let codes = PrimitiveArray::from_option_iter((0..num_codes).map(|_| {
+        rng.random_bool(fraction_valid)
+            .then(|| rng.random_range(0..num_values as usize) as u32)
+    }))
+    .into_array();
+
+    let array = DictArray::try_new(codes, values).unwrap();
+
+    bencher
+        .with_inputs(|| array.clone())
+        .bench_values(|array| array.compute_unreferenced_values_mask().unwrap());
+}
+
+/// Benchmark with sparse code coverage (many unreferenced values).
+/// This tests when only a small subset of values are actually referenced.
+#[divan::bench(args = [
+    0.01,   // Only 1% of values are referenced
+    0.1,    // 10% of values referenced
+    0.5,    // 50% of values referenced
+])]
+fn bench_sparse_coverage(bencher: Bencher, fraction_coverage: f64) {
+    let mut rng = StdRng::seed_from_u64(0);
+
+    let num_codes = 65_536;
+    let num_values = 1024i32;
+
+    // Create values array
+    let values = PrimitiveArray::from_iter(0..num_values).into_array();
+
+    // Calculate how many unique values we'll actually reference
+    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
+    let num_referenced = (num_values as f64 * fraction_coverage).max(1.0) as usize;
+
+    // Create codes that only reference a subset of values
+    #[allow(clippy::cast_possible_truncation)]
+    let codes = PrimitiveArray::from_iter(
+        (0..num_codes).map(|_| rng.random_range(0..num_referenced) as u32),
+    )
+    .into_array();
+
+    let array = DictArray::try_new(codes, values).unwrap();
+
+    bencher
+        .with_inputs(|| array.clone())
+        .bench_values(|array| array.compute_unreferenced_values_mask().unwrap());
+}

vortex-array/src/arrays/dict/array.rs

@@ -106,6 +106,45 @@ impl DictArray {
     pub fn values(&self) -> &ArrayRef {
         &self.values
     }
+
+    /// Compute a mask indicating which values in the dictionary are referenced by at least one code.
+    ///
+    /// Returns a `BitBuffer` where unset bits (false) correspond to values that are referenced
+    /// by at least one valid code, and set bits (true) correspond to unreferenced values.
+    ///
+    /// This is useful for operations like min/max that need to ignore unreferenced values.
+    pub fn compute_unreferenced_values_mask(&self) -> VortexResult<BitBuffer> {
+        let codes_validity = self.codes().validity_mask();
+        let codes_primitive = self.codes().to_primitive();
+        let values_len = self.values().len();
+
+        let mut unreferenced_vec = vec![true; values_len];
+        match codes_validity.bit_buffer() {
+            AllOr::All => {
+                match_each_integer_ptype!(codes_primitive.ptype(), |P| {
+                    #[allow(clippy::cast_possible_truncation)]
+                    for &code in codes_primitive.as_slice::<P>().iter() {
+                        unreferenced_vec[code as usize] = false;
+                    }
+                });
+            }
+            AllOr::None => {}
+            AllOr::Some(buf) => {
+                match_each_integer_ptype!(codes_primitive.ptype(), |P| {
+                    let codes = codes_primitive.as_slice::<P>();
+
+                    #[allow(clippy::cast_possible_truncation)]
+                    buf.set_indices().for_each(|idx| {
+                        unreferenced_vec[codes[idx] as usize] = false;
+                    })
+                });
+            }
+        }
+
+        Ok(BitBuffer::collect_bool(values_len, |idx| {
+            unreferenced_vec[idx]
+        }))
+    }
 }
 
 impl ArrayVTable<DictVTable> for DictVTable {

vortex-array/src/arrays/dict/compute/min_max.rs

@@ -1,14 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_buffer::BitBuffer;
-use vortex_dtype::match_each_unsigned_integer_ptype;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use super::{DictArray, DictVTable};
 use crate::compute::{MinMaxKernel, MinMaxKernelAdapter, MinMaxResult, mask, min_max};
-use crate::{Array as _, ToCanonical, register_kernel};
+use crate::{Array as _, register_kernel};
 
 impl MinMaxKernel for DictVTable {
     fn min_max(&self, array: &DictArray) -> VortexResult<Option<MinMaxResult>> {
@@ -17,25 +15,8 @@ impl MinMaxKernel for DictVTable {
             return Ok(None);
         }
 
-        let codes_primitive = array.codes().to_primitive();
-        let values_len = array.values().len();
-        match_each_unsigned_integer_ptype!(codes_primitive.ptype(), |P| {
-            codes_validity.iter_bools(|validity_iter| {
-                // mask() sets values to null where the mask is true, so we start
-                // with a fully-set bool buffer.
-                let mut unreferenced = vec![true; values_len];
-                #[allow(clippy::cast_possible_truncation)]
-                for (&code, is_valid) in codes_primitive.as_slice::<P>().iter().zip(validity_iter) {
-                    if is_valid {
-                        unreferenced[code as usize] = false;
-                    }
-                }
-
-                let unreferenced_mask =
-                    Mask::from_buffer(BitBuffer::collect_bool(values_len, |i| unreferenced[i]));
-                min_max(&mask(array.values(), &unreferenced_mask)?)
-            })
-        })
+        let unreferenced_mask = Mask::from_buffer(array.compute_unreferenced_values_mask()?);
+        min_max(&mask(array.values(), &unreferenced_mask)?)
     }
 }
 

vortex-buffer/src/bit/buf_mut.rs

@@ -268,6 +268,7 @@ impl BitBufferMut {
     /// Set a position to `false`.
     ///
     /// This operation is checked so if `index` exceeds the buffer length, this will panic.
+    #[inline]
     pub fn unset(&mut self, index: usize) {
         assert!(index < self.len, "index {index} exceeds len {}", self.len);
 
@@ -294,7 +295,8 @@ impl BitBufferMut {
     /// # Safety
     ///
     /// The caller must ensure that `index` does not exceed the largest bit index in the backing buffer.
-    unsafe fn unset_unchecked(&mut self, index: usize) {
+    #[inline]
+    pub unsafe fn unset_unchecked(&mut self, index: usize) {
         // SAFETY: checked by caller
         unsafe { unset_bit_unchecked(self.buffer.as_mut_ptr(), self.offset + index) }
     }