From 8fe92ced8eec8b01dfb30c2dfae3d55080a16473 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 31 Jul 2026 17:14:18 +0100 Subject: [PATCH 1/3] feat(expr): add bound expressions Signed-off-by: Joe Isaacs --- vortex-array/src/expr/bound_expression.rs | 390 ++++++++++++++++++++++ vortex-array/src/expr/display.rs | 95 ++++-- vortex-array/src/expr/expression.rs | 21 +- vortex-array/src/expr/mod.rs | 4 + vortex-array/src/expr/scope.rs | 45 +++ 5 files changed, 518 insertions(+), 37 deletions(-) create mode 100644 vortex-array/src/expr/bound_expression.rs create mode 100644 vortex-array/src/expr/scope.rs diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs new file mode 100644 index 00000000000..452e4df8377 --- /dev/null +++ b/vortex-array/src/expr/bound_expression.rs @@ -0,0 +1,390 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::dtype::DType; +use crate::expr::Expression; +use crate::expr::display::DisplayTreeExpr; +use crate::expr::scope::Scope; +use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::fns::root::Root; + +/// An [`Expression`] that has been type-checked against a [`Scope`]. +/// +/// Every node carries its own dtype, so reading one is a field access rather than a walk of the +/// subtree. Holding a `BoundExpression` is proof that the whole tree type-checked. +/// +/// 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 { + /// A scalar function applied to bound children. + Scalar { + /// The scalar function for this node. + scalar_fn: ScalarFnRef, + /// The bound children, in argument order. + /// + /// Sharing keeps clones cheap even though the iterative [`Drop`] implementation prevents + /// consumers from destructuring a `BoundExpression` by value. + children: Arc>, + }, + /// The scope itself. Its dtype is the scope's root dtype. + Root, +} + +/// A bound-expression wrapper that compares shared tree identity instead of structure. +#[derive(Clone, Debug)] +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::Scalar { + scalar_fn: lhs_fn, + children: lhs_children, + }, + BoundKind::Scalar { + scalar_fn: rhs_fn, + children: rhs_children, + }, + ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children), + _ => false, + } + } +} + +impl Eq for ExactBoundExpr {} + +impl Hash for ExactBoundExpr { + fn hash(&self, state: &mut H) { + self.0.dtype.hash(state); + match &self.0.kind { + BoundKind::Root => state.write_u8(0), + BoundKind::Scalar { + scalar_fn, + children, + } => { + state.write_u8(1); + scalar_fn.hash(state); + Arc::as_ptr(children).hash(state); + } + } + } +} + +impl BoundExpression { + /// Create a bound root expression with the given dtype. + pub fn new_root(dtype: DType) -> Self { + Self { + kind: BoundKind::Root, + dtype, + } + } + + /// Create a bound scalar node from a scalar function and already-bound children. + pub fn try_new( + scalar_fn: ScalarFnRef, + children: impl IntoIterator, + ) -> VortexResult { + let children = Vec::from_iter(children); + vortex_ensure!( + scalar_fn.signature().arity().matches(children.len()), + "Expression arity mismatch: expected {} children but got {}", + scalar_fn.signature().arity(), + children.len() + ); + + let arg_dtypes = children + .iter() + .map(|child| child.dtype().clone()) + .collect_vec(); + let dtype = scalar_fn.return_dtype(&arg_dtypes)?; + + Ok(Self { + kind: BoundKind::Scalar { + scalar_fn, + children: children.into(), + }, + dtype, + }) + } + + /// Rebuild this node with new bound children, recomputing its dtype. + pub fn with_children( + self, + children: impl IntoIterator, + ) -> VortexResult { + let children = Vec::from_iter(children); + let BoundKind::Scalar { scalar_fn, .. } = &self.kind else { + vortex_ensure!( + children.is_empty(), + "Root expression cannot have {} children", + children.len() + ); + return Ok(self); + }; + + Self::try_new(scalar_fn.clone(), children) + } + + /// 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 + } + + /// The bound children of this node, in argument order. Empty for [`BoundKind::Root`]. + pub fn children(&self) -> &[BoundExpression] { + match &self.kind { + BoundKind::Scalar { children, .. } => children.as_slice(), + BoundKind::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, + } + } + + /// Whether this node is the scope root. + pub fn is_root(&self) -> bool { + matches!(self.kind, BoundKind::Root) + } + + /// Display the bound expression as a formatted tree structure. + 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.kind() { + BoundKind::Root => expressions.push(crate::expr::root()), + BoundKind::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"), + ); + } + BoundKind::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 { + Display::fmt(&self.unbind(), f) + } +} + +impl Expression { + /// Bind this expression against a root dtype, type-checking every node in a single walk. + /// + /// The returned tree carries a dtype on each node, so callers needing types at more than one + /// node should bind once and read fields rather than calling + /// [`return_dtype`](Expression::return_dtype) repeatedly. + pub fn bind(&self, dtype: &DType) -> VortexResult { + self.bind_scope(&Scope::new(dtype.clone())) + } + + /// Bind this expression against an explicit [`Scope`]. + pub fn bind_scope(&self, scope: &Scope) -> VortexResult { + if self.is::() { + return Ok(BoundExpression::new_root(scope.root().clone())); + } + + let children: Vec<_> = self + .children() + .iter() + .map(|child| child.bind_scope(scope)) + .try_collect()?; + BoundExpression::try_new(self.scalar_fn().clone(), children) + } +} + +/// 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 { + return; + }; + let Some(children) = Arc::get_mut(children) else { + return; + }; + + let mut to_drop = std::mem::take(children); + while let Some(mut child) = to_drop.pop() { + if let BoundKind::Scalar { children, .. } = &mut child.kind + && let Some(grandchildren) = Arc::get_mut(children) + { + to_drop.append(grandchildren); + } + } + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::*; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::col; + use crate::expr::eq; + use crate::expr::lit; + use crate::expr::root; + use crate::expr::test_harness::struct_dtype; + + fn scope() -> Scope { + Scope::new(struct_dtype()) + } + + #[test] + fn root_binds_to_the_scope() -> VortexResult<()> { + let bound = root().bind_scope(&scope())?; + assert!(bound.is_root()); + assert_eq!(bound.dtype(), &struct_dtype()); + assert_eq!(bound.unbind(), root()); + Ok(()) + } + + #[test] + fn every_node_carries_its_dtype() -> VortexResult<()> { + let expr = eq(col("a"), lit(1_i32)); + let bound = expr.bind_scope(&scope())?; + + assert_eq!(bound.dtype(), &DType::Bool(Nullability::NonNullable)); + + let lhs = &bound.children()[0]; + assert_eq!( + lhs.dtype(), + &DType::Primitive(PType::I32, Nullability::NonNullable) + ); + assert_eq!(lhs.children()[0].dtype(), &struct_dtype()); + Ok(()) + } + + #[test] + fn bind_agrees_with_return_dtype() -> VortexResult<()> { + for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { + assert_eq!( + expr.bind(&struct_dtype())?.dtype(), + &expr.return_dtype(&struct_dtype())?, + "disagreement for {expr}" + ); + } + Ok(()) + } + + #[test] + fn bound_tree_display_matches_unbound() -> VortexResult<()> { + for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { + let bound = expr.bind_scope(&scope())?; + assert_eq!( + bound.display_tree().to_string(), + expr.display_tree().to_string() + ); + } + Ok(()) + } + + #[test] + fn clone_shares_children() -> VortexResult<()> { + 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()) + else { + unreachable!("eq is a scalar node") + }; + assert!(Arc::ptr_eq(a, b)); + Ok(()) + } + + #[test] + fn repeated_subtree_is_bound_per_occurrence() -> VortexResult<()> { + let shared = col("a"); + let bound = eq(shared.clone(), shared).bind_scope(&scope())?; + let children = bound.children(); + assert_eq!(children[0].dtype(), children[1].dtype()); + Ok(()) + } + + #[test] + fn structural_and_exact_equality_are_distinct() -> VortexResult<()> { + let expr = eq(col("a"), lit(1_i32)); + let bound = expr.bind_scope(&scope())?; + let independently_bound = expr.bind_scope(&scope())?; + + assert_eq!(bound, independently_bound); + assert_eq!(ExactBoundExpr(bound.clone()), ExactBoundExpr(bound.clone())); + assert_ne!( + ExactBoundExpr(bound.clone()), + ExactBoundExpr(independently_bound) + ); + assert_eq!(bound.unbind(), expr); + Ok(()) + } + + #[test] + fn binding_reports_a_type_error() { + let expr = eq(col("a"), lit("nope")); + assert!(expr.bind_scope(&scope()).is_err()); + } +} diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 492df401b9f..685af718134 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -1,47 +1,96 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fmt; use std::fmt::Display; use std::fmt::Formatter; -use std::ops::Deref; +use crate::expr::BoundExpression; +use crate::expr::BoundKind; use crate::expr::Expression; -use crate::scalar_fn::ScalarFnRef; +use crate::expr::root; +use crate::scalar_fn::ChildName; pub enum DisplayFormat { Compact, Tree, } -pub struct DisplayTreeExpr<'a>(pub &'a Expression); +trait DisplayTreeNode: Sized { + fn tree_children(&self) -> &[Self]; -impl Display for DisplayTreeExpr<'_> { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - pub use termtree::Tree; - fn make_tree(expr: &Expression) -> Result, std::fmt::Error> { - let scalar_fn: &ScalarFnRef = expr.deref(); - let node_name = format!("{}", scalar_fn); + fn tree_child_name(&self, index: usize) -> ChildName; + + fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result; +} + +impl DisplayTreeNode for Expression { + fn tree_children(&self) -> &[Self] { + Expression::children(self).as_slice() + } + + fn tree_child_name(&self, index: usize) -> ChildName { + self.scalar_fn().signature().child_name(index) + } + + fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { + Display::fmt(self.scalar_fn(), f) + } +} + +impl DisplayTreeNode for BoundExpression { + fn tree_children(&self) -> &[Self] { + BoundExpression::children(self) + } - // Get child names for display purposes - let child_names = (0..expr.children().len()).map(|i| expr.signature().child_name(i)); - let children = expr.children(); + 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"), + } + } + + fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self.kind() { + BoundKind::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + BoundKind::Root => Display::fmt(root().scalar_fn(), f), + } + } +} - let child_trees: Result>, std::fmt::Error> = children +struct NodeDisplay<'a, T>(&'a T); + +impl Display for NodeDisplay<'_, T> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.0.fmt_tree_node(f) + } +} + +pub struct DisplayTreeExpr<'a, T: ?Sized = Expression>(pub &'a T); + +impl Display for DisplayTreeExpr<'_, T> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + pub use termtree::Tree; + fn make_tree(expr: &T) -> Tree { + let child_trees = expr + .tree_children() .iter() - .zip(child_names) - .map(|(child, name)| { - let child_tree = make_tree(child)?; - Ok::, std::fmt::Error>( - Tree::new(format!("{}: {}", name, child_tree.root)) - .with_leaves(child_tree.leaves), - ) + .enumerate() + .map(|(index, child)| { + let child_tree = make_tree(child); + Tree::new(format!( + "{}: {}", + expr.tree_child_name(index), + child_tree.root + )) + .with_leaves(child_tree.leaves) }) - .collect(); + .collect::>(); - Ok(Tree::new(node_name).with_leaves(child_trees?)) + Tree::new(NodeDisplay(expr).to_string()).with_leaves(child_trees) } - write!(f, "{}", make_tree(self.0)?) + write!(f, "{}", make_tree(self.0)) } } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index 41cd61a369d..30f082cd84d 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -9,7 +9,6 @@ use std::hash::Hash; use std::ops::Deref; use std::sync::Arc; -use itertools::Itertools; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_session::VortexSession; @@ -17,7 +16,7 @@ use vortex_session::VortexSession; use crate::dtype::DType; use crate::expr::display::DisplayTreeExpr; use crate::scalar_fn::ScalarFnRef; -use crate::scalar_fn::fns::root::Root; +use crate::stats::rewrite::StatsRewriteCtx; /// A node in a Vortex expression tree. /// @@ -92,17 +91,11 @@ impl Expression { } /// Computes the return dtype of this expression given the input dtype. + /// + /// This binds the expression and discards everything but the root dtype. Callers needing types + /// at more than one node should bind once and read the dtypes off the bound tree. pub fn return_dtype(&self, scope: &DType) -> VortexResult { - if self.is::() { - return Ok(scope.clone()); - } - - let dtypes: Vec<_> = self - .children - .iter() - .map(|c| c.return_dtype(scope)) - .try_collect()?; - self.scalar_fn.return_dtype(&dtypes) + Ok(self.bind(scope)?.dtype().clone()) } /// Returns a new expression representing the validity mask output of this expression. @@ -123,7 +116,7 @@ impl Expression { scope: &DType, session: &VortexSession, ) -> VortexResult> { - crate::stats::rewrite::StatsRewriteCtx::new(session, scope).falsify(self) + StatsRewriteCtx::new(session, scope).falsify(self) } /// Returns an expression that proves this predicate is definitely true from stats. @@ -137,7 +130,7 @@ impl Expression { scope: &DType, session: &VortexSession, ) -> VortexResult> { - crate::stats::rewrite::StatsRewriteCtx::new(session, scope).satisfy(self) + StatsRewriteCtx::new(session, scope).satisfy(self) } /// Format the expression as a compact string. diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index e71e641a204..df7e3f59b97 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -56,6 +56,7 @@ pub mod aliases; pub mod analysis; #[cfg(feature = "arbitrary")] pub mod arbitrary; +pub mod bound_expression; pub mod display; pub(crate) mod expression; mod exprs; @@ -63,13 +64,16 @@ pub(crate) mod field; pub mod forms; mod optimize; pub mod proto; +pub mod scope; pub mod stats; pub mod transform; pub mod traversal; pub use analysis::*; +pub use bound_expression::*; pub use expression::*; pub use exprs::*; +pub use scope::*; pub trait VortexExprExt { /// Accumulate all field references from this expression and its children in a set diff --git a/vortex-array/src/expr/scope.rs b/vortex-array/src/expr/scope.rs new file mode 100644 index 00000000000..04de3202cf5 --- /dev/null +++ b/vortex-array/src/expr/scope.rs @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::dtype::DType; + +/// The context an [`Expression`](crate::expr::Expression) is bound against. +/// +/// Today a scope is just the dtype that [`root`](crate::expr::root) resolves to. It is an opaque +/// struct rather than a bare [`DType`] so that lexical bindings can be added later without changing +/// [`Expression::bind_scope`](crate::expr::Expression::bind_scope)'s signature. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Scope { + root: DType, +} + +impl Scope { + /// Create a scope in which `root` resolves to the given dtype. + pub fn new(root: DType) -> Self { + Self { root } + } + + /// The dtype that `root` resolves to. + pub fn root(&self) -> &DType { + &self.root + } +} + +impl From for Scope { + fn from(root: DType) -> Self { + Self::new(root) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dtype::Nullability; + + #[test] + fn root_round_trips() { + let dtype = DType::Bool(Nullability::Nullable); + assert_eq!(Scope::new(dtype.clone()).root(), &dtype); + assert_eq!(Scope::from(dtype.clone()).root(), &dtype); + } +} From 0d9162eeeea086523be3e47e39f926a3b9cdb08f Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 3 Aug 2026 13:32:47 +0100 Subject: [PATCH 2/3] fix Signed-off-by: Joe Isaacs --- vortex-array/src/expr/expression.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index 30f082cd84d..03fab1fce9d 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -9,6 +9,7 @@ use std::hash::Hash; use std::ops::Deref; use std::sync::Arc; +use itertools::Itertools; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_session::VortexSession; @@ -16,6 +17,7 @@ use vortex_session::VortexSession; use crate::dtype::DType; use crate::expr::display::DisplayTreeExpr; use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::fns::root::Root; use crate::stats::rewrite::StatsRewriteCtx; /// A node in a Vortex expression tree. @@ -91,18 +93,17 @@ impl Expression { } /// Computes the return dtype of this expression given the input dtype. - /// - /// This binds the expression and discards everything but the root dtype. Callers needing types - /// at more than one node should bind once and read the dtypes off the bound tree. pub fn return_dtype(&self, scope: &DType) -> VortexResult { - Ok(self.bind(scope)?.dtype().clone()) - } + if self.is::() { + return Ok(scope.clone()); + } - /// Returns a new expression representing the validity mask output of this expression. - /// - /// The returned expression evaluates to a non-nullable boolean array. - pub fn validity(&self) -> VortexResult { - self.scalar_fn.validity(self) + let dtypes: Vec<_> = self + .children + .iter() + .map(|c| c.return_dtype(scope)) + .try_collect()?; + self.scalar_fn.return_dtype(&dtypes) } /// Returns an expression that proves this predicate is definitely false from stats. From 621f73512a731c9980e2fcf9bddc49cc6c08e626 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 3 Aug 2026 13:34:45 +0100 Subject: [PATCH 3/3] fix Signed-off-by: Joe Isaacs --- vortex-array/src/expr/expression.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index 03fab1fce9d..d7f85825dbe 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -106,6 +106,13 @@ impl Expression { self.scalar_fn.return_dtype(&dtypes) } + /// Returns a new expression representing the validity mask output of this expression. + /// + /// The returned expression evaluates to a non-nullable boolean array. + pub fn validity(&self) -> VortexResult { + self.scalar_fn.validity(self) + } + /// Returns an expression that proves this predicate is definitely false from stats. /// /// `scope` is the dtype of the row this expression evaluates over.