refactor: propagate bound expressions through layouts

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

@@ -2,21 +2,16 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexExpect;
-use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::dtype::FieldName;
 use crate::dtype::StructFields;
 use crate::expr::BoundExpression;
 use crate::expr::Expression;
 use crate::expr::analysis::AnnotationFn;
-use crate::expr::analysis::Annotations;
-use crate::expr::descendent_annotations;
 use crate::scalar_fn::fns::get_item::GetItem;
 use crate::scalar_fn::fns::root::Root;
 use crate::scalar_fn::fns::select::Select;
 
-pub type FieldAccesses<'a> = Annotations<'a, FieldName>;
-
 /// Returns the "free fields" for this expression node.
 ///
 /// A "free field" is a top-level field from the root scope that this expression references—not
@@ -92,28 +87,3 @@ pub fn make_bound_free_field_annotator(
         vec![]
     }
 }
-
-/// For all subexpressions in an expression, find the fields that are accessed directly from the
-/// scope, but not any fields in those fields
-/// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = root().a.b + root().d accesses {a,d} (not b).
-///
-/// Note: This is a very naive, but simple analysis to find the fields that are accessed directly on an
-/// identity node. This is combined to provide an over-approximation of the fields that are accessed
-/// by an expression.
-pub fn immediate_scope_accesses<'a>(
-    expr: &'a Expression,
-    scope: &'a StructFields,
-) -> FieldAccesses<'a> {
-    descendent_annotations(expr, make_free_field_annotator(scope))
-}
-
-/// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`.
-pub fn immediate_scope_access<'a>(
-    expr: &'a Expression,
-    scope: &'a StructFields,
-) -> HashSet<FieldName> {
-    immediate_scope_accesses(expr, scope)
-        .get(expr)
-        .vortex_expect("Expression missing from scope accesses, this is a internal bug")
-        .clone()
-}

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

@@ -4,17 +4,15 @@
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 
-use crate::dtype::DType;
 use crate::dtype::Field;
 use crate::dtype::FieldPath;
 use crate::dtype::FieldPathSet;
-use crate::expr::Expression;
+use crate::expr::BoundExpression;
 use crate::expr::traversal::FoldDownContext;
 use crate::expr::traversal::FoldUp;
 use crate::expr::traversal::NodeExt;
 use crate::expr::traversal::NodeFolderContext;
 use crate::scalar_fn::fns::get_item::GetItem;
-use crate::scalar_fn::fns::root::Root;
 use crate::scalar_fn::fns::select::Select;
 
 /// Returns the rooted field paths referenced by an expression.
@@ -24,47 +22,13 @@ use crate::scalar_fn::fns::select::Select;
 /// expression is represented by [`FieldPath::root`], which conservatively selects all fields.
 /// Scalar functions other than `GetItem` and `Select` conservatively reference each complete child
 /// output.
-pub fn referenced_field_paths(expr: &Expression, scope: &DType) -> VortexResult<FieldPathSet> {
-    // Validate the whole expression so plain GetItem paths and Select paths behave consistently.
-    expr.return_dtype(scope)?;
-
+pub fn referenced_field_paths(expr: &BoundExpression) -> VortexResult<FieldPathSet> {
     let mut collector = ReferencedFieldPaths {
-        scope,
         field_paths: FieldPathSet::default(),
     };
     expr.clone()
         .fold_context(&vec![FieldPath::root()], &mut collector)?;
-    let field_paths = collector.field_paths;
-
-    // The top-level field of every referenced path must be one of the immediately accessed scope
-    // fields: this analysis only refines *which nested fields* are read, never which top-level
-    // fields. `FieldPath::root()` stands in for "all fields", so it expands to the whole scope.
-    #[cfg(debug_assertions)]
-    if let Some(scope_fields) = scope.as_struct_fields_opt() {
-        use vortex_utils::aliases::hash_set::HashSet;
-
-        use crate::dtype::FieldName;
-        use crate::expr::analysis::immediate_access::immediate_scope_access;
-
-        let referenced_heads: HashSet<FieldName> = if field_paths.iter().any(FieldPath::is_root) {
-            scope_fields.names().iter().cloned().collect()
-        } else {
-            field_paths
-                .iter()
-                .filter_map(|path| match path.parts().first() {
-                    Some(Field::Name(name)) => Some(name.clone()),
-                    _ => None,
-                })
-                .collect()
-        };
-        debug_assert_eq!(
-            referenced_heads,
-            immediate_scope_access(expr, scope_fields),
-            "referenced field path heads must match the immediately accessed scope fields"
-        );
-    }
-
-    Ok(field_paths)
+    Ok(collector.field_paths)
 }
 
 /// Threads the set of currently-requested field paths down the expression tree, narrowing it at
@@ -78,22 +42,21 @@ pub fn referenced_field_paths(expr: &Expression, scope: &DType) -> VortexResult<
 /// column projection). Any other function is opaque—we cannot assume it preserves a field's
 /// provenance—so its children conservatively re-request the whole scope, which is what keeps an
 /// expression like `f($).x` reading every field of `$` rather than just `x`.
-struct ReferencedFieldPaths<'a> {
-    scope: &'a DType,
+struct ReferencedFieldPaths {
     field_paths: FieldPathSet,
 }
 
-impl NodeFolderContext for ReferencedFieldPaths<'_> {
-    type NodeTy = Expression;
+impl NodeFolderContext for ReferencedFieldPaths {
+    type NodeTy = BoundExpression;
     type Result = ();
     type Context = Vec<FieldPath>;
 
     fn visit_down(
         &mut self,
         requested: &Self::Context,
-        node: &Expression,
+        node: &BoundExpression,
     ) -> VortexResult<FoldDownContext<Self::Context, ()>> {
-        if node.is::<Root>() {
+        if node.is_root() {
             self.field_paths.extend(
                 requested
                     .iter()
@@ -102,7 +65,10 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> {
             return Ok(FoldDownContext::Skip(()));
         }
 
-        if let Some(field_name) = node.as_opt::<GetItem>() {
+        if let Some(field_name) = node
+            .as_scalar()
+            .and_then(|scalar_fn| scalar_fn.as_opt::<GetItem>())
+        {
             let appended = requested
                 .iter()
                 .map(|path| path.clone().push(Field::Name(field_name.clone())))
@@ -112,9 +78,12 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> {
 
         // Keep requested paths whose head is included, expanding a whole-scope request into one
         // path per included field.
-        if let Some(selection) = node.as_opt::<Select>() {
-            let child_dtype = node.child(0).return_dtype(self.scope)?;
-            let child_fields = child_dtype
+        if let Some(selection) = node
+            .as_scalar()
+            .and_then(|scalar_fn| scalar_fn.as_opt::<Select>())
+        {
+            let child_fields = node.children()[0]
+                .dtype()
                 .as_struct_fields_opt()
                 .ok_or_else(|| vortex_err!("Select child is not a struct"))?;
             let included_fields = selection.normalize_to_included_fields(child_fields.names())?;
@@ -146,7 +115,7 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> {
 
     fn visit_up(
         &mut self,
-        _node: Expression,
+        _node: BoundExpression,
         _requested: &Self::Context,
         _children: Vec<()>,
     ) -> VortexResult<FoldUp<()>> {
@@ -159,9 +128,11 @@ mod tests {
     use vortex_utils::aliases::hash_set::HashSet;
 
     use super::*;
+    use crate::dtype::DType;
     use crate::dtype::Nullability::NonNullable;
     use crate::dtype::PType::I32;
     use crate::dtype::StructFields;
+    use crate::expr::Expression;
     use crate::expr::get_item;
     use crate::expr::pack;
     use crate::expr::root;
@@ -183,7 +154,7 @@ mod tests {
 
     /// Collects the prefix-minimal field paths referenced by `expr` against [`scope`].
     fn referenced(expr: &Expression) -> VortexResult<HashSet<FieldPath>> {
-        Ok(referenced_field_paths(expr, &scope())?
+        Ok(referenced_field_paths(&expr.bind(&scope())?)?
             .into_iter()
             .collect())
     }
@@ -259,6 +230,9 @@ mod tests {
 
     #[test]
     fn invalid_get_item_path_returns_error() {
-        assert!(referenced_field_paths(&get_item("missing", root()), &scope()).is_err());
+        let result = get_item("missing", root())
+            .bind(&scope())
+            .and_then(|expr| referenced_field_paths(&expr));
+        assert!(result.is_err());
     }
 }

vortex-array/src/expr/bound_expression.rs

@@ -57,12 +57,8 @@ pub struct ExactBoundExpr(pub BoundExpression);
 
 impl PartialEq for ExactBoundExpr {
     fn eq(&self, other: &Self) -> bool {
-        if self.0.dtype != other.0.dtype {
-            return false;
-        }
-
         match (&self.0.kind, &other.0.kind) {
-            (BoundKind::Root, BoundKind::Root) => true,
+            (BoundKind::Root, BoundKind::Root) => self.0.dtype == other.0.dtype,
             (
                 BoundKind::Scalar {
                     scalar_fn: lhs_fn,
@@ -72,7 +68,11 @@ impl PartialEq for ExactBoundExpr {
                     scalar_fn: rhs_fn,
                     children: rhs_children,
                 },
-            ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children),
+            ) => {
+                lhs_fn == rhs_fn
+                    && Arc::ptr_eq(lhs_children, rhs_children)
+                    && self.0.dtype == other.0.dtype
+            }
             _ => false,
         }
     }
@@ -82,7 +82,8 @@ impl Eq for ExactBoundExpr {}
 
 impl Hash for ExactBoundExpr {
     fn hash<H: Hasher>(&self, state: &mut H) {
-        self.0.dtype.hash(state);
+        // DType differences are resolved by equality. Omitting the potentially lazy dtype keeps
+        // identity-keyed cache lookups from deserializing an entire schema just to compute a hash.
         match &self.0.kind {
             BoundKind::Root => state.write_u8(0),
             BoundKind::Scalar {

vortex-array/src/expr/transform/bound_partition.rs

@@ -102,7 +102,10 @@ where
     })
 }
 
-/// The result of partitioning an expression.
+/// The result of partitioning a bound expression.
+///
+/// The root and partitions remain bound so callers can cache and reuse their shared tree identity
+/// without an unbind/rebind round trip.
 #[derive(Debug)]
 pub struct BoundPartitionedExpr<A> {
     /// The root expression used to re-assemble the results.

vortex-array/src/expression.rs

@@ -8,12 +8,39 @@ use crate::ArrayRef;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
 use crate::arrays::ScalarFnArray;
+use crate::expr::BoundExpression;
+use crate::expr::BoundKind;
 use crate::expr::Expression;
 use crate::optimizer::ArrayOptimizer;
 use crate::scalar_fn::fns::literal::Literal;
 use crate::scalar_fn::fns::root::Root;
 
 impl ArrayRef {
+    /// Apply a bound expression to this array, producing a new array in constant time.
+    pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult<ArrayRef> {
+        let BoundKind::Scalar {
+            scalar_fn,
+            children,
+        } = expr.kind()
+        else {
+            return Ok(self);
+        };
+
+        if let Some(scalar) = scalar_fn.as_opt::<Literal>() {
+            return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array());
+        }
+
+        let children: Vec<_> = children
+            .iter()
+            .map(|child| self.clone().apply_bound(child))
+            .try_collect()?;
+
+        let array =
+            ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array();
+
+        array.optimize()
+    }
+
     /// Apply the expression to this array, producing a new array in constant time.
     pub fn apply(self, expr: &Expression) -> VortexResult<ArrayRef> {
         // If the expression is a root, return self.

vortex-cuda/src/layout.rs

@@ -19,7 +19,7 @@ use vortex::array::MaskFuture;
 use vortex::array::ProstMetadata;
 use vortex::array::VortexSessionExecute;
 use vortex::array::arrays::Constant;
-use vortex::array::expr::Expression;
+use vortex::array::expr::BoundExpression;
 use vortex::array::expr::stats::Precision;
 use vortex::array::expr::stats::Stat;
 use vortex::array::expr::stats::StatsProvider;
@@ -268,7 +268,7 @@ impl LayoutReader for CudaFlatReader {
     fn pruning_evaluation(
         &self,
         _row_range: &Range<u64>,
-        _expr: &Expression,
+        _expr: &BoundExpression,
         mask: Mask,
     ) -> VortexResult<MaskFuture> {
         Ok(MaskFuture::ready(mask))
@@ -277,7 +277,7 @@ impl LayoutReader for CudaFlatReader {
     fn filter_evaluation(
         &self,
         row_range: &Range<u64>,
-        expr: &Expression,
+        expr: &BoundExpression,
         mask: MaskFuture,
     ) -> VortexResult<MaskFuture> {
         let row_range = usize::try_from(row_range.start)
@@ -299,13 +299,13 @@ impl LayoutReader for CudaFlatReader {
 
             let mask_density = mask.density();
             let array_mask = if mask_density < EXPR_EVAL_THRESHOLD {
-                let array = array.apply(&expr)?;
+                let array = array.apply_bound(&expr)?;
                 let array = array.filter(mask.clone())?;
                 let mut ctx = session.create_execution_ctx();
                 let array_mask = array.null_as_false().execute(&mut ctx)?;
                 mask.intersect_by_rank(&array_mask)
             } else {
-                let array = array.apply(&expr)?;
+                let array = array.apply_bound(&expr)?;
                 let mut ctx = session.create_execution_ctx();
                 let array_mask = array.null_as_false().execute(&mut ctx)?;
                 mask.bitand(&array_mask)
@@ -326,7 +326,7 @@ impl LayoutReader for CudaFlatReader {
     fn projection_evaluation(
         &self,
         row_range: &Range<u64>,
-        expr: &Expression,
+        expr: &BoundExpression,
         mask: MaskFuture,
     ) -> VortexResult<BoxFuture<'static, VortexResult<ArrayRef>>> {
         let row_range = usize::try_from(row_range.start)
@@ -351,7 +351,7 @@ impl LayoutReader for CudaFlatReader {
                 array = array.filter(mask)?;
             }
 
-            array = array.apply(&expr)?;
+            array = array.apply_bound(&expr)?;
 
             Ok(array)
         }