Pushdown some expressions to Dict layout reader (#8341)

myrrc · web-flow · commit c9c8fb7bb85b · 2026-06-17T10:21:13.000Z
When we access values of Dict layout reader, it canonicalizes them and stores them in a SharedArray. This means we always pay the cost of canonicalization which in turn means we can't do #8310 . In order to solve this issue, we need to apply some expressions to the values array before canonicalizing it. However, we can't push down arbitrary expressions as it may be beneficial to apply them over canonicalized array. One example of such expressions is LIKE over a Dict array with few codes used. Applying LIKE to whole values array is not beneficial. This PR adds a hardcoded internal is_negative_cost estimation for expressions that we want to push before canonicalization. A hint for these are expressions which don't depend on individual input size. As an example, for every string, len(string) doesn't read the string itself but reads the metadata and thus is O(1) on individual input. We don't push down fallible (like cast) or null sensitive (like IS NULL) expressions as well because we want to propagate the errors at call site rather than upfront. Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>
diff --git a/vortex-array/src/expr/analysis/annotation.rs b/vortex-array/src/expr/analysis/annotation.rs
@@ -42,6 +42,25 @@ pub fn descendent_annotations<A: AnnotationFn>(
     let mut visitor = AnnotationVisitor {
         annotations: Default::default(),
         annotate,
+        propagate_up: true,
+    };
+    expr.accept(&mut visitor).vortex_expect("Infallible");
+    visitor.annotations
+}
+
+/// Walk the expression tree and annotate each expression with zero or more
+/// annotations.
+///
+/// Returns a map of each expression to all annotations. Annotations of
+/// children are not propagated to parents.
+pub fn direct_annotations<A: AnnotationFn>(
+    expr: &Expression,
+    annotate: A,
+) -> Annotations<'_, A::Annotation> {
+    let mut visitor = AnnotationVisitor {
+        annotations: Default::default(),
+        annotate,
+        propagate_up: false,
     };
     expr.accept(&mut visitor).vortex_expect("Infallible");
     visitor.annotations
@@ -50,6 +69,7 @@ pub fn descendent_annotations<A: AnnotationFn>(
 struct AnnotationVisitor<'a, A: AnnotationFn> {
     annotations: Annotations<'a, A::Annotation>,
     annotate: A,
+    propagate_up: bool,
 }
 
 impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> {
@@ -70,6 +90,9 @@ impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> {
     }
 
     fn visit_up(&mut self, node: &'a Expression) -> VortexResult<TraversalOrder> {
+        if !self.propagate_up {
+            return Ok(TraversalOrder::Continue);
+        }
         let child_annotations = node
             .children()
             .iter()
diff --git a/vortex-array/src/expr/transform/partition.rs b/vortex-array/src/expr/transform/partition.rs
@@ -3,6 +3,7 @@
 
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::hash::Hash;
 
 use itertools::Itertools;
 use vortex_error::VortexExpect;
@@ -49,11 +50,22 @@ where
 {
     // Annotate each expression with the annotations that any of its descendent expressions have.
     let annotations = descendent_annotations(&expr, annotate_fn);
+    partition_annotations(expr.clone(), scope, annotations)
+}
 
+pub fn partition_annotations<A>(
+    expr: Expression,
+    scope: &DType,
+    annotations: Annotations<A>,
+) -> VortexResult<PartitionedExpr<A>>
+where
+    A: Display + Clone + Eq + Hash,
+    FieldName: From<A>,
+{
     // Now we split the original expression into sub-expressions based on the annotations, and
     // generate a root expression to re-assemble the results.
-    let mut splitter = StructFieldExpressionSplitter::<A::Annotation>::new(&annotations);
-    let root = expr.clone().rewrite(&mut splitter)?.value;
+    let mut splitter = StructFieldExpressionSplitter::<A>::new(&annotations);
+    let root = expr.rewrite(&mut splitter)?.value;
 
     let mut partitions = Vec::with_capacity(splitter.sub_expressions.len());
     let mut partition_annotations = Vec::with_capacity(splitter.sub_expressions.len());
diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs
@@ -9,6 +9,10 @@
 
 use vortex_session::registry::Id;
 
+use crate::scalar_fn::fns::byte_length::ByteLength;
+use crate::scalar_fn::fns::get_item::GetItem;
+use crate::scalar_fn::fns::literal::Literal;
+
 mod vtable;
 pub use vtable::*;
 
@@ -48,3 +52,17 @@ mod sealed {
     /// This can be the **only** implementor for [`super::typed::DynScalarFn`].
     impl<V: ScalarFnVTable> Sealed for TypedScalarFnInstance<V> {}
 }
+
+/// A scalar function has a negative cost if applying it to an array and
+/// canonicalizing is cheaper than canonicalizing an array and applying it.
+///
+/// Example of negative cost expressions are byte_length() and get_item() since
+/// they don't depend on input size.
+///
+/// Example of non-negative cost expression is like() as it's linear over
+/// individual input.
+pub fn is_negative_cost(id: ScalarFnId) -> bool {
+    id == ScalarFnVTable::id(&ByteLength)
+        || id == ScalarFnVTable::id(&GetItem)
+        || id == ScalarFnVTable::id(&Literal)
+}
diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs
