From a9d023b8324e1e96b743c3c2760ab310e6a72274 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 6 Aug 2026 19:04:55 -0700 Subject: [PATCH 01/30] flatten BoundKind into BoundExpression Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 137 ++++++++++++++-------- vortex-array/src/expr/display.rs | 13 +- vortex-array/src/expr/traversal/mod.rs | 17 ++- vortex-array/src/expression.rs | 6 +- 4 files changed, 104 insertions(+), 69 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 45040ab2cd1..2da193398fe 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -32,17 +32,11 @@ use crate::stats::rewrite::StatsRewriteCtx; /// 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. @@ -52,7 +46,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. @@ -61,23 +58,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, } } } @@ -88,11 +93,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); @@ -105,10 +111,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. @@ -130,12 +133,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(), }) } @@ -145,7 +146,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", @@ -159,19 +160,16 @@ 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 { .. } => &[], } } @@ -182,9 +180,9 @@ impl BoundExpression { /// 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, } } @@ -223,7 +221,7 @@ impl BoundExpression { /// Whether this node is the scope root. pub fn is_root(&self) -> bool { - matches!(self.kind, BoundKind::Root) + matches!(self, Self::Root { .. }) } /// Return whether every scope root in this expression has `dtype`. @@ -256,13 +254,50 @@ impl BoundExpression { pub fn display_tree(&self) -> impl Display { DisplayTreeExpr(self) } + + /// Convert this bound tree back into its unbound logical representation. + /// + /// This rebuilds the expression iteratively; the bound representation does not retain a + /// second expression tree. + // TODO: This is temporary artifact of the migration from using `Expression`s to + // `BoundExpression`s + pub fn unbind(&self) -> Expression { + let mut pending = vec![(self, false)]; + let mut expressions = Vec::new(); + + while let Some((node, visited)) = pending.pop() { + 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); + expressions.push( + Expression::try_new(scalar_fn.clone(), child_expressions) + .vortex_expect("a bound expression always has valid arity"), + ); + } + BoundExpression::Scalar { children, .. } => { + pending.push((node, true)); + pending.extend(children.iter().rev().map(|child| (child, false))); + } + } + } + + expressions + .pop() + .vortex_expect("binding always produces one expression root") + } } 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("$"), } } } @@ -298,7 +333,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 { @@ -307,7 +342,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); @@ -411,8 +446,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 345f393e779..250d7063834 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -9,7 +9,6 @@ use vortex_utils::tree::TreeDisplayAdapter; use vortex_utils::tree::write_branch_tree; use crate::expr::BoundExpression; -use crate::expr::BoundKind; use crate::expr::Expression; use crate::scalar_fn::ChildName; @@ -87,16 +86,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 19877cb5b06e8d7337322ff740debdbaf749cdd5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 11 Aug 2026 09:29:32 -0700 Subject: [PATCH 02/30] fix Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 42 +---------------------- 1 file changed, 1 insertion(+), 41 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 2da193398fe..ad95ae61c9d 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -79,10 +79,7 @@ impl PartialEq for ExactBoundExpr { && Arc::ptr_eq(lhs_children, rhs_children) && lhs_dtype == rhs_dtype } - // 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, + _ => false, } } } @@ -254,43 +251,6 @@ impl BoundExpression { pub fn display_tree(&self) -> impl Display { DisplayTreeExpr(self) } - - /// Convert this bound tree back into its unbound logical representation. - /// - /// This rebuilds the expression iteratively; the bound representation does not retain a - /// second expression tree. - // TODO: This is temporary artifact of the migration from using `Expression`s to - // `BoundExpression`s - pub fn unbind(&self) -> Expression { - let mut pending = vec![(self, false)]; - let mut expressions = Vec::new(); - - while let Some((node, visited)) = pending.pop() { - 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); - expressions.push( - Expression::try_new(scalar_fn.clone(), child_expressions) - .vortex_expect("a bound expression always has valid arity"), - ); - } - BoundExpression::Scalar { children, .. } => { - pending.push((node, true)); - pending.extend(children.iter().rev().map(|child| (child, false))); - } - } - } - - expressions - .pop() - .vortex_expect("binding always produces one expression root") - } } impl Display for BoundExpression { From 26eb8e724323fa71bd936da1231dbf167410f526 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 6 Aug 2026 19:04:55 -0700 Subject: [PATCH 03/30] flatten BoundKind into BoundExpression Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 42 ++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index ad95ae61c9d..2da193398fe 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -79,7 +79,10 @@ impl PartialEq for ExactBoundExpr { && Arc::ptr_eq(lhs_children, rhs_children) && 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, } } } @@ -251,6 +254,43 @@ impl BoundExpression { pub fn display_tree(&self) -> impl Display { DisplayTreeExpr(self) } + + /// Convert this bound tree back into its unbound logical representation. + /// + /// This rebuilds the expression iteratively; the bound representation does not retain a + /// second expression tree. + // TODO: This is temporary artifact of the migration from using `Expression`s to + // `BoundExpression`s + pub fn unbind(&self) -> Expression { + let mut pending = vec![(self, false)]; + let mut expressions = Vec::new(); + + while let Some((node, visited)) = pending.pop() { + 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); + expressions.push( + Expression::try_new(scalar_fn.clone(), child_expressions) + .vortex_expect("a bound expression always has valid arity"), + ); + } + BoundExpression::Scalar { children, .. } => { + pending.push((node, true)); + pending.extend(children.iter().rev().map(|child| (child, false))); + } + } + } + + expressions + .pop() + .vortex_expect("binding always produces one expression root") + } } impl Display for BoundExpression { From 1eb60a689f6513a421406df666564ad1a83ae223 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 6 Aug 2026 19:48:30 -0700 Subject: [PATCH 04/30] 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 | 249 ++++++++- vortex-array/src/expr/display.rs | 19 +- vortex-array/src/expr/expression.rs | 118 ++++- vortex-array/src/expr/exprs.rs | 477 +----------------- vortex-array/src/expr/lambda.rs | 74 +++ vortex-array/src/expr/mod.rs | 59 +-- 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, 855 insertions(+), 590 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 6c2e4975a92..8f97c0e862a 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -67,7 +67,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::