fix[vortex-dict]: avoid full materialization in min_max

Prior to this commit, min_max on dictionaries would materialize the full
dictionary. This commit iterates over the codes and checks whether they fully
cover the values slice, in which case min_max can be run on the values slice.
Even if the codes do not fully cover the values, we can materialize without
duplicates, by filtering the values slice to only the covered indices.

Signed-off-by: Alfonso Subiotto Marques <[email protected]>

Author: asubiotto

Cargo.lock

@@ -20,8 +20,6 @@ arrow = ["dep:arrow-array"]
 
 [dependencies]
 arrow-array = { workspace = true, optional = true }
-arrow-buffer = { workspace = true }
-num-traits = { workspace = true }
 prost = { workspace = true }
 # test-harness
 rand = { workspace = true, optional = true }

encodings/dict/Cargo.toml

@@ -1,16 +1,140 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_array::compute::{MinMaxKernel, MinMaxKernelAdapter, MinMaxResult, min_max, take};
-use vortex_array::register_kernel;
+use vortex_array::compute::{MinMaxKernel, MinMaxKernelAdapter, MinMaxResult, mask, min_max};
+use vortex_array::{Array as _, ToCanonical, register_kernel};
+use vortex_buffer::BitBufferMut;
+use vortex_dtype::match_each_unsigned_integer_ptype;
 use vortex_error::VortexResult;
+use vortex_mask::Mask;
 
 use crate::{DictArray, DictVTable};
 
 impl MinMaxKernel for DictVTable {
     fn min_max(&self, array: &DictArray) -> VortexResult<Option<MinMaxResult>> {
-        min_max(&take(array.values(), array.codes())?)
+        let codes_validity = array.codes().validity_mask();
+        if codes_validity.all_false() {
+            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 bit buffer.
+                let mut unreferenced = BitBufferMut::new_set(values_len);
+                for (&code, is_valid) in codes_primitive.as_slice::<P>().iter().zip(validity_iter) {
+                    if is_valid {
+                        // SAFETY: code is valid, so it must be in range.
+                        #[allow(clippy::cast_possible_truncation)]
+                        unsafe {
+                            unreferenced.unset_unchecked(code as usize);
+                        }
+                    }
+                }
+
+                let unreferenced_mask = Mask::from_buffer(unreferenced.freeze());
+                min_max(&mask(array.values(), &unreferenced_mask)?)
+            })
+        })
     }
 }
 
 register_kernel!(MinMaxKernelAdapter(DictVTable).lift());
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::compute::min_max;
+    use vortex_array::{Array, IntoArray};
+    use vortex_buffer::buffer;
+
+    use crate::DictArray;
+    use crate::builders::dict_encode;
+
+    fn assert_min_max(array: &dyn Array, expected: Option<(i32, i32)>) {
+        match (min_max(array).unwrap(), expected) {
+            (Some(result), Some((expected_min, expected_max))) => {
+                assert_eq!(i32::try_from(result.min).unwrap(), expected_min);
+                assert_eq!(i32::try_from(result.max).unwrap(), expected_max);
+            }
+            (None, None) => {}
+            (got, expected) => panic!(
+                "min_max mismatch: expected {:?}, got {:?}",
+                expected,
+                got.as_ref().map(|r| (
+                    i32::try_from(r.min.clone()).ok(),
+                    i32::try_from(r.max.clone()).ok()
+                ))
+            ),
+        }
+    }
+
+    #[rstest]
+    #[case::covering(
+        DictArray::try_new(
+            buffer![0u32, 1, 2, 3, 0, 1].into_array(),
+            buffer![10i32, 20, 30, 40].into_array(),
+        ).unwrap(),
+        (10, 40)
+    )]
+    #[case::non_covering_duplicates(
+        DictArray::try_new(
+            buffer![1u32, 1, 1, 3, 3].into_array(),
+            buffer![1i32, 2, 3, 4, 5].into_array(),
+        ).unwrap(),
+        (2, 4)
+    )]
+    // Non-covering: codes with gaps
+    #[case::non_covering_gaps(
+        DictArray::try_new(
+            buffer![0u32, 2, 4].into_array(),
+            buffer![1i32, 2, 3, 4, 5].into_array(),
+        ).unwrap(),
+        (1, 5)
+    )]
+    #[case::single(dict_encode(&buffer![42i32].into_array()).unwrap(), (42, 42))]
+    #[case::nullable_codes(
+        DictArray::try_new(
+            PrimitiveArray::from_option_iter([Some(0u32), None, Some(1), Some(2)]).into_array(),
+            buffer![10i32, 20, 30].into_array(),
+        ).unwrap(),
+        (10, 30)
+    )]
+    #[case::nullable_values(
+        dict_encode(
+            PrimitiveArray::from_option_iter([Some(1i32), None, Some(2), Some(1), None]).as_ref()
+        ).unwrap(),
+        (1, 2)
+    )]
+    fn test_min_max(#[case] dict: DictArray, #[case] expected: (i32, i32)) {
+        assert_min_max(dict.as_ref(), Some(expected));
+    }
+
+    #[test]
+    fn test_sliced_dict() {
+        let reference = PrimitiveArray::from_iter([1, 5, 10, 50, 100]);
+        let dict = dict_encode(reference.as_ref()).unwrap();
+        let sliced = dict.slice(1..3);
+        assert_min_max(&sliced, Some((5, 10)));
+    }
+
+    #[rstest]
+    #[case::empty(
+        DictArray::try_new(
+            PrimitiveArray::from_iter(Vec::<u32>::new()).into_array(),
+            buffer![10i32, 20, 30].into_array(),
+        ).unwrap()
+    )]
+    #[case::all_null_codes(
+        DictArray::try_new(
+            PrimitiveArray::from_option_iter([Option::<u32>::None, None, None]).into_array(),
+            buffer![10i32, 20, 30].into_array(),
+        ).unwrap()
+    )]
+    fn test_min_max_none(#[case] dict: DictArray) {
+        assert_min_max(dict.as_ref(), None);
+    }
+}

encodings/dict/src/compute/min_max.rs

@@ -234,7 +234,7 @@ impl BitBufferMut {
     /// # Safety
     ///
     /// The caller must ensure that `index` does not exceed the largest bit index in the backing buffer.
-    unsafe fn set_unchecked(&mut self, index: usize) {
+    pub unsafe fn set_unchecked(&mut self, index: usize) {
         // SAFETY: checked by caller
         unsafe { set_bit_unchecked(self.buffer.as_mut_ptr(), self.offset + index) }
     }
@@ -246,7 +246,7 @@ 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) {
+    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) }
     }