Pushdown some expressions to Dict layout reader

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

vortex-array/src/expr/analysis/annotation.rs

@@ -47,6 +47,23 @@ pub fn descendent_annotations<A: AnnotationFn>(
     visitor.annotations
 }
 
+pub fn annotations<A: AnnotationFn>(
+    expr: &Expression,
+    annotate: A,
+) -> Annotations<'_, A::Annotation> {
+    let mut visitor = DirectAnnotationVisitor {
+        annotations: Default::default(),
+        annotate,
+    };
+    expr.accept(&mut visitor).vortex_expect("Infallible");
+    visitor.annotations
+}
+
+struct DirectAnnotationVisitor<'a, A: AnnotationFn> {
+    annotations: Annotations<'a, A::Annotation>,
+    annotate: A,
+}
+
 struct AnnotationVisitor<'a, A: AnnotationFn> {
     annotations: Annotations<'a, A::Annotation>,
     annotate: A,
@@ -84,3 +101,19 @@ impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> {
         Ok(TraversalOrder::Continue)
     }
 }
+
+impl<'a, A: AnnotationFn> NodeVisitor<'a> for DirectAnnotationVisitor<'a, A> {
+    type NodeTy = Expression;
+    fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
+        let annotations = (self.annotate)(node);
+        if annotations.is_empty() {
+            Ok(TraversalOrder::Continue)
+        } else {
+            self.annotations
+                .entry(node)
+                .or_default()
+                .extend(annotations);
+            Ok(TraversalOrder::Skip)
+        }
+    }
+}

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;
@@ -19,13 +20,50 @@ use crate::expr::analysis::Annotation;
 use crate::expr::analysis::AnnotationFn;
 use crate::expr::analysis::Annotations;
 use crate::expr::analysis::descendent_annotations;
+use crate::expr::annotations;
 use crate::expr::get_item;
+use crate::expr::is_root;
+use crate::expr::label_tree;
 use crate::expr::pack;
 use crate::expr::root;
 use crate::expr::traversal::NodeExt;
 use crate::expr::traversal::NodeRewriter;
 use crate::expr::traversal::Transformed;
 use crate::expr::traversal::TraversalOrder;
+use crate::scalar_fn::is_negative_cost;
+
+/// Split expression into two parts:
+///
+/// left is the outer part that we want to apply to array after canonicalizing.
+/// right is the optional inner part that we want to apply to array before
+/// canonicalizing.
+///
+/// We want to push to array only if expression has a negative cost, is
+/// infallible and null-insensitive.
+pub fn split_expression_for_pushdown(
+    expr: Expression,
+    dtype: &DType,
+) -> VortexResult<(Expression, Option<Expression>)> {
+    let references_root = label_tree(&expr, is_root, |acc, &child| acc | child);
+    let annotations = annotations(&expr, |expr| {
+        let signature = expr.signature();
+        if !signature.is_fallible()
+            && !signature.is_null_sensitive()
+            && is_negative_cost(expr.id())
+            && references_root.get(&expr).copied().unwrap_or(true)
+        {
+            vec![""] // we want a single annotation so an empty literal is fine
+        } else {
+            vec![]
+        }
+    });
+    let partition = partition_annotations(expr.clone(), dtype, annotations)?;
+    if partition.partitions.is_empty() {
+        Ok((partition.root, None))
+    } else {
+        Ok((partition.root, Some(partition.partitions[0].clone())))
+    }
+}
 
 /// Partition an expression into sub-expressions that are uniquely associated with an annotation.
 /// A root expression is also returned that can be used to recombine the results of the partitions
@@ -49,11 +87,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());
@@ -205,10 +254,16 @@ where
 mod tests {
     use rstest::fixture;
     use rstest::rstest;
+    use vortex_array::expr::Expression;
+    use vortex_array::expr::byte_length;
+    use vortex_array::expr::cast;
+    use vortex_array::expr::like;
 
+    use super::split_expression_for_pushdown;
     use super::*;
     use crate::dtype::DType;
     use crate::dtype::Nullability::NonNullable;
+    use crate::dtype::PType;
     use crate::dtype::PType::I32;
     use crate::dtype::StructFields;
     use crate::expr::analysis::make_free_field_annotator;
@@ -348,4 +403,57 @@ mod tests {
         let expected_b = pack([("b_0", pack([("b", col("b"))], NonNullable))], NonNullable);
         assert_eq!(part_b, &expected_b, "{part_b} {expected_b}");
     }
+
+    fn pushed_inner(exprs: impl IntoIterator<Item = Expression>) -> Expression {
+        pack(
+            exprs
+                .into_iter()
+                .enumerate()
+                .map(|(idx, e)| (format!("_{idx}"), e)),
+            NonNullable,
+        )
+    }
+
+    fn pushed_ref(idx: usize) -> Expression {
+        get_item(format!("_{idx}"), get_item("", root()))
+    }
+
+    #[test]
+    fn split_expr_root() {
+        let (outer, inner) = split_expression_for_pushdown(root(), &DType::Null).unwrap();
+        assert_eq!(outer, root());
+        assert_eq!(inner, None);
+    }
+
+    #[test]
+    fn split_expr_partial_pushdown() {
+        // cast is fallible, thus not pushed
+        let target = DType::Primitive(PType::I64, NonNullable);
+        let expr = cast(byte_length(root()), target.clone());
+        let (outer, inner) =
+            split_expression_for_pushdown(expr, &DType::Utf8(false.into())).unwrap();
+        // [0] = cast([1], dtype)
+        // [1] = byte_length(root)
+        assert_eq!(outer, cast(pushed_ref(0), target));
+        assert_eq!(inner, Some(pushed_inner([byte_length(root())])));
+    }
+
+    #[test]
+    fn split_expr_full_pushdown() {
+        let expr = byte_length(root());
+        let (outer, inner) =
+            split_expression_for_pushdown(expr, &DType::Utf8(false.into())).unwrap();
+        assert_eq!(outer, pushed_ref(0));
+        assert_eq!(inner, Some(pushed_inner([byte_length(root())])));
+    }
+
+    #[test]
+    fn split_expr_no_pushdown() {
+        // like is fallible, thus not pushed. lit() does not reference root()
+        let expr = like(root(), lit(1u64));
+        let (outer, inner) =
+            split_expression_for_pushdown(expr.clone(), &DType::Utf8(true.into())).unwrap();
+        assert_eq!(outer, expr);
+        assert_eq!(inner, None);
+    }
 }

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,16 @@ 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()
+pub fn is_negative_cost(id: ScalarFnId) -> bool {
+    id == ScalarFnVTable::id(&ByteLength)
+        || id == ScalarFnVTable::id(&GetItem)
+        || id == ScalarFnVTable::id(&Literal)
+}

vortex-layout/src/layouts/dict/reader.rs

@@ -20,6 +20,7 @@ use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldMask;
 use vortex_array::expr::Expression;
 use vortex_array::expr::root;
+use vortex_array::expr::transform::split_expression_for_pushdown;
 use vortex_array::optimizer::ArrayOptimizer;
 use vortex_error::VortexError;
 use vortex_error::VortexExpect;
@@ -100,10 +101,7 @@ impl DictReader {
                     )
                     .vortex_expect("must construct dict values array evaluation")
                     .map_err(Arc::new)
-                    .map(move |array| {
-                        let array = array?;
-                        Ok(SharedArray::new(array).into_array())
-                    })
+                    .map(move |array| Ok(SharedArray::new(array?).into_array()))
                     .boxed()
                     .shared()
             })
@@ -229,13 +227,18 @@ impl LayoutReader for DictReader {
         mask: MaskFuture,
     ) -> VortexResult<BoxFuture<'static, VortexResult<ArrayRef>>> {
         // TODO: fix up expr partitioning with fallible & null sensitive annotations
-        let values_eval = self.values_array();
         let codes_eval = self
             .codes
             .projection_evaluation(row_range, &root(), mask)
             .map_err(|err| err.with_context("While evaluating projection on codes"))?;
-        let expr = expr.clone();
 
+        let (expr_outer, expr_inner) = split_expression_for_pushdown(expr.clone(), self.dtype())?;
+
+        let values_eval = if let Some(inner) = expr_inner {
+            self.values_eval(inner)
+        } else {
+            self.values_array()
+        };
         let all_values_referenced = self.layout.has_all_values_referenced();
         Ok(async move {
             let (values, codes) = try_join!(values_eval.map_err(VortexError::from), codes_eval)?;
@@ -252,7 +255,7 @@ impl LayoutReader for DictReader {
             .into_array()
             .optimize()?;
 
-            array.apply(&expr)
+            array.apply(&expr_outer)
         }
         .boxed())
     }