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

@@ -109,7 +109,10 @@ impl VTable for DictVTable {
         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_with_metadata(codes, values, all_values_referenced)
+        // SAFETY: We've validated the metadata and children
+        Ok(unsafe {
+            DictArray::new_unchecked(codes, values).set_all_values_referenced(all_values_referenced)
+        })
     }
 }
 
@@ -135,11 +138,7 @@ 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,
-        all_values_referenced: bool,
-    ) -> Self {
+    pub unsafe fn new_unchecked(codes: ArrayRef, values: ArrayRef) -> Self {
         let dtype = values
             .dtype()
             .union_nullability(codes.dtype().nullability());
@@ -148,8 +147,35 @@ impl DictArray {
             values,
             stats_set: Default::default(),
             dtype,
-            all_values_referenced,
+            all_values_referenced: false,
+        }
+    }
+
+    /// Set whether all dictionary values are definitely referenced.
+    ///
+    /// # Safety
+    /// The caller must ensure that when setting `all_values_referenced = true`, ALL dictionary
+    /// values are actually referenced by at least one valid code. Setting this incorrectly can
+    /// lead to incorrect query results in operations like min/max.
+    ///
+    /// This is typically only set to `true` during dictionary encoding when we know for certain
+    /// that all values are referenced.
+    pub unsafe fn set_all_values_referenced(mut self, all_values_referenced: bool) -> Self {
+        // In debug builds, verify the claim when setting to true
+        #[cfg(debug_assertions)]
+        if all_values_referenced {
+            if let Ok(unreferenced_mask) = self.compute_unreferenced_values_mask(false) {
+                let has_unreferenced = unreferenced_mask.iter().any(|b| b);
+                debug_assert!(
+                    !has_unreferenced,
+                    "set_all_values_referenced(true) called but {} unreferenced values found",
+                    unreferenced_mask.iter().filter(|&b| b).count()
+                );
+            }
         }
+
+        self.all_values_referenced = all_values_referenced;
+        self
     }
 
     /// Build a new `DictArray` from its components, `codes` and `values`.
@@ -189,7 +215,9 @@ impl DictArray {
             vortex_bail!(MismatchedTypes: "unsigned int", codes.dtype());
         }
 
-        Ok(unsafe { Self::new_unchecked(codes, values, all_values_referenced) })
+        Ok(unsafe {
+            Self::new_unchecked(codes, values).set_all_values_referenced(all_values_referenced)
+        })
     }
 
     #[inline]
@@ -212,24 +240,56 @@ impl DictArray {
         self.all_values_referenced
     }
 
+    /// Validates that the `all_values_referenced` flag matches reality.
+    ///
+    /// Returns `Ok(())` if the flag is consistent with the actual referenced values,
+    /// or an error describing the mismatch.
+    ///
+    /// This is primarily useful for testing and debugging.
+    #[cfg(debug_assertions)]
+    pub fn validate_all_values_referenced(&self) -> VortexResult<()> {
+        let unreferenced_mask = self.compute_unreferenced_values_mask(false)?;
+        let has_unreferenced = unreferenced_mask.iter().any(|b| b);
+        let actual_all_referenced = !has_unreferenced;
+
+        if self.all_values_referenced && !actual_all_referenced {
+            let unreferenced_count = unreferenced_mask.iter().filter(|&b| b).count();
+            vortex_bail!(
+                "all_values_referenced=true but {} unreferenced values found",
+                unreferenced_count
+            );
+        }
+
+        Ok(())
+    }
+
     /// 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.
+    /// When `referenced = true`, returns a `BitBuffer` where set bits (true) correspond to
+    /// referenced values, and unset bits (false) correspond to unreferenced values.
+    ///
+    /// When `referenced = false` (default for unreferenced values), returns the inverse:
+    /// set bits (true) correspond to unreferenced values, and unset bits (false) correspond
+    /// to referenced values.
     ///
     /// This is useful for operations like min/max that need to ignore unreferenced values.
-    pub fn compute_unreferenced_values_mask(&self) -> VortexResult<BitBuffer> {
+    pub fn compute_unreferenced_values_mask(&self, referenced: bool) -> 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];
+        // Initialize with the starting value: false for referenced, true for unreferenced
+        let init_value = !referenced;
+        // Value to set when we find a referenced code: true for referenced, false for unreferenced
+        let referenced_value = referenced;
+
+        let mut values_vec = vec![init_value; 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;
+                        values_vec[code as usize] = referenced_value;
                     }
                 });
             }
@@ -240,15 +300,13 @@ impl DictArray {
 
                     #[allow(clippy::cast_possible_truncation)]
                     buf.set_indices().for_each(|idx| {
-                        unreferenced_vec[codes[idx] as usize] = false;
+                        values_vec[codes[idx] as usize] = referenced_value;
                     })
                 });
             }
         }
 
-        Ok(BitBuffer::collect_bool(values_len, |idx| {
-            unreferenced_vec[idx]
-        }))
+        Ok(BitBuffer::collect_bool(values_len, |idx| values_vec[idx]))
     }
 }
 

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

@@ -14,7 +14,6 @@ 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
-        // We conservatively set all_values_referenced to false since Arrow doesn't provide this info
-        unsafe { DictArray::new_unchecked(keys, values, false) }
+        unsafe { DictArray::new_unchecked(keys, values) }
     }
 }

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

@@ -24,13 +24,11 @@ impl CastKernel for DictVTable {
         };
 
         // SAFETY: casting does not alter invariants of the codes
+        // Preserve all_values_referenced since casting only changes values, not which are referenced
         Ok(Some(
             unsafe {
-                DictArray::new_unchecked(
-                    casted_codes,
-                    casted_values,
-                    array.has_all_values_referenced(),
-                )
+                DictArray::new_unchecked(casted_codes, casted_values)
+                    .set_all_values_referenced(array.has_all_values_referenced())
             }
             .into_array(),
         ))

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

@@ -30,12 +30,9 @@ 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,
-                    lhs.has_all_values_referenced(),
-                )
-                .into_array()
+                DictArray::new_unchecked(lhs.codes().clone(), compare_result)
+                    .set_all_values_referenced(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

@@ -43,10 +43,11 @@ impl FillNullKernel for DictVTable {
 
         // SAFETY: invariants are still satisfied after patching nulls
         unsafe {
-            Ok(
-                DictArray::new_unchecked(codes, values, array.has_all_values_referenced())
-                    .into_array(),
-            )
+            Ok(DictArray::new_unchecked(codes, values)
+                // Preserve all_values_referenced since filling nulls cannot make values
+                // unreferenced.
+                .set_all_values_referenced(array.has_all_values_referenced())
+                .into_array())
         }
     }
 }

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

@@ -28,12 +28,9 @@ impl LikeKernel for DictVTable {
             // Preserve all_values_referenced since codes are unchanged
             unsafe {
                 Ok(Some(
-                    DictArray::new_unchecked(
-                        array.codes().clone(),
-                        values,
-                        array.has_all_values_referenced(),
-                    )
-                    .into_array(),
+                    DictArray::new_unchecked(array.codes().clone(), values)
+                        .set_all_values_referenced(array.has_all_values_referenced())
+                        .into_array(),
                 ))
             }
         } else {

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

@@ -21,7 +21,7 @@ impl MinMaxKernel for DictVTable {
         }
 
         // Slow path: compute which values are unreferenced and mask them out
-        let unreferenced_mask = Mask::from_buffer(array.compute_unreferenced_values_mask()?);
+        let unreferenced_mask = Mask::from_buffer(array.compute_unreferenced_values_mask(false)?);
         min_max(&mask(array.values(), &unreferenced_mask)?)
     }
 }

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

@@ -22,8 +22,7 @@ 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
-        // 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())
+        Ok(unsafe { DictArray::new_unchecked(codes, array.values().clone()) }.into_array())
     }
 }
 
@@ -34,8 +33,7 @@ impl FilterKernel for DictVTable {
         let codes = filter(array.codes(), mask)?;
 
         // SAFETY: filtering codes doesn't change invariants
-        // 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()) }
+        unsafe { Ok(DictArray::new_unchecked(codes, array.values().clone()).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(), false).into_array() }
+        unsafe { DictArray::new_unchecked(sliced_code, array.values().clone()).into_array() }
     }
 
     fn scalar_at(array: &DictArray, index: usize) -> Scalar {

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

@@ -53,14 +53,13 @@ pub fn dict_encode_with_constraints(
     let mut encoder = dict_encoder(array, constraints);
     let codes = encoder.encode(array).to_primitive().narrow()?;
     // SAFETY: The encoding process will produce a value set of codes and values
+    // 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
     unsafe {
-        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,
-        ))
+        Ok(
+            DictArray::new_unchecked(codes.into_array(), encoder.reset())
+                .set_all_values_referenced(true),
+        )
     }
 }