dict: avoid full materialization in min_max

Prior to this commit, min_max on dictionaries would materialize the full
dictionary. This commit deduplicates 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, which should make take()ing the values slice cheaper.

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

Author: asubiotto

encodings/dict/Cargo.toml

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

encodings/dict/src/compute/min_max.rs

@@ -1,16 +1,148 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use itertools::Itertools;
+use num_traits::Zero;
 use vortex_array::compute::{MinMaxKernel, MinMaxKernelAdapter, MinMaxResult, min_max, take};
-use vortex_array::register_kernel;
+use vortex_array::{Array as _, ToCanonical, register_kernel};
+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 Mask::AllTrue(_) = array.codes().validity_mask() else {
+            return min_max(&take(array.values(), array.codes().as_ref())?);
+        };
+
+        let codes_primitive = array.codes().to_primitive();
+        let values_len = array.values().len();
+        match_each_unsigned_integer_ptype!(codes_primitive.ptype(), |P| {
+            let mut codes = codes_primitive
+                .as_slice::<P>()
+                .iter()
+                // First dedup of consecutive codes before collecting should
+                // make this operation a little cheaper.
+                .dedup()
+                .copied()
+                .collect::<Vec<_>>();
+            codes.sort_unstable();
+            codes.dedup();
+
+            let max_expected = P::try_from(values_len - 1).ok();
+            if codes.len() == values_len
+                && codes.first() == Some(&P::zero())
+                && codes.last() == max_expected.as_ref()
+            {
+                // Codes fully cover all values, compute min_max directly on
+                // values.
+                return min_max(array.values());
+            }
+
+            // Codes do not fully cover the values, we need to take only
+            // referenced values, but can do so using only the unique codes,
+            // avoiding a full materialization.
+            min_max(&take(array.values(), array.codes().as_ref())?)
+        })
     }
 }
 
 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);
+    }
+}