perf[dict]: all_values_referenced to allow pushdown optimisation.

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

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

@@ -32,6 +32,11 @@ pub struct DictMetadata {
     // nullable codes are optional since they were added after stabilisation
     #[prost(optional, bool, tag = "3")]
     pub(super) is_nullable_codes: Option<bool>,
+    // all_values_referenced is optional for backward compatibility
+    // true = all dictionary values are definitely referenced by at least one code
+    // false/None = unknown whether all values are referenced (conservative default)
+    #[prost(optional, bool, tag = "4")]
+    pub(super) all_values_referenced: Option<bool>,
 }
 
 impl VTable for DictVTable {
@@ -66,6 +71,7 @@ impl VTable for DictVTable {
                 )
             })?,
             is_nullable_codes: Some(array.codes().dtype().is_nullable()),
+            all_values_referenced: Some(array.all_values_referenced),
         }))
     }
 
@@ -101,8 +107,9 @@ impl VTable for DictVTable {
         let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable);
         let codes = children.get(0, &codes_dtype, len)?;
         let values = children.get(1, dtype, metadata.values_len as usize)?;
+        let all_values_referenced = metadata.all_values_referenced.unwrap_or(false);
 
-        DictArray::try_new(codes, values)
+        DictArray::try_new_with_metadata(codes, values, all_values_referenced)
     }
 }
 
@@ -112,6 +119,10 @@ pub struct DictArray {
     values: ArrayRef,
     stats_set: ArrayStats,
     dtype: DType,
+    /// Indicates whether all dictionary values are definitely referenced by at least one code.
+    /// `true` = all values are referenced (computed during encoding).
+    /// `false` = unknown/might have unreferenced values (conservative default).
+    all_values_referenced: bool,
 }
 
 #[derive(Clone, Debug)]
@@ -124,7 +135,11 @@ impl DictArray {
     /// This should be called only when you can guarantee the invariants checked
     /// by the safe [`DictArray::try_new`] constructor are valid, for example when
     /// you are filtering or slicing an existing valid `DictArray`.
-    pub unsafe fn new_unchecked(codes: ArrayRef, values: ArrayRef) -> Self {
+    pub unsafe fn new_unchecked(
+        codes: ArrayRef,
+        values: ArrayRef,
+        all_values_referenced: bool,
+    ) -> Self {
         let dtype = values
             .dtype()
             .union_nullability(codes.dtype().nullability());
@@ -133,6 +148,7 @@ impl DictArray {
             values,
             stats_set: Default::default(),
             dtype,
+            all_values_referenced,
         }
     }
 
@@ -156,11 +172,24 @@ impl DictArray {
     ///
     /// It is an error to provide a nullable `codes` with non-nullable `values`.
     pub fn try_new(codes: ArrayRef, values: ArrayRef) -> VortexResult<Self> {
+        Self::try_new_with_metadata(codes, values, false)
+    }
+
+    /// Build a new `DictArray` from its components with explicit metadata.
+    ///
+    /// Same as [`DictArray::try_new`] but allows specifying whether all values are referenced.
+    /// This is typically only set to `true` during dictionary encoding when we know for certain
+    /// that all dictionary values are referenced by at least one code.
+    pub fn try_new_with_metadata(
+        codes: ArrayRef,
+        values: ArrayRef,
+        all_values_referenced: bool,
+    ) -> VortexResult<Self> {
         if !codes.dtype().is_unsigned_int() {
             vortex_bail!(MismatchedTypes: "unsigned int", codes.dtype());
         }
 
-        Ok(unsafe { Self::new_unchecked(codes, values) })
+        Ok(unsafe { Self::new_unchecked(codes, values, all_values_referenced) })
     }
 
     #[inline]
@@ -173,6 +202,16 @@ impl DictArray {
         &self.values
     }
 
+    /// Returns `true` if all dictionary values are definitely referenced by at least one code.
+    ///
+    /// When `true`, operations like min/max can safely operate on all values without needing to
+    /// compute which values are actually referenced. When `false`, it is unknown whether all
+    /// values are referenced (conservative default).
+    #[inline]
+    pub fn has_all_values_referenced(&self) -> bool {
+        self.all_values_referenced
+    }
+
     /// 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
@@ -467,6 +506,7 @@ mod test {
                 codes_ptype: PType::U64 as i32,
                 values_len: u32::MAX,
                 is_nullable_codes: None,
+                all_values_referenced: None,
             }),
         );
     }

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

@@ -14,6 +14,7 @@ impl<K: ArrowDictionaryKeyType> FromArrowArray<&DictionaryArray<K>> for DictArra
         let keys = ArrayRef::from_arrow(keys, keys.is_nullable());
         let values = ArrayRef::from_arrow(array.values().as_ref(), nullable);
         // SAFETY: we assume that Arrow has checked the invariants on construction
-        unsafe { DictArray::new_unchecked(keys, values) }
+        // We conservatively set all_values_referenced to false since Arrow doesn't provide this info
+        unsafe { DictArray::new_unchecked(keys, values, false) }
     }
 }

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

@@ -25,7 +25,14 @@ impl CastKernel for DictVTable {
 
         // SAFETY: casting does not alter invariants of the codes
         Ok(Some(
-            unsafe { DictArray::new_unchecked(casted_codes, casted_values) }.into_array(),
+            unsafe {
+                DictArray::new_unchecked(
+                    casted_codes,
+                    casted_values,
+                    array.has_all_values_referenced(),
+                )
+            }
+            .into_array(),
         ))
     }
 }

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

@@ -30,7 +30,12 @@ impl CompareKernel for DictVTable {
 
             // SAFETY: values len preserved, codes all still point to valid values
             let result = unsafe {
-                DictArray::new_unchecked(lhs.codes().clone(), compare_result).into_array()
+                DictArray::new_unchecked(
+                    lhs.codes().clone(),
+                    compare_result,
+                    lhs.has_all_values_referenced(),
+                )
+                .into_array()
             };
 
             // We canonicalize the result because dictionary-encoded bools is dumb.

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

@@ -42,7 +42,12 @@ impl FillNullKernel for DictVTable {
         let values = fill_null(array.values(), fill_value)?;
 
         // SAFETY: invariants are still satisfied after patching nulls
-        unsafe { Ok(DictArray::new_unchecked(codes, values).into_array()) }
+        unsafe {
+            Ok(
+                DictArray::new_unchecked(codes, values, array.has_all_values_referenced())
+                    .into_array(),
+            )
+        }
     }
 }
 

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

@@ -25,9 +25,15 @@ impl LikeKernel for DictVTable {
 
             // SAFETY: LIKE preserves the len of the values, so codes are still pointing at
             //  valid positions.
+            // Preserve all_values_referenced since codes are unchanged
             unsafe {
                 Ok(Some(
-                    DictArray::new_unchecked(array.codes().clone(), values).into_array(),
+                    DictArray::new_unchecked(
+                        array.codes().clone(),
+                        values,
+                        array.has_all_values_referenced(),
+                    )
+                    .into_array(),
                 ))
             }
         } else {

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

@@ -15,6 +15,12 @@ impl MinMaxKernel for DictVTable {
             return Ok(None);
         }
 
+        // Fast path: if all values are referenced, directly compute min/max on values
+        if array.has_all_values_referenced() {
+            return min_max(array.values());
+        }
+
+        // Slow path: compute which values are unreferenced and mask them out
         let unreferenced_mask = Mask::from_buffer(array.compute_unreferenced_values_mask()?);
         min_max(&mask(array.values(), &unreferenced_mask)?)
     }

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

@@ -22,7 +22,8 @@ impl TakeKernel for DictVTable {
     fn take(&self, array: &DictArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
         let codes = take(array.codes(), indices)?;
         // SAFETY: selecting codes doesn't change the invariants of DictArray
-        Ok(unsafe { DictArray::new_unchecked(codes, array.values().clone()) }.into_array())
+        // We conservatively set all_values_referenced to false since taking may leave values unreferenced
+        Ok(unsafe { DictArray::new_unchecked(codes, array.values().clone(), false) }.into_array())
     }
 }
 
@@ -33,7 +34,8 @@ impl FilterKernel for DictVTable {
         let codes = filter(array.codes(), mask)?;
 
         // SAFETY: filtering codes doesn't change invariants
-        unsafe { Ok(DictArray::new_unchecked(codes, array.values().clone()).into_array()) }
+        // We conservatively set all_values_referenced to false since filtering may leave values unreferenced
+        unsafe { Ok(DictArray::new_unchecked(codes, array.values().clone(), false).into_array()) }
     }
 }
 

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

@@ -24,7 +24,7 @@ impl OperationsVTable<DictVTable> for DictVTable {
             };
         }
         // SAFETY: slicing the codes preserves invariants
-        unsafe { DictArray::new_unchecked(sliced_code, array.values().clone()).into_array() }
+        unsafe { DictArray::new_unchecked(sliced_code, array.values().clone(), false).into_array() }
     }
 
     fn scalar_at(array: &DictArray, index: usize) -> Scalar {

vortex-array/src/builders/dict/mod.rs

@@ -57,6 +57,9 @@ pub fn dict_encode_with_constraints(
         Ok(DictArray::new_unchecked(
             codes.into_array(),
             encoder.reset(),
+            // All values in the dictionary are guaranteed to be referenced by at least one code
+            // since we build the dictionary from the codes we observe during encoding
+            true,
         ))
     }
 }