From d8e744f9a4f8ce24178a5f3f656542ce64c21a1c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 31 Aug 2026 13:01:35 -0400 Subject: [PATCH] Add lambda expressions Signed-off-by: Matt Katz --- .../src/expr/analysis/immediate_access.rs | 4 +- vortex-array/src/expr/analysis/infallible.rs | 4 +- vortex-array/src/expr/analysis/strict.rs | 5 +- vortex-array/src/expr/bound_expression.rs | 288 +++++++++++++++++- vortex-array/src/expr/display.rs | 58 +++- vortex-array/src/expr/expression.rs | 110 +++++-- vortex-array/src/expr/exprs.rs | 11 + vortex-array/src/expr/lambda.rs | 96 ++++++ vortex-array/src/expr/mod.rs | 9 + vortex-array/src/expr/optimize.rs | 36 ++- vortex-array/src/expr/proto.rs | 69 +++++ vortex-array/src/expr/scope.rs | 22 ++ vortex-array/src/expr/traversal/mod.rs | 18 +- vortex-array/src/expression.rs | 16 + vortex-proto/proto/expr.proto | 5 + vortex-proto/src/generated/vortex.expr.rs | 6 + 16 files changed, 696 insertions(+), 61 deletions(-) create mode 100644 vortex-array/src/expr/lambda.rs diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index eacece3489e..64ec360b52a 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -42,7 +42,7 @@ pub fn make_free_field_annotator( ) -> impl AnnotationFn { move |expr: &Expression| match expr { Expression::Root => scope.names().iter().cloned().collect(), - Expression::Variable(_) => vec![], + Expression::Lambda(_) | Expression::Variable(_) => vec![], Expression::Scalar { scalar_fn, children, @@ -72,7 +72,7 @@ pub fn make_bound_free_field_annotator( ) -> impl AnnotationFn { move |expr: &BoundExpression| match expr { BoundExpression::Root { .. } => scope.names().iter().cloned().collect(), - BoundExpression::Variable(_) => vec![], + BoundExpression::Lambda(_) | BoundExpression::Variable(_) => vec![], BoundExpression::Scalar { scalar_fn, children, diff --git a/vortex-array/src/expr/analysis/infallible.rs b/vortex-array/src/expr/analysis/infallible.rs index 46b7a5af697..82a8c236a89 100644 --- a/vortex-array/src/expr/analysis/infallible.rs +++ b/vortex-array/src/expr/analysis/infallible.rs @@ -17,8 +17,8 @@ pub fn label_infallible(expr: &Expression) -> BooleanLabels<'_> { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_infallible(), // The scope itself cannot fail. Expression::Root => true, - // References are vacuously infallible. - Expression::Variable(_) => true, + // Fallibility is determined by the enclosing HOF. + Expression::Lambda(_) | Expression::Variable(_) => true, }, |acc, &child| acc & child, ) diff --git a/vortex-array/src/expr/analysis/strict.rs b/vortex-array/src/expr/analysis/strict.rs index 5e72e5541c8..008a6e1eaa7 100644 --- a/vortex-array/src/expr/analysis/strict.rs +++ b/vortex-array/src/expr/analysis/strict.rs @@ -14,8 +14,9 @@ pub fn label_strict(expr: &Expression) -> BooleanLabels<'_> { expr, |expr| match expr { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_strict(), - // Vacuously strict. - Expression::Root | Expression::Variable(_) => true, + Expression::Root => true, + // Strictness is determined by the enclosing HOF. + Expression::Variable(_) | Expression::Lambda(_) => true, }, |acc, &child| acc & child, ) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index f8be1d2c9e7..e4cbce4f302 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -13,10 +13,12 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_session::VortexSession; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::Lambda; use crate::expr::display::DisplayTreeExpr; use crate::expr::scope::Scope; use crate::expr::scope::VariableRef; @@ -48,6 +50,11 @@ pub enum BoundExpression { /// consumers from destructuring a `BoundExpression` by value. children: Arc>, }, + /// A bound lambda whose dtype is the dtype of its body. + /// + /// A lambda is not independently executable; only an enclosing higher-order function may + /// close it over captures and apply it to arguments. + Lambda(BoundLambda), /// The scope itself. Its dtype is the scope's root dtype. Root { /// The dtype this node evaluates to. @@ -88,6 +95,137 @@ impl Display for BoundVariable { } } +/// A lambda whose parameters and body have been resolved against a lexical [`Scope`]. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct BoundLambda { + params: Box<[Variable]>, + param_dtypes: Box<[DType]>, + param_refs: Box<[VariableRef]>, + parameter_frame: usize, + body: Arc, +} + +impl BoundLambda { + /// Bind `lambda` against a scope containing its parameter bindings in the innermost frame. + /// + /// The enclosing higher-order function determines the parameter dtypes and constructs this + /// scope before binding the lambda. + pub fn bind(lambda: &Lambda, scope: &Scope) -> VortexResult { + vortex_ensure!( + scope.depth() > 0, + "lambda parameters must be bound in a lexical frame" + ); + let parameter_frame = scope.depth() - 1; + let parameter_bindings = lambda + .params() + .iter() + .map(|param| { + scope + .resolve(param) + .map(|(dtype, variable_ref)| (dtype.clone(), variable_ref)) + .ok_or_else(|| { + vortex_err!("lambda parameter '{param}' is not bound in its scope") + }) + }) + .collect::>>()?; + vortex_ensure!( + parameter_bindings + .iter() + .all(|(_, variable_ref)| variable_ref.frame() == parameter_frame), + "lambda parameters must be bound in the innermost lexical frame" + ); + + Ok(Self { + params: lambda.params().into(), + param_dtypes: parameter_bindings + .iter() + .map(|(dtype, _)| dtype.clone()) + .collect(), + param_refs: parameter_bindings + .into_iter() + .map(|(_, variable_ref)| variable_ref) + .collect(), + parameter_frame, + body: Arc::new(lambda.body().bind_scope(scope)?), + }) + } + + /// The variables this lambda binds, in declaration order. + pub fn params(&self) -> &[Variable] { + &self.params + } + + /// The dtypes of the parameters, in declaration order. + pub fn param_dtypes(&self) -> &[DType] { + &self.param_dtypes + } + + /// The lexical locations assigned to the parameters, in declaration order. + pub fn param_refs(&self) -> &[VariableRef] { + &self.param_refs + } + + /// The lexical frame containing the parameters. + pub fn parameter_frame(&self) -> usize { + self.parameter_frame + } + + /// The bound body. + pub fn body(&self) -> &BoundExpression { + &self.body + } + + /// The dtype the body evaluates to. + pub fn body_dtype(&self) -> &DType { + self.body.dtype() + } + + /// The outer lexical bindings read by this lambda body. + pub fn free_variables(&self) -> Vec { + fn collect( + expression: &BoundExpression, + parameter_frame: usize, + variables: &mut Vec, + ) { + match expression { + BoundExpression::Variable(variable) + if variable.variable_ref().frame() < parameter_frame + && !variables.contains(&variable.variable_ref()) => + { + variables.push(variable.variable_ref()); + } + BoundExpression::Scalar { children, .. } => { + for child in children.iter() { + collect(child, parameter_frame, variables); + } + } + BoundExpression::Lambda(_) + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => {} + } + } + + let mut variables = Vec::new(); + collect(&self.body, self.parameter_frame, &mut variables); + variables.sort_by_key(|variable_ref| (variable_ref.frame(), variable_ref.slot())); + variables + } + + fn take_body(&mut self) -> Option { + Arc::try_unwrap(std::mem::replace( + &mut self.body, + Arc::new(BoundExpression::new_root(DType::Null)), + )) + .ok() + } +} + +impl Display for BoundLambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "({}) -> {}", self.params.iter().join(", "), self.body) + } +} + /// A bound-expression wrapper that compares shared tree identity instead of structure. #[derive(Clone, Debug)] pub struct ExactBoundExpr(pub BoundExpression); @@ -115,9 +253,11 @@ impl PartialEq for ExactBoundExpr { && Arc::ptr_eq(lhs_children, rhs_children) && lhs_dtype == rhs_dtype } + (BoundExpression::Lambda(lhs), BoundExpression::Lambda(rhs)) => lhs == rhs, (BoundExpression::Variable(lhs), BoundExpression::Variable(rhs)) => lhs == rhs, (BoundExpression::Root { .. }, _) | (BoundExpression::Scalar { .. }, _) + | (BoundExpression::Lambda(_), _) | (BoundExpression::Variable(_), _) => false, } } @@ -136,6 +276,10 @@ impl Hash for ExactBoundExpr { variable.variable().hash(state); variable.variable_ref().hash(state); } + BoundExpression::Lambda(lambda) => { + state.write_u8(3); + lambda.hash(state); + } BoundExpression::Scalar { scalar_fn, children, @@ -170,6 +314,10 @@ impl BoundExpression { scalar_fn.signature().arity(), children.len() ); + vortex_ensure!( + children.iter().all(|child| !child.is_lambda()), + "a scalar function cannot take a lambda as an ordinary argument" + ); let arg_dtypes = children .iter() @@ -194,7 +342,9 @@ impl BoundExpression { BoundExpression::Scalar { scalar_fn, .. } => { Self::try_new_vec(scalar_fn.clone(), children) } - BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + BoundExpression::Lambda(_) + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => { vortex_ensure!( children.is_empty(), "{self} cannot have {} children", @@ -209,15 +359,18 @@ impl BoundExpression { pub fn dtype(&self) -> &DType { match self { Self::Scalar { dtype, .. } | Self::Root { dtype } => dtype, + Self::Lambda(lambda) => lambda.body_dtype(), Self::Variable(variable) => variable.dtype(), } } - /// The bound children of this node, in argument order. Empty for leaf nodes. + /// The ordinary bound children of this node, in argument order. + /// + /// A bound lambda body is available through [`BoundLambda::body`] instead. pub fn children(&self) -> &[BoundExpression] { match self { Self::Scalar { children, .. } => children.as_slice(), - Self::Root { .. } | Self::Variable(_) => &[], + Self::Lambda(_) | Self::Root { .. } | Self::Variable(_) => &[], } } @@ -230,15 +383,28 @@ impl BoundExpression { pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root { .. } | Self::Variable(_) => None, + Self::Lambda(_) | Self::Root { .. } | Self::Variable(_) => None, + } + } + + /// Return this node's bound lambda, if it is a lambda. + pub fn as_lambda(&self) -> Option<&BoundLambda> { + match self { + Self::Lambda(lambda) => Some(lambda), + Self::Scalar { .. } | Self::Root { .. } | Self::Variable(_) => None, } } + /// Whether this node is a bound lambda. + pub fn is_lambda(&self) -> bool { + self.as_lambda().is_some() + } + /// Return this node's bound variable, if it is a variable. pub fn as_variable(&self) -> Option<&BoundVariable> { match self { Self::Variable(variable) => Some(variable), - Self::Scalar { .. } | Self::Root { .. } => None, + Self::Lambda(_) | Self::Scalar { .. } | Self::Root { .. } => None, } } @@ -316,6 +482,7 @@ impl Display for BoundExpression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + Self::Lambda(lambda) => Display::fmt(lambda, f), Self::Root { .. } => f.write_str("$"), Self::Variable(variable) => write!(f, "${variable}"), } @@ -346,6 +513,9 @@ impl Expression { variable_ref, })) } + Expression::Lambda(_) => { + vortex_bail!("a lambda can be bound only as an argument to a higher-order function") + } Expression::Scalar { scalar_fn, children, @@ -363,19 +533,34 @@ impl Expression { /// Iterative drop to avoid stack overflows on deep trees. impl Drop for BoundExpression { fn drop(&mut self) { - let Self::Scalar { children, .. } = self else { - return; - }; - let Some(children) = Arc::get_mut(children) else { - return; - }; + let mut to_drop = Vec::new(); + match self { + Self::Scalar { children, .. } => { + if let Some(children) = Arc::get_mut(children) { + to_drop.append(children); + } + } + Self::Lambda(lambda) => { + if let Some(body) = lambda.take_body() { + to_drop.push(body); + } + } + Self::Root { .. } | Self::Variable(_) => return, + } - let mut to_drop = std::mem::take(children); while let Some(mut child) = to_drop.pop() { - if let BoundExpression::Scalar { children, .. } = &mut child - && let Some(grandchildren) = Arc::get_mut(children) - { - to_drop.append(grandchildren); + match &mut child { + BoundExpression::Scalar { children, .. } => { + if let Some(grandchildren) = Arc::get_mut(children) { + to_drop.append(grandchildren); + } + } + BoundExpression::Lambda(lambda) => { + if let Some(body) = lambda.take_body() { + to_drop.push(body); + } + } + BoundExpression::Root { .. } | BoundExpression::Variable(_) => {} } } } @@ -392,6 +577,7 @@ mod tests { use crate::expr::col; use crate::expr::eq; use crate::expr::is_not_null; + use crate::expr::lambda; use crate::expr::lit; use crate::expr::root; use crate::expr::test_harness::struct_dtype; @@ -515,6 +701,76 @@ mod tests { Ok(()) } + #[test] + fn lambda_signature_comes_from_its_parameter_frame() -> VortexResult<()> { + let expression = lambda(["value"], var("value"))?; + let value_dtype = DType::Primitive(PType::I64, Nullability::Nullable); + let lambda_scope = + scope().with_bindings([(Variable::new("value"), value_dtype.clone())])?; + + assert!(expression.return_dtype(&struct_dtype()).is_err()); + assert!(expression.bind_scope(&lambda_scope).is_err()); + + let lambda = expression + .as_lambda() + .vortex_expect("the lambda factory must produce lambda syntax"); + let bound = BoundLambda::bind(lambda, &lambda_scope)?; + + assert_eq!(bound.params(), &[Variable::new("value")]); + assert_eq!(bound.param_dtypes(), std::slice::from_ref(&value_dtype)); + assert_eq!(bound.parameter_frame(), 0); + assert_eq!(bound.param_refs()[0].frame(), 0); + assert_eq!(bound.param_refs()[0].slot(), 0); + assert_eq!(bound.body_dtype(), &value_dtype); + assert_eq!( + bound + .body() + .as_variable() + .vortex_expect("the lambda body must resolve its parameter") + .variable_ref(), + bound.param_refs()[0] + ); + Ok(()) + } + + #[test] + fn lambda_parameter_must_be_in_the_innermost_frame() -> VortexResult<()> { + let expression = lambda(["value"], var("value"))?; + let lambda = expression + .as_lambda() + .vortex_expect("the lambda factory must produce lambda syntax"); + let scope = scope() + .with_bindings([(Variable::new("value"), DType::Null)])? + .with_bindings([(Variable::new("other"), DType::Null)])?; + + assert!(BoundLambda::bind(lambda, &scope).is_err()); + Ok(()) + } + + #[test] + fn lambda_tracks_outer_captures_separately_from_parameters() -> VortexResult<()> { + let expression = lambda(["parameter"], eq(var("captured"), var("parameter")))?; + let lambda = expression + .as_lambda() + .vortex_expect("the lambda factory must produce lambda syntax"); + let value_dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let scope = scope() + .with_bindings([(Variable::new("captured"), value_dtype.clone())])? + .with_bindings([(Variable::new("parameter"), value_dtype)])?; + + let bound = BoundLambda::bind(lambda, &scope)?; + + assert_eq!(bound.parameter_frame(), 1); + assert_eq!(bound.param_refs().len(), 1); + assert_eq!(bound.param_refs()[0].frame(), 1); + assert_eq!(bound.param_refs()[0].slot(), 0); + let captures = bound.free_variables(); + assert_eq!(captures.len(), 1); + assert_eq!(captures[0].frame(), 0); + assert_eq!(captures[0].slot(), 0); + Ok(()) + } + #[test] fn clone_shares_children() -> VortexResult<()> { let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 1ab745c0169..2a151716877 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -62,50 +62,77 @@ const ROOT_DISPLAY: &str = "vortex.root()"; impl DisplayTreeNode for Expression { fn tree_children(&self) -> &[Self] { - Expression::children(self) + match self { + Expression::Lambda(lambda) => std::slice::from_ref(lambda.body()), + Expression::Scalar { .. } | Expression::Root | Expression::Variable(_) => { + Expression::children(self) + } + } } fn tree_child_name(&self, index: usize) -> ChildName { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), - Expression::Root => unreachable!("the scope root has no children"), - Expression::Variable(_) => unreachable!("a variable has no children"), + Expression::Lambda(_) => ChildName::from("body"), + Expression::Root | Expression::Variable(_) => { + unreachable!("a leaf expression has no children") + } } } fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + Expression::Lambda(lambda) => fmt_lambda(lambda.params(), f), Expression::Root => write!(f, "{ROOT_DISPLAY}"), - Expression::Variable(var) => write!(f, "vortex.var({var})"), + Expression::Variable(variable) => write!(f, "vortex.var({variable})"), } } } impl DisplayTreeNode for BoundExpression { fn tree_children(&self) -> &[Self] { - BoundExpression::children(self) + match self { + BoundExpression::Lambda(lambda) => std::slice::from_ref(lambda.body()), + BoundExpression::Scalar { .. } + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => BoundExpression::children(self), + } } fn tree_child_name(&self, index: usize) -> ChildName { match self { BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), - BoundExpression::Root { .. } => unreachable!("the scope root has no children"), - BoundExpression::Variable { .. } => unreachable!("a variable has no children"), + BoundExpression::Lambda(_) => ChildName::from("body"), + BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + unreachable!("a leaf bound expression has no children") + } } } fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { BoundExpression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + BoundExpression::Lambda(lambda) => fmt_lambda(lambda.params(), f), BoundExpression::Root { .. } => write!(f, "{ROOT_DISPLAY}"), - BoundExpression::Variable(var) => write!(f, "vortex.var({var})"), + BoundExpression::Variable(variable) => write!(f, "vortex.var({variable})"), } } } pub struct DisplayTreeExpr<'a, T: ?Sized = Expression>(pub &'a T); +fn fmt_lambda(params: &[P], f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "vortex.lambda(params: [")?; + for (index, param) in params.iter().enumerate() { + if index > 0 { + write!(f, ", ")?; + } + write!(f, "vortex.var({param})")?; + } + write!(f, "])") +} + impl TreeDisplayAdapter for DisplayTreeExpr<'_, T> { type Context = (); type Node = T; @@ -141,6 +168,8 @@ impl Display for DisplayTreeExpr<'_, T> { #[cfg(test)] mod tests { + use vortex_error::VortexResult; + use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -150,6 +179,7 @@ mod tests { use crate::expr::eq; use crate::expr::get_item; use crate::expr::gt; + use crate::expr::lambda; use crate::expr::lit; use crate::expr::not; use crate::expr::pack; @@ -203,6 +233,18 @@ mod tests { assert_snapshot!(var("value").display_tree().to_string(), @"vortex.var(value)"); } + #[test] + fn test_display_tree_lambda() -> VortexResult<()> { + use insta::assert_snapshot; + + let expression = lambda(["value"], var("value"))?; + assert_snapshot!(expression.display_tree().to_string(), @r" + vortex.lambda(params: [vortex.var(value)]) + └── body: vortex.var(value) + "); + Ok(()) + } + #[test] fn test_display_tree_literal() { use insta::assert_snapshot; diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index cf9ce929fb0..dfe1ec326f2 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -18,6 +18,7 @@ use vortex_error::vortex_ensure; use crate::dtype::DType; use crate::expr::display::DisplayTreeExpr; use crate::expr::is_not_null; +use crate::expr::lambda::Lambda; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; use crate::expr::variable::Variable; @@ -31,8 +32,9 @@ const NO_CHILDREN: &[Expression] = &[]; /// /// Most nodes are a scalar function applied to child expressions. [`Expression::Root`] is the scope /// itself: a language primitive rather than a registered function, because its dtype comes from the -/// scope rather than from children and it is not executable. [`Expression::Variable`] is likewise -/// resolved by the surrounding scope when the expression is bound. +/// scope rather than from children and it is not executable. [`Expression::Variable`] is resolved +/// by the surrounding scope when the expression is bound, while [`Expression::Lambda`] introduces +/// a new lexical frame when an enclosing higher-order function binds it. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Expression { /// A scalar function applied to child expressions. @@ -42,6 +44,8 @@ pub enum Expression { /// Any children of this expression. children: Arc>, }, + /// Lambda syntax owned and invoked by an enclosing higher-order function. + Lambda(Lambda), /// The full scope of the expression evaluation. Root, /// A named value resolved from the surrounding scope when the expression is bound. @@ -62,6 +66,10 @@ impl Expression { scalar_fn.signature().arity(), children.len() ); + vortex_ensure!( + children.iter().all(|child| !child.is_lambda()), + "a scalar function cannot take a lambda as an ordinary argument" + ); Ok(Self::Scalar { scalar_fn, @@ -78,15 +86,28 @@ impl Expression { pub fn as_variable(&self) -> Option<&Variable> { match self { Self::Variable(variable) => Some(variable), - Self::Root | Self::Scalar { .. } => None, + Self::Lambda(_) | Self::Root | Self::Scalar { .. } => None, + } + } + + /// The lambda syntax held by this node, if it is a lambda. + pub fn as_lambda(&self) -> Option<&Lambda> { + match self { + Self::Lambda(lambda) => Some(lambda), + Self::Root | Self::Scalar { .. } | Self::Variable(_) => None, } } + /// Whether this node is lambda syntax. + pub fn is_lambda(&self) -> bool { + self.as_lambda().is_some() + } + /// Returns the scalar fn for this expression, or `None` if it is not a scalar node. pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root | Self::Variable(_) => None, + Self::Lambda(_) | Self::Root | Self::Variable(_) => None, } } @@ -110,11 +131,13 @@ impl Expression { .vortex_expect("Expression options type mismatch") } - /// Returns the children of this expression. + /// Returns the ordinary scalar children of this expression. + /// + /// A lambda body is binder-owned syntax and is available through [`Lambda::body`] instead. pub fn children(&self) -> &[Expression] { match self { Self::Scalar { children, .. } => children.as_slice(), - Self::Root | Self::Variable(_) => NO_CHILDREN, + Self::Lambda(_) | Self::Root | Self::Variable(_) => NO_CHILDREN, } } @@ -123,14 +146,14 @@ impl Expression { &self.children()[n] } - /// Replace the children of this expression with the provided new children. + /// Replace the ordinary scalar children of this expression with the provided new children. pub fn with_children( self, children: impl IntoIterator, ) -> VortexResult { let children = Vec::from_iter(children); match &self { - Self::Root | Self::Variable(_) => { + Self::Lambda(_) | Self::Root | Self::Variable(_) => { vortex_ensure!( children.is_empty(), "Expression arity mismatch: a leaf expects 0 children but got {}", @@ -163,6 +186,9 @@ impl Expression { Self::Variable(variable) => { vortex_bail!("cannot determine dtype of unbound variable '{variable}'") } + Self::Lambda(_) => vortex_bail!( + "a lambda has no standalone dtype; it must be bound by a higher-order function" + ), Self::Scalar { scalar_fn, children, @@ -185,6 +211,9 @@ impl Expression { Self::Root => Ok(Self::Root), // The binding supplies the variable's actual dtype later. Self::Variable(_) => Ok(is_not_null(self.clone())), + Self::Lambda(_) => vortex_bail!( + "a lambda has no standalone validity expression; it must be applied by a higher-order function" + ), Self::Scalar { scalar_fn, .. } => scalar_fn.validity(self), } } @@ -197,6 +226,7 @@ impl Expression { match self { Self::Root => write!(f, "$"), Self::Variable(variable) => write!(f, "${variable}"), + Self::Lambda(lambda) => Display::fmt(lambda, f), Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), } } @@ -319,26 +349,37 @@ impl Drop for DropDepthGuard { impl Drop for Expression { fn drop(&mut self) { - let Self::Scalar { children, .. } = self else { - return; - }; - let Some(children) = Arc::get_mut(children) else { - return; - }; - if children.is_empty() { - return; + let mut children_to_drop = Vec::new(); + match self { + Self::Scalar { children, .. } => { + if let Some(children) = Arc::get_mut(children) { + children_to_drop.append(children); + } + } + Self::Lambda(lambda) => { + if let Some(body) = lambda.take_body() { + children_to_drop.push(body); + } + } + Self::Root | Self::Variable(_) => return, } - let mut children_to_drop = std::mem::take(children); - match DropDepthGuard::enter() { Some(_guard) => drop(children_to_drop), None => { while let Some(mut child) = children_to_drop.pop() { - if let Self::Scalar { children, .. } = &mut child - && let Some(expr_children) = Arc::get_mut(children) - { - children_to_drop.append(expr_children); + match &mut child { + Self::Scalar { children, .. } => { + if let Some(expr_children) = Arc::get_mut(children) { + children_to_drop.append(expr_children); + } + } + Self::Lambda(lambda) => { + if let Some(body) = lambda.take_body() { + children_to_drop.push(body); + } + } + Self::Root | Self::Variable(_) => {} } } } @@ -351,8 +392,10 @@ mod tests { use std::thread; use super::*; + use crate::expr::lambda; use crate::expr::lit; use crate::expr::not; + use crate::expr::var; fn deep_expression(depth: usize) -> Expression { let mut expr = lit(true); @@ -388,4 +431,27 @@ mod tests { assert_eq!(shared.children().len(), 1); } + + #[test] + fn lambda_has_no_standalone_array_type_or_validity() -> VortexResult<()> { + let expression = lambda(["value"], var("value"))?; + let root_dtype = DType::Bool(crate::dtype::Nullability::NonNullable); + + assert!(expression.return_dtype(&root_dtype).is_err()); + assert!(expression.validity().is_err()); + Ok(()) + } + + #[test] + fn scalar_functions_reject_lambda_children() -> VortexResult<()> { + let scalar = not(lit(true)); + let scalar_fn = scalar + .as_scalar() + .vortex_expect("not must be a scalar function") + .clone(); + let lambda = lambda(["value"], var("value"))?; + + assert!(Expression::try_new(scalar_fn, [lambda]).is_err()); + Ok(()) + } } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 1391ce2b427..256872ad956 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -17,6 +17,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::Lambda; use crate::expr::Variable; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -76,6 +77,16 @@ pub fn var(name: impl AsRef) -> Expression { Variable::new(name).into() } +/// Creates a lambda expression binding `params` over `body`. +/// +/// Returns an error when a parameter name is repeated in the same lambda. +pub fn lambda( + params: impl IntoIterator>, + body: Expression, +) -> VortexResult { + Ok(Lambda::try_new(params, body)?.into()) +} + /// Return whether the expression is a root expression. pub fn is_root(expr: &Expression) -> bool { expr.is_root() diff --git a/vortex-array/src/expr/lambda.rs b/vortex-array/src/expr/lambda.rs new file mode 100644 index 00000000000..0d633b32abe --- /dev/null +++ b/vortex-array/src/expr/lambda.rs @@ -0,0 +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::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::expr::Expression; +use crate::expr::Variable; + +/// A function-like expression that binds `params` in `body`. +/// +/// A lambda is binder-owned syntax rather than an array-valued expression. A higher-order function +/// supplies the parameter dtypes and invocation semantics before the body can be bound or applied. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Lambda { + params: Arc>, + body: Arc, +} + +impl Lambda { + /// Create a lambda binding `params` over `body`. + /// + /// Returns an error when a parameter name is repeated. + pub fn try_new( + params: impl IntoIterator>, + body: Expression, + ) -> VortexResult { + let mut variables = Vec::new(); + let mut seen = HashSet::new(); + + for param in params { + let variable = param.into(); + if !seen.insert(variable.clone()) { + vortex_bail!("duplicate lambda parameter '{variable}'"); + } + variables.push(variable); + } + + Ok(Self { + params: Arc::new(variables), + body: Arc::new(body), + }) + } + + /// The variables this lambda binds, in declaration order. + pub fn params(&self) -> &[Variable] { + &self.params + } + + /// The expression evaluated under the parameter bindings. + pub fn body(&self) -> &Expression { + &self.body + } + + /// Take the body when this lambda is its sole owner. + /// + /// This supports expression's iterative drop implementation for deeply nested binder bodies. + pub(crate) fn take_body(&mut self) -> Option { + Arc::try_unwrap(std::mem::replace( + &mut self.body, + Arc::new(Expression::Root), + )) + .ok() + } +} + +impl Display for Lambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "({}) -> {}", self.params.iter().join(", "), self.body) + } +} + +impl From for Expression { + fn from(lambda: Lambda) -> Self { + Self::Lambda(lambda) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn duplicate_parameters_are_rejected() { + let error = Lambda::try_new(["x", "x"], Expression::Root) + .expect_err("duplicate parameters must be rejected"); + assert!(error.to_string().contains("'x'")); + } +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index aff195d818c..d8d16aa695e 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -62,6 +62,7 @@ pub(crate) mod expression; mod exprs; pub(crate) mod field; pub mod forms; +pub mod lambda; mod optimize; pub mod proto; pub mod scope; @@ -96,6 +97,7 @@ pub use exprs::ilike; pub use exprs::is_not_null; pub use exprs::is_null; pub use exprs::is_root; +pub use exprs::lambda; pub use exprs::like; pub use exprs::list_contains; pub use exprs::list_length; @@ -122,6 +124,7 @@ pub use exprs::union_child_validities; pub use exprs::var; pub use exprs::variant_get; pub use exprs::zip_expr; +pub use lambda::*; pub use scope::*; pub use variable::*; @@ -167,6 +170,7 @@ impl PartialEq for ExactExpr { match (&self.0, &other.0) { (Expression::Root, Expression::Root) => true, (Expression::Variable(lhs), Expression::Variable(rhs)) => lhs == rhs, + (Expression::Lambda(lhs), Expression::Lambda(rhs)) => lhs == rhs, ( Expression::Scalar { scalar_fn: lhs_fn, @@ -179,6 +183,7 @@ impl PartialEq for ExactExpr { ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children), (Expression::Root, _) | (Expression::Scalar { .. }, _) + | (Expression::Lambda(_), _) | (Expression::Variable(_), _) => false, } } @@ -193,6 +198,10 @@ impl Hash for ExactExpr { state.write_u8(2); variable.hash(state); } + Expression::Lambda(lambda) => { + state.write_u8(3); + lambda.hash(state); + } Expression::Scalar { scalar_fn, children, diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index b47350007db..aecd2c74114 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -32,7 +32,7 @@ impl Expression { fn simplify_untyped_node(&self) -> VortexResult> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify_untyped(self), - Expression::Root | Expression::Variable(_) => Ok(None), + Expression::Lambda(_) | Expression::Root | Expression::Variable(_) => Ok(None), } } @@ -40,7 +40,7 @@ impl Expression { fn simplify_node(&self, ctx: &dyn SimplifyCtx) -> VortexResult> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self, ctx), - Expression::Root | Expression::Variable(_) => Ok(None), + Expression::Lambda(_) | Expression::Root | Expression::Variable(_) => Ok(None), } } @@ -51,7 +51,7 @@ impl Expression { ) -> VortexResult>> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node), - Expression::Root | Expression::Variable(_) => Ok(None), + Expression::Lambda(_) | Expression::Root | Expression::Variable(_) => Ok(None), } } @@ -71,6 +71,15 @@ impl Expression { loop_counter += 1; let expr = current.as_ref().unwrap_or(self); + match expr { + Expression::Variable(variable) => vortex_bail!( + "cannot optimize variable '{variable}' outside a higher-order function" + ), + Expression::Lambda(_) => { + vortex_bail!("cannot optimize a lambda outside a higher-order function") + } + Expression::Root | Expression::Scalar { .. } => {} + } let mut changed = false; // Try simplify_untyped @@ -190,6 +199,9 @@ impl SimplifyCtx for SimplifyCache<'_> { Expression::Variable(variable) => { vortex_bail!("cannot determine dtype of unbound variable '{variable}'") } + Expression::Lambda(_) => vortex_bail!( + "cannot determine the standalone dtype of a lambda; it must be bound by a higher-order function" + ), Expression::Scalar { scalar_fn, children, @@ -222,10 +234,12 @@ mod tests { use crate::expr::cast; use crate::expr::eq; use crate::expr::get_item; + use crate::expr::lambda; use crate::expr::lit; use crate::expr::lt_eq; use crate::expr::or; use crate::expr::root; + use crate::expr::var; use crate::scalar::Scalar; use crate::scalar_fn::fns::literal::Literal; @@ -278,4 +292,20 @@ mod tests { assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable)); Ok(()) } + + #[test] + fn detached_lexical_nodes_error_when_optimization_visits_them() -> VortexResult<()> { + assert!(var("value").optimize(&DType::Null).is_err()); + assert!( + lambda(["value"], var("value"))? + .optimize(&DType::Null) + .is_err() + ); + assert!( + eq(var("value"), lit(42_i32)) + .optimize_recursive(&DType::Null) + .is_err() + ); + Ok(()) + } } diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index b6ec9bc771b..8e76a70011f 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -10,6 +10,7 @@ use vortex_proto::expr as pb; use vortex_session::VortexSession; use crate::expr::Expression; +use crate::expr::Lambda; use crate::expr::Variable; use crate::scalar_fn::ForeignScalarFnVTable; use crate::scalar_fn::ScalarFnId; @@ -28,6 +29,43 @@ pub(crate) const ROOT_ID: &str = "vortex.root"; /// The wire id for [`Expression::Variable`]. pub(crate) const VARIABLE_ID: &str = "vortex.var"; +/// The wire id for [`Expression::Lambda`]. +pub(crate) const LAMBDA_ID: &str = "vortex.lambda"; + +impl Lambda { + /// Serialize this lambda to its protobuf representation. + fn serialize_proto(&self) -> VortexResult { + Ok(pb::Expr { + id: LAMBDA_ID.to_string(), + children: vec![self.body().serialize_proto()?], + metadata: Some( + pb::LambdaOpts { + params: self + .params() + .iter() + .map(|variable| variable.name().to_string()) + .collect(), + } + .encode_to_vec(), + ), + }) + } + + /// Deserialize a lambda expression whose id is [`LAMBDA_ID`]. + fn from_proto(expr: &pb::Expr, session: &VortexSession) -> VortexResult { + vortex_ensure!( + expr.children.len() == 1, + "a lambda must have exactly one child, its body, got {}", + expr.children.len() + ); + let options = pb::LambdaOpts::decode(expr.metadata())?; + Self::try_new( + options.params.into_iter().map(Variable::new), + Expression::from_proto(&expr.children[0], session)?, + ) + } +} + impl Variable { /// Serialize this variable to its protobuf representation. fn serialize_proto(&self) -> pb::Expr { @@ -88,6 +126,7 @@ impl ExprSerializeProtoExt for Expression { metadata: Some(vec![]), }), Expression::Variable(variable) => Ok(variable.serialize_proto()), + Expression::Lambda(lambda) => lambda.serialize_proto(), Expression::Scalar { scalar_fn, children, @@ -112,6 +151,10 @@ impl Expression { return Ok(Variable::from_proto(expr)?.into()); } + if expr.id == LAMBDA_ID { + return Ok(Lambda::from_proto(expr, session)?.into()); + } + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let expr_id = ScalarFnId::new(expr.id.as_str()); let children = expr @@ -155,6 +198,7 @@ mod tests { use crate::expr::between; use crate::expr::eq; use crate::expr::get_item; + use crate::expr::lambda; use crate::expr::lit; use crate::expr::or; use crate::expr::root; @@ -211,6 +255,31 @@ mod tests { Ok(()) } + #[test] + fn lambda_serde() -> VortexResult<()> { + let expression = lambda(["x", "y"], eq(var("x"), var("y")))?; + let encoded = expression.serialize_proto()?.encode_to_vec(); + let proto = pb::Expr::decode(encoded.as_slice())?; + + assert_eq!( + Expression::from_proto(&proto, &array_session())?, + expression + ); + Ok(()) + } + + #[test] + fn lambda_requires_exactly_one_body() -> VortexResult<()> { + let mut without_body = lambda(["x"], var("x"))?.serialize_proto()?; + without_body.children.clear(); + assert!(Expression::from_proto(&without_body, &array_session()).is_err()); + + let mut with_two_bodies = lambda(["x"], var("x"))?.serialize_proto()?; + with_two_bodies.children.push(root().serialize_proto()?); + assert!(Expression::from_proto(&with_two_bodies, &array_session()).is_err()); + Ok(()) + } + #[test] fn unknown_expression_id_allow_unknown() { let session = VortexSession::empty().with::(); diff --git a/vortex-array/src/expr/scope.rs b/vortex-array/src/expr/scope.rs index 3c71401d69a..d4168fa8b07 100644 --- a/vortex-array/src/expr/scope.rs +++ b/vortex-array/src/expr/scope.rs @@ -116,6 +116,17 @@ impl Scope { Ok(self.push_frame(Frame::try_new(bindings)?)) } + /// Return this scope with the same lexical frames and a different root dtype. + /// + /// Higher-order functions use this when a lambda has a new implicit input while retaining + /// access to bindings captured from surrounding scopes. + pub fn with_root(&self, root: DType) -> Self { + Self { + root, + frames: self.frames.clone(), + } + } + /// Resolve `name`, searching innermost-first so inner bindings shadow outer ones. pub fn resolve(&self, name: &Variable) -> Option<(&DType, VariableRef)> { self.frames @@ -217,4 +228,15 @@ mod tests { ); Ok(()) } + + #[test] + fn replacing_the_root_preserves_lexical_bindings() -> VortexResult<()> { + let variable = Variable::new("captured"); + let scope = Scope::new(i32_()).with_bindings([(variable.clone(), utf8())])?; + let with_new_root = scope.with_root(utf8()); + + assert_eq!(with_new_root.root(), &utf8()); + assert_eq!(with_new_root.resolve(&variable), scope.resolve(&variable)); + Ok(()) + } } diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index 987d420c9a0..839f8d7622f 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -535,7 +535,9 @@ impl Node for BoundExpression { ) -> VortexResult { let children = match self { BoundExpression::Scalar { children, .. } => children, - BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + BoundExpression::Lambda(_) + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => { return Ok(TraversalOrder::Continue); } }; @@ -556,7 +558,9 @@ impl Node for BoundExpression { ) -> VortexResult> { let children = match &self { BoundExpression::Scalar { children, .. } => children, - BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + BoundExpression::Lambda(_) + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => { return Ok(Transformed::no(self)); } }; @@ -599,16 +603,18 @@ impl Node for BoundExpression { fn iter_children(&self, f: impl FnOnce(&mut dyn Iterator) -> T) -> T { match self { BoundExpression::Scalar { children, .. } => f(&mut children.iter()), - BoundExpression::Root { .. } | BoundExpression::Variable(_) => { - f(&mut std::iter::empty()) - } + BoundExpression::Lambda(_) + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => f(&mut std::iter::empty()), } } fn children_count(&self) -> usize { match self { BoundExpression::Scalar { children, .. } => children.len(), - BoundExpression::Root { .. } | BoundExpression::Variable(_) => 0, + BoundExpression::Lambda(_) + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => 0, } } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index 83c2e26f804..d5cd410b63c 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -20,6 +20,9 @@ impl ArrayRef { pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult { match expr { BoundExpression::Root { .. } => Ok(self), + BoundExpression::Lambda(_) => { + vortex_bail!("cannot apply a lambda outside a higher-order function") + } BoundExpression::Variable(variable) => { vortex_bail!("cannot apply variable '{variable}' without a provided value") } @@ -35,6 +38,9 @@ impl ArrayRef { pub fn apply(self, expr: &Expression) -> VortexResult { match expr { Expression::Root => Ok(self), + Expression::Lambda(_) => { + vortex_bail!("cannot apply a lambda outside a higher-order function") + } Expression::Variable(variable) => { vortex_bail!("cannot apply unbound variable '{variable}'") } @@ -90,6 +96,7 @@ mod tests { use crate::IntoArray; use crate::expr::Scope; use crate::expr::Variable; + use crate::expr::lambda; use crate::expr::var; #[test] @@ -104,4 +111,13 @@ mod tests { assert!(root.apply_bound(&bound).is_err()); Ok(()) } + + #[test] + fn lambda_application_requires_a_higher_order_function() -> VortexResult<()> { + let root = buffer![1_i32, 2, 3].into_array(); + let expression = lambda(["value"], var("value"))?; + + assert!(root.apply(&expression).is_err()); + Ok(()) + } } diff --git a/vortex-proto/proto/expr.proto b/vortex-proto/proto/expr.proto index d506214c4a2..482bb204475 100644 --- a/vortex-proto/proto/expr.proto +++ b/vortex-proto/proto/expr.proto @@ -130,3 +130,8 @@ message CaseWhenOpts { message VariableOpts { string name = 1; } + +// Options for `vortex.lambda`. The body is the expression's single child. +message LambdaOpts { + repeated string params = 1; +} diff --git a/vortex-proto/src/generated/vortex.expr.rs b/vortex-proto/src/generated/vortex.expr.rs index 2660ce852b7..605f30c07ad 100644 --- a/vortex-proto/src/generated/vortex.expr.rs +++ b/vortex-proto/src/generated/vortex.expr.rs @@ -213,3 +213,9 @@ pub struct VariableOpts { #[prost(string, tag = "1")] pub name: ::prost::alloc::string::String, } +/// Options for `vortex.lambda`. The body is the expression's single child. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct LambdaOpts { + #[prost(string, repeated, tag = "1")] + pub params: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +}