diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index 5c013b2d321..8bea82e6132 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -32,11 +32,13 @@ impl TakeExecute for Dict { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let codes = array.codes().take(indices.clone())?; - // SAFETY: selecting codes doesn't change the invariants of DictArray - // Preserve all_values_referenced since taking codes doesn't affect which values are referenced - Ok(Some(unsafe { - DictArray::new_unchecked(codes, array.values().clone()).into_array() - })) + + // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code + // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because + // selection can remove a value's last reference. + let taken = unsafe { DictArray::new_unchecked(codes, array.values().clone()) }; + + Ok(Some(taken.into_array())) } } @@ -44,11 +46,12 @@ impl FilterReduce for Dict { fn filter(array: ArrayView<'_, Dict>, mask: &Mask) -> VortexResult> { let codes = array.codes().filter(mask.clone())?; - // SAFETY: filtering codes doesn't change invariants - // Preserve all_values_referenced since filtering codes doesn't affect which values are referenced - Ok(Some(unsafe { - DictArray::new_unchecked(codes, array.values().clone()).into_array() - })) + // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code + // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because + // selection can remove a value's last reference. + let filtered = unsafe { DictArray::new_unchecked(codes, array.values().clone()) }; + + Ok(Some(filtered.into_array())) } } diff --git a/vortex-array/src/arrays/dict/compute/rules.rs b/vortex-array/src/arrays/dict/compute/rules.rs index cfb467dbba5..d43c6deeb75 100644 --- a/vortex-array/src/arrays/dict/compute/rules.rs +++ b/vortex-array/src/arrays/dict/compute/rules.rs @@ -106,101 +106,104 @@ impl ArrayParentReduceRule for DictionaryScalarFnValuesPushDownRule { parent: ArrayView<'_, ScalarFn>, child_idx: usize, ) -> VortexResult> { - // Check that the scalar function can actually be pushed down. - let sig = parent.scalar_fn().signature(); + let scalar_fn = parent.scalar_fn(); + let signature = scalar_fn.signature(); - // Don't push down pack expressions since we might want to unpack them in exporters - // later. - if parent.scalar_fn().is::() { + // Preserve pack expressions so exporters can unpack them later. + if scalar_fn.is::() { return Ok(None); } - // Don't push down cast operations — CastReduceAdaptor handles these eagerly. - // If it declined (returned None), we must fall through to the canonical path - // rather than creating a lazy cast inside the dictionary values. - if parent.scalar_fn().is::() { + // CastReduceAdaptor handles casts eagerly. If it declines the rewrite, leave the cast on + // the dictionary instead of creating a lazy cast over its values. + if scalar_fn.is::() { return Ok(None); } - // If the dictionary has less codes than values don't push down this might - // happen if the dictionary is sliced. + // A sliced dictionary can have more values than code rows. Do not increase the work in + // that case. if array.values().len() > array.codes().len() { return Ok(None); } - // If the scalar function is fallible, we cannot push it down since it may fail over a - // value that isn't referenced by any code. - if !array.all_values_referenced && !sig.is_infallible() { + // A fallible function could fail on an unreferenced value that row-wise evaluation would + // never visit. + if !array.has_all_values_referenced() && !signature.is_infallible() { tracing::trace!( "Not pushing down fallible scalar function {} over dictionary with sparse codes {}", - parent.scalar_fn(), + scalar_fn, Dict.id(), ); return Ok(None); } - // Check that all siblings are constant - // TODO(ngates): we can also support other dictionaries if the values are the same! - if !parent + // TODO(ngates): Support dictionary siblings when their values match. + let other_children_are_constant = parent .iter_children() .enumerate() - .all(|(idx, c)| idx == child_idx || c.is::()) - { + .all(|(idx, child)| idx == child_idx || child.is::()); + if !other_children_are_constant { return Ok(None); } // Before this rewrite, a null code supplies null for this argument while the constant // arguments retain their values. After the rewrite, the null code masks the function's // result. Those are equivalent only for a strict function. - if array.codes().dtype().is_nullable() + let codes_have_nulls = array.codes().dtype().is_nullable() && !matches!( array.codes().validity()?, Validity::NonNullable | Validity::AllValid - ) - && !sig.is_strict() - { + ); + if codes_have_nulls && !signature.is_strict() { tracing::trace!( "Not pushing down non-strict scalar function {} over dictionary with null codes {}", - parent.scalar_fn(), + scalar_fn, Dict.id(), ); return Ok(None); } - // Now we push the parent scalar function into the dictionary values. let values_len = array.values().len(); - let mut new_children = Vec::with_capacity(parent.nchildren()); + let mut value_children = Vec::with_capacity(parent.nchildren()); for (idx, child) in parent.iter_children().enumerate() { if idx == child_idx { - new_children.push(array.values().clone()); + value_children.push(array.values().clone()); } else { let scalar = child.as_::().scalar().clone(); - new_children.push(ConstantArray::new(scalar, values_len).into_array()); + value_children.push(ConstantArray::new(scalar, values_len).into_array()); } } - let new_values = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)? + let transformed_values = ScalarFnArray::try_new(scalar_fn.clone(), value_children)? .into_array() .optimize()?; // A non-strict function reaches this point only when the codes are all valid, but their // dtype may still be nullable. Remove that declared nullability while rebuilding the // dictionary, then cast its output to the function's declared dtype. - if !sig.is_strict() && array.codes().dtype().is_nullable() { - let new_codes = array.codes().cast(array.codes().dtype().as_nonnullable())?; - let new_dict = unsafe { - DictArray::new_unchecked(new_codes, new_values) + if !signature.is_strict() && array.codes().dtype().is_nullable() { + let non_nullable_codes = array.codes().cast(array.codes().dtype().as_nonnullable())?; + + // SAFETY: The validity guard proves that the codes contain no nulls. Removing their + // declared nullability preserves every code, and `transformed_values` has one entry + // for each original dictionary value. + let transformed_dict = unsafe { + DictArray::new_unchecked(non_nullable_codes, transformed_values) .set_all_values_referenced(array.has_all_values_referenced()) } .into_array(); - return Ok(Some(new_dict.cast(parent.dtype().clone())?)); + + return Ok(Some(transformed_dict.cast(parent.dtype().clone())?)); } - Ok(Some(unsafe { - DictArray::new_unchecked(array.codes().clone(), new_values) + // SAFETY: The codes are unchanged and `transformed_values` has one entry for each original + // dictionary value, so code bounds and `all_values_referenced` remain unchanged. + let transformed_dict = unsafe { + DictArray::new_unchecked(array.codes().clone(), transformed_values) .set_all_values_referenced(array.has_all_values_referenced()) - .into_array() - })) + }; + + Ok(Some(transformed_dict.into_array())) } }