From 93a164fb86bccfc47046947ed0be4a2623dd75dc Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 6 Aug 2026 19:04:55 -0700 Subject: [PATCH 1/2] flatten BoundKind into BoundExpression Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 109 +++++++++++----------- vortex-array/src/expr/display.rs | 13 ++- vortex-array/src/expr/traversal/mod.rs | 17 ++-- vortex-array/src/expression.rs | 6 +- 4 files changed, 72 insertions(+), 73 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index d93da7d2580..51aa5917dec 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -27,17 +27,11 @@ use crate::scalar_fn::ScalarFnRef; /// Binding is purely logical: it deals only in [`DType`]s and never sees an array, a length, or an /// encoding. #[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct BoundExpression { - kind: BoundKind, - dtype: DType, -} - -/// The per-variant contents of a [`BoundExpression`], mirroring the logical variants of -/// [`Expression`]. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub enum BoundKind { +pub enum BoundExpression { /// A scalar function applied to bound children. Scalar { + /// The dtype this node evaluates to. + dtype: DType, /// The scalar function for this node. scalar_fn: ScalarFnRef, /// The bound children, in argument order. @@ -47,7 +41,10 @@ pub enum BoundKind { children: Arc>, }, /// The scope itself. Its dtype is the scope's root dtype. - Root, + Root { + /// The dtype this node evaluates to. + dtype: DType, + }, } /// A bound-expression wrapper that compares shared tree identity instead of structure. @@ -56,23 +53,31 @@ pub struct ExactBoundExpr(pub BoundExpression); impl PartialEq for ExactBoundExpr { fn eq(&self, other: &Self) -> bool { - match (&self.0.kind, &other.0.kind) { - (BoundKind::Root, BoundKind::Root) => self.0.dtype == other.0.dtype, + match (&self.0, &other.0) { + ( + BoundExpression::Root { dtype: lhs_dtype }, + BoundExpression::Root { dtype: rhs_dtype }, + ) => lhs_dtype == rhs_dtype, ( - BoundKind::Scalar { + BoundExpression::Scalar { + dtype: lhs_dtype, scalar_fn: lhs_fn, children: lhs_children, }, - BoundKind::Scalar { + BoundExpression::Scalar { + dtype: rhs_dtype, scalar_fn: rhs_fn, children: rhs_children, }, ) => { lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children) - && self.0.dtype == other.0.dtype + && lhs_dtype == rhs_dtype } - _ => false, + // No catch-all: a new variant must state its own identity rather than silently + // comparing unequal, which would put `eq` out of step with `hash`. + (BoundExpression::Root { .. }, BoundExpression::Scalar { .. }) + | (BoundExpression::Scalar { .. }, BoundExpression::Root { .. }) => false, } } } @@ -83,11 +88,12 @@ impl Hash for ExactBoundExpr { fn hash(&self, state: &mut H) { // 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 { + match &self.0 { + BoundExpression::Root { .. } => state.write_u8(0), + BoundExpression::Scalar { scalar_fn, children, + .. } => { state.write_u8(1); scalar_fn.hash(state); @@ -100,10 +106,7 @@ impl Hash for ExactBoundExpr { impl BoundExpression { /// Create a bound root expression with the given dtype. pub fn new_root(dtype: DType) -> Self { - Self { - kind: BoundKind::Root, - dtype, - } + Self::Root { dtype } } /// Create a bound scalar node from a scalar function and already-bound children. @@ -125,12 +128,10 @@ impl BoundExpression { .collect_vec(); let dtype = scalar_fn.return_dtype(&arg_dtypes)?; - Ok(Self { - kind: BoundKind::Scalar { - scalar_fn, - children: children.into(), - }, + Ok(Self::Scalar { dtype, + scalar_fn, + children: children.into(), }) } @@ -140,7 +141,7 @@ impl BoundExpression { children: impl IntoIterator, ) -> VortexResult { let children = Vec::from_iter(children); - let BoundKind::Scalar { scalar_fn, .. } = &self.kind else { + let BoundExpression::Scalar { scalar_fn, .. } = &self else { vortex_ensure!( children.is_empty(), "Root expression cannot have {} children", @@ -154,33 +155,30 @@ impl BoundExpression { /// The dtype this expression evaluates to. pub fn dtype(&self) -> &DType { - &self.dtype - } - - /// The per-variant contents of this node. - pub fn kind(&self) -> &BoundKind { - &self.kind + match self { + Self::Scalar { dtype, .. } | Self::Root { dtype } => dtype, + } } - /// The bound children of this node, in argument order. Empty for [`BoundKind::Root`]. + /// The bound children of this node, in argument order. Empty for [`BoundExpression::Root`]. pub fn children(&self) -> &[BoundExpression] { - match &self.kind { - BoundKind::Scalar { children, .. } => children.as_slice(), - BoundKind::Root => &[], + match self { + Self::Scalar { children, .. } => children.as_slice(), + Self::Root { .. } => &[], } } /// The scalar function for this node, or `None` if it is the scope root. pub fn as_scalar(&self) -> Option<&ScalarFnRef> { - match &self.kind { - BoundKind::Scalar { scalar_fn, .. } => Some(scalar_fn), - BoundKind::Root => None, + match self { + Self::Scalar { scalar_fn, .. } => Some(scalar_fn), + Self::Root { .. } => None, } } /// Whether this node is the scope root. pub fn is_root(&self) -> bool { - matches!(self.kind, BoundKind::Root) + matches!(self, Self::Root { .. }) } /// Display the bound expression as a formatted tree structure. @@ -199,11 +197,12 @@ impl BoundExpression { let mut expressions = Vec::new(); while let Some((node, visited)) = pending.pop() { - match node.kind() { - BoundKind::Root => expressions.push(crate::expr::root()), - BoundKind::Scalar { + match node { + BoundExpression::Root { .. } => expressions.push(crate::expr::root()), + BoundExpression::Scalar { scalar_fn, children, + .. } if visited => { let child_start = expressions.len() - children.len(); let child_expressions = expressions.split_off(child_start); @@ -212,7 +211,7 @@ impl BoundExpression { .vortex_expect("a bound expression always has valid arity"), ); } - BoundKind::Scalar { children, .. } => { + BoundExpression::Scalar { children, .. } => { pending.push((node, true)); pending.extend(children.iter().rev().map(|child| (child, false))); } @@ -227,9 +226,9 @@ impl BoundExpression { impl Display for BoundExpression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - match self.kind() { - BoundKind::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), - BoundKind::Root => f.write_str("$"), + match self { + Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + Self::Root { .. } => f.write_str("$"), } } } @@ -265,7 +264,7 @@ impl Expression { /// Iterative drop to avoid stack overflows on deep trees. impl Drop for BoundExpression { fn drop(&mut self) { - let BoundKind::Scalar { children, .. } = &mut self.kind else { + let Self::Scalar { children, .. } = self else { return; }; let Some(children) = Arc::get_mut(children) else { @@ -274,7 +273,7 @@ impl Drop for BoundExpression { let mut to_drop = std::mem::take(children); while let Some(mut child) = to_drop.pop() { - if let BoundKind::Scalar { children, .. } = &mut child.kind + if let BoundExpression::Scalar { children, .. } = &mut child && let Some(grandchildren) = Arc::get_mut(children) { to_drop.append(grandchildren); @@ -355,8 +354,10 @@ mod tests { let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; let cloned = bound.clone(); - let (BoundKind::Scalar { children: a, .. }, BoundKind::Scalar { children: b, .. }) = - (bound.kind(), cloned.kind()) + let ( + BoundExpression::Scalar { children: a, .. }, + BoundExpression::Scalar { children: b, .. }, + ) = (&bound, &cloned) else { unreachable!("eq is a scalar node") }; diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 202ec857bd0..a7642325b67 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -6,7 +6,6 @@ use std::fmt::Display; use std::fmt::Formatter; use crate::expr::BoundExpression; -use crate::expr::BoundKind; use crate::expr::Expression; use crate::scalar_fn::ChildName; @@ -84,16 +83,16 @@ impl DisplayTreeNode for BoundExpression { } fn tree_child_name(&self, index: usize) -> ChildName { - match self.kind() { - BoundKind::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), - BoundKind::Root => unreachable!("the scope root has no children"), + match self { + BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), + BoundExpression::Root { .. } => unreachable!("the scope root has no children"), } } fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { - match self.kind() { - BoundKind::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), - BoundKind::Root => write!(f, "{ROOT_DISPLAY}"), + match self { + BoundExpression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + BoundExpression::Root { .. } => write!(f, "{ROOT_DISPLAY}"), } } } diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index b97ae2536f7..952a73f2657 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -24,7 +24,6 @@ use vortex_error::VortexResult; use crate::expr::BoundExpression; use crate::expr::Expression; -use crate::expr::bound_expression::BoundKind; use crate::expr::traversal::fold::NodeFolderContextWrapper; /// Signal to control a traversal's flow @@ -534,7 +533,7 @@ impl Node for BoundExpression { &'a self, mut f: F, ) -> VortexResult { - let BoundKind::Scalar { children, .. } = self.kind() else { + let BoundExpression::Scalar { children, .. } = self else { return Ok(TraversalOrder::Continue); }; @@ -552,7 +551,7 @@ impl Node for BoundExpression { self, mut f: F, ) -> VortexResult> { - let BoundKind::Scalar { children, .. } = self.kind() else { + let BoundExpression::Scalar { children, .. } = &self else { return Ok(Transformed::no(self)); }; @@ -583,16 +582,16 @@ impl Node for BoundExpression { } fn iter_children(&self, f: impl FnOnce(&mut dyn Iterator) -> T) -> T { - match self.kind() { - BoundKind::Scalar { children, .. } => f(&mut children.iter()), - BoundKind::Root => f(&mut std::iter::empty()), + match self { + BoundExpression::Scalar { children, .. } => f(&mut children.iter()), + BoundExpression::Root { .. } => f(&mut std::iter::empty()), } } fn children_count(&self) -> usize { - match self.kind() { - BoundKind::Scalar { children, .. } => children.len(), - BoundKind::Root => 0, + match self { + BoundExpression::Scalar { children, .. } => children.len(), + BoundExpression::Root { .. } => 0, } } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index a1a46196766..d0590f2bf58 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -10,7 +10,6 @@ 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; @@ -18,10 +17,11 @@ use crate::scalar_fn::fns::literal::Literal; impl ArrayRef { /// Apply a bound expression to this array, producing a new array in constant time. pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult { - let BoundKind::Scalar { + let BoundExpression::Scalar { scalar_fn, children, - } = expr.kind() + .. + } = expr else { return Ok(self); }; From e41b4842db063f5ce2d0c3bf2f22b3abc53334d2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 6 Aug 2026 19:48:30 -0700 Subject: [PATCH 2/2] lambdas and variables Signed-off-by: Matt Katz --- vortex-array/src/expr/analysis/fallible.rs | 31 ++- .../src/expr/analysis/immediate_access.rs | 44 +++- .../expr/analysis/referenced_field_paths.rs | 2 +- vortex-array/src/expr/analysis/strict.rs | 7 +- vortex-array/src/expr/bound_expression.rs | 240 ++++++++++++++++-- vortex-array/src/expr/display.rs | 20 +- vortex-array/src/expr/expression.rs | 118 +++++++-- vortex-array/src/expr/exprs.rs | 19 ++ vortex-array/src/expr/lambda.rs | 74 ++++++ vortex-array/src/expr/mod.rs | 14 + vortex-array/src/expr/optimize.rs | 6 +- vortex-array/src/expr/proto.rs | 42 ++- vortex-array/src/expr/scope.rs | 143 ++++++++++- .../src/expr/transform/bound_partition.rs | 9 +- vortex-array/src/expr/traversal/mod.rs | 8 +- vortex-array/src/expr/variable.rs | 96 +++++++ vortex-array/src/expression.rs | 45 ++-- vortex-layout/src/layouts/chunked/reader.rs | 4 +- vortex-layout/src/layouts/list/reader.rs | 2 +- vortex-layout/src/layouts/partitioned.rs | 2 +- vortex-layout/src/scan/filter.rs | 8 +- 21 files changed, 844 insertions(+), 90 deletions(-) create mode 100644 vortex-array/src/expr/lambda.rs create mode 100644 vortex-array/src/expr/variable.rs diff --git a/vortex-array/src/expr/analysis/fallible.rs b/vortex-array/src/expr/analysis/fallible.rs index ff43d51603b..4b39f9c09b5 100644 --- a/vortex-array/src/expr/analysis/fallible.rs +++ b/vortex-array/src/expr/analysis/fallible.rs @@ -10,8 +10,11 @@ pub fn label_is_fallible(expr: &Expression) -> BooleanLabels<'_> { expr, |expr| match expr { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_fallible(), - // The scope itself cannot fail. - Expression::Root => false, + // These add no fallibility of their own. Note this is the *self* label: a lambda's + // body is one of its children, so the folded label at a lambda node is the body's + // fallibility. A higher-order function therefore picks the body up through the + // ordinary fold instead of walking it by hand. + Expression::Root | Expression::Variable(_) | Expression::Lambda(_) => false, }, |acc, &child| acc | child, ) @@ -82,3 +85,27 @@ mod tests { assert_eq!(labels.get(&expr), Some(&false)); } } + +#[cfg(test)] +mod lambda_tests { + use super::*; + use crate::expr::checked_add; + use crate::expr::lambda; + use crate::expr::lit; + use crate::expr::var; + + /// A lambda contributes no fallibility of its own, but its body is one of its children, so the + /// label at the lambda node is the body's. That is what lets a future higher-order function + /// pick the body up through the ordinary fold rather than walking it by hand. + #[test] + fn a_lambdas_label_is_its_bodys_fallibility() { + let fallible = Expression::from(lambda(["x"], checked_add(var("x"), lit(1i32)))); + assert_eq!(label_is_fallible(&fallible).get(&fallible), Some(&true)); + + let infallible = Expression::from(lambda(["x"], var("x"))); + assert_eq!( + label_is_fallible(&infallible).get(&infallible), + Some(&false) + ); + } +} diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 612552ce4e5..eb6b1c3ddad 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -66,7 +66,13 @@ pub fn make_bound_free_field_annotator( ) -> impl AnnotationFn { move |expr: &BoundExpression| { let Some(scalar_fn) = expr.as_scalar() else { - return scope.names().iter().cloned().collect(); + // Only the scope root reads every field. A variable resolves against a frame, so it + // reads none of them, and saying otherwise would defeat column pruning. + return if expr.is_root() { + scope.names().iter().cloned().collect() + } else { + vec![] + }; }; if let Some(selection) = scalar_fn.as_opt::