From 78ce19363f7a1cf50722443fcbfff18bb8966a2b Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 12 Oct 2026 15:56:47 -0400 Subject: [PATCH 1/2] higher-order expressions with lambdas Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 44 +++++ vortex-array/src/expr/display.rs | 28 +++- vortex-array/src/expr/expression.rs | 84 ++++++++-- vortex-array/src/expr/exprs.rs | 12 ++ vortex-array/src/expr/lambda/bound.rs | 182 +++++++++++++++++++++ vortex-array/src/expr/lambda/mod.rs | 176 ++++++++++++++++++++ vortex-array/src/expr/lambda/tests.rs | 165 +++++++++++++++++++ vortex-array/src/expr/mod.rs | 3 + vortex-array/src/expr/proto.rs | 7 + vortex-array/src/expr/scope.rs | 2 +- vortex-array/src/expression.rs | 8 +- vortex-array/src/scalar_fn/erased.rs | 24 +++ vortex-array/src/scalar_fn/higher_order.rs | 38 +++++ vortex-array/src/scalar_fn/mod.rs | 3 + vortex-array/src/scalar_fn/typed.rs | 36 ++++ vortex-array/src/scalar_fn/vtable.rs | 10 ++ 16 files changed, 807 insertions(+), 15 deletions(-) create mode 100644 vortex-array/src/expr/lambda/bound.rs create mode 100644 vortex-array/src/expr/lambda/mod.rs create mode 100644 vortex-array/src/expr/lambda/tests.rs create mode 100644 vortex-array/src/scalar_fn/higher_order.rs diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index f5fce8f5130..9bf0fa6c0da 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -17,6 +17,7 @@ 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; @@ -73,6 +74,14 @@ pub struct BoundVariable { } impl BoundVariable { + pub(crate) fn new(dtype: DType, variable: Variable, variable_ref: VariableRef) -> Self { + Self { + dtype, + variable, + variable_ref, + } + } + /// The dtype of the value bound to this variable. pub fn dtype(&self) -> &DType { &self.dtype @@ -377,6 +386,11 @@ impl Expression { variable_ref, })) } + Expression::HigherOrder { + scalar_fn, + children, + lambdas, + } => bind_higher_order(scalar_fn, children, lambdas, scope), Expression::Scalar { scalar_fn, children, @@ -391,6 +405,36 @@ impl Expression { } } +/// Bind a higher-order node into an ordinary scalar node. +/// +/// The function's options take the bound lambdas, and its children are the bound value arguments +/// followed by every lambda's captures. +fn bind_higher_order( + scalar_fn: &ScalarFnRef, + children: &[Expression], + lambdas: &[Lambda], + scope: &Scope, +) -> VortexResult { + let mut children: Vec<_> = children + .iter() + .map(|child| child.bind_scope(scope)) + .try_collect()?; + let arg_dtypes = children + .iter() + .map(|child| child.dtype().clone()) + .collect_vec(); + + let mut bound_lambdas = Vec::with_capacity(lambdas.len()); + for (index, lambda) in lambdas.iter().enumerate() { + let param_dtypes = scalar_fn.lambda_params(index, lambda.params().len(), &arg_dtypes)?; + let (bound_lambda, captures) = lambda.bind(scope, param_dtypes)?; + bound_lambdas.push(bound_lambda); + children.extend(captures); + } + + BoundExpression::try_new(scalar_fn.with_lambdas(bound_lambdas)?, children) +} + /// Iterative drop to avoid stack overflows on deep trees. impl Drop for BoundExpression { fn drop(&mut self) { diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 1ab745c0169..f86f0428159 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -5,11 +5,13 @@ use std::fmt; use std::fmt::Display; use std::fmt::Formatter; +use itertools::Itertools; use vortex_utils::tree::TreeDisplayAdapter; use vortex_utils::tree::write_branch_tree; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::Lambda; use crate::scalar_fn::ChildName; pub enum DisplayFormat { @@ -29,13 +31,28 @@ pub trait ExprDisplay: Display { fn display_children_count(&self) -> usize; } +/// An unbound node displays its children followed by its lambdas. impl ExprDisplay for Expression { fn display_child(&self, index: usize) -> &dyn ExprDisplay { - Expression::child(self, index) + let children = self.children(); + match children.get(index) { + Some(child) => child, + None => &self.lambdas()[index - children.len()], + } } fn display_children_count(&self) -> usize { - self.children().len() + self.children().len() + self.lambdas().len() + } +} + +impl ExprDisplay for Lambda { + fn display_child(&self, _index: usize) -> &dyn ExprDisplay { + unreachable!("a lambda has no display children") + } + + fn display_children_count(&self) -> usize { + 0 } } @@ -67,7 +84,9 @@ impl DisplayTreeNode for Expression { fn tree_child_name(&self, index: usize) -> ChildName { match self { - Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), + Expression::Scalar { scalar_fn, .. } | Expression::HigherOrder { 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"), } @@ -76,6 +95,9 @@ impl DisplayTreeNode for Expression { fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + Expression::HigherOrder { + scalar_fn, lambdas, .. + } => write!(f, "{scalar_fn} {}", lambdas.iter().join(", ")), Expression::Root => write!(f, "{ROOT_DISPLAY}"), Expression::Variable(var) => write!(f, "vortex.var({var})"), } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index e1caac161fc..90f6b9c2cb0 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -16,6 +16,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use crate::dtype::DType; +use crate::expr::Lambda; use crate::expr::display::DisplayTreeExpr; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; @@ -31,7 +32,8 @@ 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. +/// resolved by the surrounding scope when the expression is bound. [`Expression::HigherOrder`] is a +/// scalar function that also takes [`Lambda`]s. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Expression { /// A scalar function applied to child expressions. @@ -45,6 +47,18 @@ pub enum Expression { Root, /// A named value resolved from the surrounding scope when the expression is bound. Variable(Variable), + /// A scalar function applied to child expressions and then to lambdas. + /// + /// The function implements [`HigherOrderVTable`](crate::scalar_fn::HigherOrderVTable). Binding + /// produces an ordinary scalar node whose options hold the bound lambdas. + HigherOrder { + /// The scalar fn for this node. + scalar_fn: ScalarFnRef, + /// The value arguments. + children: Arc>, + /// The lambda arguments, which follow the value arguments. + lambdas: Arc>, + }, } impl Expression { @@ -68,6 +82,31 @@ impl Expression { }) } + /// Create a node applying a higher-order scalar function to child expressions and lambdas. + pub fn try_new_higher_order( + scalar_fn: ScalarFnRef, + children: impl IntoIterator, + lambdas: impl IntoIterator, + ) -> VortexResult { + let children = Vec::from_iter(children); + vortex_ensure!( + scalar_fn.is_higher_order(), + "{} does not take lambda arguments", + scalar_fn.id() + ); + vortex_ensure!( + scalar_fn.signature().arity().matches(children.len()), + "Expression arity mismatch: expected {} children but got {}", + scalar_fn.signature().arity(), + children.len() + ); + Ok(Self::HigherOrder { + scalar_fn, + children: children.into(), + lambdas: Arc::new(lambdas.into_iter().collect()), + }) + } + /// Whether this expression is the scope root. pub fn is_root(&self) -> bool { matches!(self, Self::Root) @@ -77,15 +116,25 @@ impl Expression { pub fn as_variable(&self) -> Option<&Variable> { match self { Self::Variable(variable) => Some(variable), - Self::Root | Self::Scalar { .. } => None, + Self::Root | Self::Scalar { .. } | Self::HigherOrder { .. } => None, + } + } + + /// The lambdas of a higher-order node, and none for any other node. + pub fn lambdas(&self) -> &[Lambda] { + match self { + Self::HigherOrder { lambdas, .. } => lambdas.as_slice(), + Self::Root | Self::Variable(_) | Self::Scalar { .. } => &[], } } /// Returns the scalar fn for this expression, or `None` if it is not a scalar node. + /// + /// A higher-order node is not a scalar node, because its lambdas are not among its children. pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root | Self::Variable(_) => None, + Self::Root | Self::Variable(_) | Self::HigherOrder { .. } => None, } } @@ -112,7 +161,9 @@ impl Expression { /// Returns the children of this expression. pub fn children(&self) -> &[Expression] { match self { - Self::Scalar { children, .. } => children.as_slice(), + Self::Scalar { children, .. } | Self::HigherOrder { children, .. } => { + children.as_slice() + } Self::Root | Self::Variable(_) => NO_CHILDREN, } } @@ -137,16 +188,25 @@ impl Expression { ); Ok(self) } - Self::Scalar { scalar_fn, .. } => { + Self::Scalar { scalar_fn, .. } | Self::HigherOrder { scalar_fn, .. } => { vortex_ensure!( scalar_fn.signature().arity().matches(children.len()), "Expression arity mismatch: expected {} children but got {}", scalar_fn.signature().arity(), children.len() ); - Ok(Self::Scalar { - scalar_fn: scalar_fn.clone(), - children: children.into(), + let scalar_fn = scalar_fn.clone(); + let children = children.into(); + Ok(match &self { + Self::HigherOrder { lambdas, .. } => Self::HigherOrder { + scalar_fn, + children, + lambdas: Arc::clone(lambdas), + }, + _ => Self::Scalar { + scalar_fn, + children, + }, }) } } @@ -155,10 +215,12 @@ impl Expression { /// Computes the return dtype of this expression given the root dtype. /// /// Returns an error for variables because a root dtype alone cannot resolve lexical bindings; - /// use [`Expression::bind_scope`] when the expression may contain variables. + /// use [`Expression::bind_scope`] when the expression may contain variables. A higher-order + /// node is typed by binding it. pub fn return_dtype(&self, root_dtype: &DType) -> VortexResult { match self { Self::Root => Ok(root_dtype.clone()), + Self::HigherOrder { .. } => Ok(self.bind(root_dtype)?.dtype().clone()), Self::Variable(variable) => { vortex_bail!("cannot determine dtype of unbound variable '{variable}'") } @@ -183,7 +245,9 @@ impl Expression { match self { Self::Root => write!(f, "$"), Self::Variable(variable) => write!(f, "${variable}"), - Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + Self::Scalar { scalar_fn, .. } | Self::HigherOrder { scalar_fn, .. } => { + scalar_fn.fmt_sql(self, f) + } } } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 7311d51dc86..7da0d3ead4b 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; @@ -16,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; @@ -79,6 +81,16 @@ pub fn var(name: impl AsRef) -> Expression { Variable::new(name).into() } +/// Creates a lambda binding `params` over `body`, the argument of a higher-order function. +/// +/// Returns an error when a parameter name is repeated. +pub fn lambda( + params: impl IntoIterator>, + body: Expression, +) -> VortexResult { + Lambda::try_new(params, body) +} + /// 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/bound.rs b/vortex-array/src/expr/lambda/bound.rs new file mode 100644 index 00000000000..07264631d4d --- /dev/null +++ b/vortex-array/src/expr/lambda/bound.rs @@ -0,0 +1,182 @@ +// 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::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::ScalarFnArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::expr::BoundExpression; +use crate::expr::Variable; +use crate::expr::traversal::TraversalOrder; +use crate::expr::traversal::pre_order_visit_down; +use crate::optimizer::ArrayOptimizer; +use crate::scalar_fn::fns::literal::Literal; + +/// A lambda bound by its higher-order function: typed, with a body that reads only its inputs. +/// +/// The inputs are the parameters followed by the captures, the outer references that binding moved +/// out of the body. The body reads input `i` through a variable at frame `0`, slot `i`. The +/// function holds this in its options, takes the captures as extra children, and evaluates the body +/// with [`BoundLambda::evaluate`]. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct BoundLambda { + params: Arc<[(Variable, DType)]>, + captures: Arc<[DType]>, + body: Arc, +} + +impl BoundLambda { + pub(crate) fn new( + params: Vec<(Variable, DType)>, + captures: Vec, + body: BoundExpression, + ) -> Self { + Self { + params: params.into(), + captures: captures.into(), + body: Arc::new(body), + } + } + + /// The parameters with their dtypes, in declaration order. + pub fn params(&self) -> &[(Variable, DType)] { + &self.params + } + + /// The dtypes of the captures, the inputs after the parameters. + pub fn capture_dtypes(&self) -> &[DType] { + &self.captures + } + + /// The body. + pub fn body(&self) -> &BoundExpression { + &self.body + } + + /// The dtype the body evaluates to. + pub fn body_dtype(&self) -> &DType { + self.body.dtype() + } + + /// Whether every function in the body is infallible, so the body may run over values that no + /// row observes, such as the elements of a null list. + pub fn is_infallible(&self) -> bool { + let mut infallible = true; + pre_order_visit_down(self.body.as_ref(), |node| { + infallible &= node + .as_scalar() + .is_none_or(|scalar_fn| scalar_fn.signature().is_infallible()); + Ok(if infallible { + TraversalOrder::Continue + } else { + TraversalOrder::Stop + }) + }) + .vortex_expect("bound expression traversal cannot fail"); + infallible + } + + /// Check that inputs of `dtypes`, in order, fit this lambda. + /// + /// An input of a nullable dtype also accepts its non-nullable form, because rewrites after + /// binding can narrow a child's nullability, such as a layout reader stepping into the fields of + /// a nullable struct. [`BoundLambda::evaluate`] widens such inputs back. + pub fn check_inputs<'a>( + &self, + dtypes: impl IntoIterator, + ) -> VortexResult<()> { + let dtypes = dtypes.into_iter().collect_vec(); + vortex_ensure!( + dtypes.len() == self.params.len() + self.captures.len(), + "lambda {self} takes {} input(s) but was given {}", + self.params.len() + self.captures.len(), + dtypes.len() + ); + for (actual, expected) in dtypes.into_iter().zip(self.input_dtypes()) { + vortex_ensure!( + actual == expected || actual.as_nullable() == *expected, + "lambda {self} input has dtype {actual}, expected {expected}" + ); + } + Ok(()) + } + + /// Evaluate the body over `inputs`, arrays of equal length, into a lazy array of that length. + pub fn evaluate(&self, inputs: &[ArrayRef]) -> VortexResult { + self.check_inputs(inputs.iter().map(ArrayRef::dtype))?; + let Some(len) = inputs.first().map(ArrayRef::len) else { + vortex_bail!("lambda {self} has no inputs to evaluate over"); + }; + let inputs: Vec = inputs + .iter() + .zip(self.input_dtypes()) + .map(|(input, dtype)| { + vortex_ensure!( + input.len() == len, + "lambda {self} inputs have lengths {} and {len}", + input.len() + ); + if input.dtype() == dtype { + Ok(input.clone()) + } else { + input.cast(dtype.clone()) + } + }) + .try_collect()?; + eval(&self.body, &inputs, len) + } + + fn input_dtypes(&self) -> impl Iterator { + self.params + .iter() + .map(|(_, dtype)| dtype) + .chain(self.captures.iter()) + } +} + +/// Evaluate a bound lambda body, whose variables read `inputs`, into a lazy array of length `len`. +fn eval(node: &BoundExpression, inputs: &[ArrayRef], len: usize) -> VortexResult { + match node { + BoundExpression::Variable(variable) => match inputs.get(variable.variable_ref().slot()) { + Some(input) => Ok(input.clone()), + None => vortex_bail!("lambda body reads input {variable}, which it was not given"), + }, + BoundExpression::Root { .. } => vortex_bail!("a bound lambda body cannot read root()"), + BoundExpression::Scalar { + scalar_fn, + children, + .. + } => { + if let Some(scalar) = scalar_fn.as_opt::() { + return Ok(ConstantArray::new(scalar.clone(), len).into_array()); + } + let children: Vec<_> = children + .iter() + .map(|child| eval(child, inputs, len)) + .try_collect()?; + ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)? + .into_array() + .optimize() + } + } +} + +impl Display for BoundLambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let params = self.params.iter().map(|(name, _)| name).join(", "); + write!(f, "({params}) -> {}", self.body) + } +} diff --git a/vortex-array/src/expr/lambda/mod.rs b/vortex-array/src/expr/lambda/mod.rs new file mode 100644 index 00000000000..2738ae39fb2 --- /dev/null +++ b/vortex-array/src/expr/lambda/mod.rs @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lambdas, the function arguments of higher-order scalar functions. +//! +//! A [`Lambda`] is syntax and only appears in an [`Expression::HigherOrder`] node. Binding the node +//! turns each lambda into a [`BoundLambda`]: the body is typed with the parameter dtypes the +//! function supplies, and every outer reference in it, such as `$.a`, becomes an extra input that +//! the node takes as a child. The bound body reads only its inputs, so the function can evaluate it +//! over plain arrays. + +mod bound; +#[cfg(test)] +mod tests; + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::sync::Arc; + +pub use bound::BoundLambda; +use itertools::Itertools; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::dtype::DType; +use crate::expr::BoundExpression; +use crate::expr::BoundVariable; +use crate::expr::Expression; +use crate::expr::Frame; +use crate::expr::Scope; +use crate::expr::Variable; +use crate::expr::VariableRef; +use crate::scalar_fn::fns::get_item::GetItem; + +/// Parameter names over a body: a lambda argument of an [`Expression::HigherOrder`] node. +/// +/// Inside the body, `root()` is the same row as outside the lambda, and a parameter is read through +/// its [`Variable`]. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Lambda { + params: Arc<[Variable]>, + 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 params: Vec = params.into_iter().map(Into::into).collect(); + if let Some(param) = params.iter().duplicates().next() { + vortex_bail!("duplicate lambda parameter '{param}'"); + } + Ok(Self { + params: params.into(), + body: Arc::new(body), + }) + } + + /// The parameters, in declaration order. + pub fn params(&self) -> &[Variable] { + &self.params + } + + /// The body. + pub fn body(&self) -> &Expression { + &self.body + } + + /// Bind this lambda in `scope`, with the parameter dtypes its function supplies. + /// + /// Returns the bound lambda and its captures: the outer references of the body, bound in + /// `scope`, which the function takes as trailing children. + pub fn bind( + &self, + scope: &Scope, + param_dtypes: Vec, + ) -> VortexResult<(BoundLambda, Vec)> { + vortex_ensure!( + param_dtypes.len() == self.params.len(), + "lambda {self} declares {} parameter(s) but was given {} dtype(s)", + self.params.len(), + param_dtypes.len() + ); + let params = self.params.iter().cloned().zip(param_dtypes).collect_vec(); + let body = self + .body + .bind_scope(&scope.push_frame(Frame::try_new(params.clone())?))?; + + let mut inputs = Inputs { + param_frame: scope.depth(), + param_count: params.len(), + captures: Vec::new(), + }; + let body = inputs.rewrite(&body)?; + let capture_dtypes = inputs + .captures + .iter() + .map(|capture| capture.dtype().clone()) + .collect(); + Ok(( + BoundLambda::new(params, capture_dtypes, body), + inputs.captures, + )) + } +} + +impl Display for Lambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "({}) -> {}", self.params.iter().join(", "), self.body) + } +} + +/// Rewrites a bound lambda body to read only its inputs, the parameters followed by the captures. +struct Inputs { + /// The scope frame holding the lambda's parameters. + param_frame: usize, + param_count: usize, + /// The outer references moved out of the body, in input order. + captures: Vec, +} + +impl Inputs { + fn rewrite(&mut self, node: &BoundExpression) -> VortexResult { + if let Some(variable) = node.as_variable() + && variable.variable_ref().frame() == self.param_frame + { + return Ok(input( + variable.dtype().clone(), + variable.variable().clone(), + variable.variable_ref().slot(), + )); + } + if is_outer_reference(node) { + let index = match self.captures.iter().position(|capture| capture == node) { + Some(index) => index, + None => { + self.captures.push(node.clone()); + self.captures.len() - 1 + } + }; + return Ok(input( + node.dtype().clone(), + Variable::new(format!("capture{index}")), + self.param_count + index, + )); + } + let children: Vec<_> = node + .children() + .iter() + .map(|child| self.rewrite(child)) + .try_collect()?; + node.clone().with_children(children) + } +} + +/// Whether `node` reads only the enclosing scope: `root()`, a variable of an enclosing lambda, or a +/// field path into one of them. +/// +/// Parameters of the lambda being bound are handled before this is asked. +fn is_outer_reference(node: &BoundExpression) -> bool { + match node { + BoundExpression::Root { .. } | BoundExpression::Variable(_) => true, + BoundExpression::Scalar { .. } => node.is::() && is_outer_reference(node.child(0)), + } +} + +/// Input `slot` of a bound lambda body. +fn input(dtype: DType, name: Variable, slot: usize) -> BoundExpression { + BoundExpression::Variable(BoundVariable::new(dtype, name, VariableRef::new(0, slot))) +} diff --git a/vortex-array/src/expr/lambda/tests.rs b/vortex-array/src/expr/lambda/tests.rs new file mode 100644 index 00000000000..f273aa63b9b --- /dev/null +++ b/vortex-array/src/expr/lambda/tests.rs @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::StructFields; +use crate::expr::Expression; +use crate::expr::Scope; +use crate::expr::Variable; +use crate::expr::checked_add; +use crate::expr::col; +use crate::expr::eq; +use crate::expr::lambda; +use crate::expr::lit; +use crate::expr::not; +use crate::expr::var; + +fn i32_() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) +} + +/// The row `{a: i32, b: i32}`. +fn row() -> DType { + DType::Struct( + StructFields::from_iter([("a", i32_()), ("b", i32_())]), + Nullability::NonNullable, + ) +} + +#[test] +fn duplicate_parameters_are_rejected() { + assert!(lambda(["x", "x"], var("x")).is_err()); +} + +#[test] +fn outer_references_become_shared_captures() -> VortexResult<()> { + let scope = Scope::new(row()); + let lambda = lambda( + ["x"], + checked_add(checked_add(var("x"), col("a")), col("a")), + )?; + let (bound, captures) = lambda.bind(&scope, vec![i32_()])?; + + assert_eq!(captures, [col("a").bind(&row())?]); + assert_eq!(bound.capture_dtypes(), [i32_()]); + assert_eq!(bound.to_string(), "(x) -> (($x + $capture0) + $capture0)"); + Ok(()) +} + +#[test] +fn parameters_and_literals_stay_in_the_body() -> VortexResult<()> { + let (bound, captures) = + lambda(["x"], checked_add(var("x"), lit(1_i32)))?.bind(&Scope::new(row()), vec![i32_()])?; + assert!(captures.is_empty()); + assert_eq!(bound.to_string(), "(x) -> ($x + 1i32)"); + Ok(()) +} + +#[test] +fn variables_of_an_enclosing_lambda_are_captured() -> VortexResult<()> { + let scope = Scope::new(row()).with_bindings([(Variable::new("y"), i32_())])?; + let (bound, captures) = + lambda(["x"], checked_add(var("x"), var("y")))?.bind(&scope, vec![i32_()])?; + assert_eq!(captures, [var("y").bind_scope(&scope)?]); + assert_eq!(bound.to_string(), "(x) -> ($x + $capture0)"); + Ok(()) +} + +#[test] +fn bind_checks_parameter_count_and_dtypes() -> VortexResult<()> { + let scope = Scope::new(row()); + assert!( + lambda(["x", "y"], var("x"))? + .bind(&scope, vec![i32_()]) + .is_err() + ); + assert!( + lambda(["x"], eq(var("x"), lit("nope")))? + .bind(&scope, vec![i32_()]) + .is_err() + ); + Ok(()) +} + +#[test] +fn evaluate_runs_the_body_over_its_inputs() -> VortexResult<()> { + let (bound, _) = + lambda(["x"], checked_add(var("x"), col("a")))?.bind(&Scope::new(row()), vec![i32_()])?; + let x = buffer![1_i32, 2].into_array(); + let a = buffer![10_i32, 20].into_array(); + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + bound.evaluate(&[x.clone(), a])?, + buffer![11_i32, 22].into_array(), + &mut ctx + ); + + assert!(bound.evaluate(&[]).is_err()); + assert!(bound.evaluate(std::slice::from_ref(&x)).is_err()); + assert!( + bound + .evaluate(&[x.clone(), buffer![1_i64, 2].into_array()]) + .is_err() + ); + assert!(bound.evaluate(&[x, buffer![10_i32].into_array()]).is_err()); + Ok(()) +} + +#[test] +fn a_nullable_input_accepts_its_non_nullable_form() -> VortexResult<()> { + // In a nullable row `$.a` is nullable, but a layout reader stepping into the row's fields + // hands over the field's own non-nullable array. + let (bound, _) = lambda(["x"], checked_add(var("x"), col("a")))? + .bind(&Scope::new(row().as_nullable()), vec![i32_()])?; + assert_eq!(bound.capture_dtypes(), [i32_().as_nullable()]); + + let x = buffer![1_i32, 2].into_array(); + let a = buffer![10_i32, 20].into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + bound.evaluate(&[x, a.clone()])?, + PrimitiveArray::from_option_iter([Some(11_i32), Some(22)]).into_array(), + &mut ctx + ); + + // The reverse would let nulls into an input typed non-nullable. + let nullable_x = PrimitiveArray::from_option_iter([Some(1_i32), None]).into_array(); + assert!(bound.evaluate(&[nullable_x, a]).is_err()); + Ok(()) +} + +#[test] +fn infallibility_comes_from_the_body() -> VortexResult<()> { + let scope = Scope::new(row()); + let (bound, _) = + lambda(["x", "i"], eq(var("x"), lit(1_i32)))?.bind(&scope, vec![i32_(), i32_()])?; + assert!(bound.is_infallible()); + + let (bound, _) = + lambda(["x"], checked_add(var("x"), lit(1_i32)))?.bind(&scope, vec![i32_()])?; + assert!(!bound.is_infallible()); + Ok(()) +} + +#[test] +fn ordinary_functions_cannot_take_lambdas() -> VortexResult<()> { + let not_fn = not(lit(true)) + .as_scalar() + .cloned() + .ok_or_else(|| vortex_error::vortex_err!("not is a scalar function"))?; + assert!( + Expression::try_new_higher_order(not_fn, [lit(true)], [lambda(["x"], var("x"))?]).is_err() + ); + Ok(()) +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 827179c108d..a3074af5299 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -49,6 +49,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; @@ -85,6 +86,7 @@ pub use exprs::is_nan; 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_contains_opts; @@ -112,6 +114,7 @@ pub use exprs::select_exclude; pub use exprs::var; pub use exprs::variant_get; pub use exprs::zip_expr; +pub use lambda::*; pub use scope::*; pub use variable::*; diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 951751e1ac0..a4d489ef96c 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use prost::Message; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -88,6 +89,12 @@ impl ExprSerializeProtoExt for Expression { metadata: Some(vec![]), }), Expression::Variable(variable) => Ok(variable.serialize_proto()), + Expression::HigherOrder { scalar_fn, .. } => { + vortex_bail!( + "higher-order expression '{}' is not serializable", + scalar_fn.id() + ) + } Expression::Scalar { scalar_fn, children, diff --git a/vortex-array/src/expr/scope.rs b/vortex-array/src/expr/scope.rs index 3c71401d69a..d55cd4bc223 100644 --- a/vortex-array/src/expr/scope.rs +++ b/vortex-array/src/expr/scope.rs @@ -20,7 +20,7 @@ pub struct VariableRef { } impl VariableRef { - fn new(frame: usize, slot: usize) -> Self { + pub(crate) fn new(frame: usize, slot: usize) -> Self { Self { frame, slot } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index 1387fe01e93..917a25db038 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -56,6 +56,12 @@ impl ArrayRef { vortex_bail!("cannot apply unbound variable '{variable}'"); } + // A higher-order node needs the dtypes of its arguments to bind its lambdas. + if !expr.lambdas().is_empty() { + let bound = expr.bind(self.dtype())?; + return self.apply_bound(&bound); + } + // Manually convert literals to ConstantArray. if let Some(scalar) = expr.as_opt::() { return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array()); @@ -71,7 +77,7 @@ impl ArrayRef { // And wrap the scalar function up in an array. let scalar_fn = expr.as_scalar().vortex_expect( - "root, variable and literal were handled above, so this is a scalar node", + "root, variable, higher-order and literal were handled above, so this is a scalar node", ); let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array(); diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index d1b09ce5f1f..3b7ca2ea88a 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -20,6 +20,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::BoundExpression; +use crate::expr::BoundLambda; use crate::expr::display::ExprDisplay; use crate::scalar_fn::ArrayReduceNode; use crate::scalar_fn::ExecutionArgs; @@ -163,6 +164,29 @@ impl ScalarFnRef { self.0.reduce_array(node) } + /// Whether this function takes lambda arguments. + /// See [`ScalarFnVTable::higher_order`]. + pub fn is_higher_order(&self) -> bool { + self.0.is_higher_order() + } + + /// The parameter dtypes of a lambda argument. + /// See [`HigherOrderVTable::lambda_params`](crate::scalar_fn::HigherOrderVTable::lambda_params). + pub fn lambda_params( + &self, + lambda_index: usize, + param_count: usize, + arg_dtypes: &[DType], + ) -> VortexResult> { + self.0.lambda_params(lambda_index, param_count, arg_dtypes) + } + + /// This scalar function with its options holding the given bound lambdas. + /// See [`HigherOrderVTable::with_lambdas`](crate::scalar_fn::HigherOrderVTable::with_lambdas). + pub fn with_lambdas(&self, lambdas: Vec) -> VortexResult { + self.0.with_lambdas(lambdas) + } + // ------------------------------------------------------------------ // Expression-taking methods — used by expr/ module via pub(crate) // ------------------------------------------------------------------ diff --git a/vortex-array/src/scalar_fn/higher_order.rs b/vortex-array/src/scalar_fn/higher_order.rs new file mode 100644 index 00000000000..4e21eb52103 --- /dev/null +++ b/vortex-array/src/scalar_fn/higher_order.rs @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::dtype::DType; +use crate::expr::BoundLambda; +use crate::scalar_fn::ScalarFnVTable; + +/// The lambda hooks of a scalar function that takes lambda arguments. +/// +/// A higher-order function implements this alongside [`ScalarFnVTable`] and returns itself from +/// [`ScalarFnVTable::higher_order`]. Binding an +/// [`Expression::HigherOrder`](crate::expr::Expression::HigherOrder) node asks the function for each +/// lambda's parameter dtypes, binds the lambdas, and stores them in the function's options. The +/// bound node is an ordinary scalar node whose children are the value arguments followed by every +/// lambda's captures. +pub trait HigherOrderVTable { + /// The parameter dtypes of the `lambda_index`th lambda, given the dtypes of the value + /// arguments. + /// + /// `param_count` is the number of parameters the lambda declares, so a function can offer an + /// optional parameter. + fn lambda_params( + &self, + options: &V::Options, + lambda_index: usize, + param_count: usize, + arg_dtypes: &[DType], + ) -> VortexResult>; + + /// These options holding the bound lambdas, in argument order. + fn with_lambdas( + &self, + options: &V::Options, + lambdas: Vec, + ) -> VortexResult; +} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 6be34ce1f34..6e2b069e46d 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -29,6 +29,9 @@ pub use plugin::*; mod foreign; pub use foreign::*; +mod higher_order; +pub use higher_order::*; + mod typed; pub use typed::*; diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index aa5ba2d36b6..dd94ff8bec1 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -19,11 +19,13 @@ use std::hash::Hasher; use std::sync::Arc; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::BoundExpression; +use crate::expr::BoundLambda; use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ArrayReduceNode; @@ -94,6 +96,16 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { fn is_strict(&self) -> bool; fn is_infallible(&self) -> bool; + // Lambda arguments — options in self, see `ScalarFnVTable::higher_order` + fn is_higher_order(&self) -> bool; + fn lambda_params( + &self, + lambda_index: usize, + param_count: usize, + arg_dtypes: &[DType], + ) -> VortexResult>; + fn with_lambdas(&self, lambdas: Vec) -> VortexResult; + // Expression methods — take expressions for tree traversal fn fmt_sql(&self, expression: &dyn ExprDisplay, f: &mut Formatter<'_>) -> fmt::Result; fn simplify(&self, expression: &BoundExpression) -> VortexResult>; @@ -199,6 +211,30 @@ impl DynScalarFn for TypedScalarFnInstance { V::is_infallible(&self.vtable, &self.options) } + fn is_higher_order(&self) -> bool { + V::higher_order(&self.vtable).is_some() + } + + fn lambda_params( + &self, + lambda_index: usize, + param_count: usize, + arg_dtypes: &[DType], + ) -> VortexResult> { + let Some(higher_order) = V::higher_order(&self.vtable) else { + vortex_bail!("{} does not take lambda arguments", V::id(&self.vtable)); + }; + higher_order.lambda_params(&self.options, lambda_index, param_count, arg_dtypes) + } + + fn with_lambdas(&self, lambdas: Vec) -> VortexResult { + let Some(higher_order) = V::higher_order(&self.vtable) else { + vortex_bail!("{} does not take lambda arguments", V::id(&self.vtable)); + }; + let options = higher_order.with_lambdas(&self.options, lambdas)?; + Ok(TypedScalarFnInstance::new(self.vtable.clone(), options).erased()) + } + fn fmt_sql(&self, expression: &dyn ExprDisplay, f: &mut Formatter<'_>) -> fmt::Result { V::fmt_sql(&self.vtable, &self.options, expression, f) } diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index cf13ba98bcf..04669a80f17 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -28,6 +28,7 @@ use crate::expr::Expression; use crate::expr::bound; use crate::expr::display::ExprDisplay; use crate::scalar::Scalar; +use crate::scalar_fn::HigherOrderVTable; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::TypedScalarFnInstance; @@ -277,6 +278,15 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { _ = options; false } + + /// The lambda hooks of this function, if it takes lambda arguments. + /// + /// A function that returns `Some` can be applied to lambdas in an + /// [`Expression::HigherOrder`](crate::expr::Expression::HigherOrder) node. The default is + /// `None`, so an ordinary function rejects lambdas when such a node is built. + fn higher_order(&self) -> Option<&dyn HigherOrderVTable> { + None + } } /// A node used for implementing abstract reduction rules over a tree of scalar functions. From 676b8877a3fb49d93ca77e6f7b7b4d22693b3346 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 12 Oct 2026 16:11:25 -0400 Subject: [PATCH 2/2] lambda hooks as defaulted ScalarFnVTable methods Signed-off-by: Matt Katz --- vortex-array/src/expr/expression.rs | 10 ++---- vortex-array/src/expr/lambda/tests.rs | 7 +++- vortex-array/src/scalar_fn/erased.rs | 10 ++---- vortex-array/src/scalar_fn/higher_order.rs | 38 ---------------------- vortex-array/src/scalar_fn/mod.rs | 3 -- vortex-array/src/scalar_fn/typed.rs | 24 +++++--------- vortex-array/src/scalar_fn/vtable.rs | 36 ++++++++++++++++---- 7 files changed, 49 insertions(+), 79 deletions(-) delete mode 100644 vortex-array/src/scalar_fn/higher_order.rs diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index 90f6b9c2cb0..d611a1a7021 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -49,8 +49,9 @@ pub enum Expression { Variable(Variable), /// A scalar function applied to child expressions and then to lambdas. /// - /// The function implements [`HigherOrderVTable`](crate::scalar_fn::HigherOrderVTable). Binding - /// produces an ordinary scalar node whose options hold the bound lambdas. + /// Binding asks the function for each lambda's parameter dtypes (see + /// [`ScalarFnVTable::lambda_params`]) and produces an ordinary scalar node whose options hold + /// the bound lambdas. HigherOrder { /// The scalar fn for this node. scalar_fn: ScalarFnRef, @@ -89,11 +90,6 @@ impl Expression { lambdas: impl IntoIterator, ) -> VortexResult { let children = Vec::from_iter(children); - vortex_ensure!( - scalar_fn.is_higher_order(), - "{} does not take lambda arguments", - scalar_fn.id() - ); vortex_ensure!( scalar_fn.signature().arity().matches(children.len()), "Expression arity mismatch: expected {} children but got {}", diff --git a/vortex-array/src/expr/lambda/tests.rs b/vortex-array/src/expr/lambda/tests.rs index f273aa63b9b..1213f700695 100644 --- a/vortex-array/src/expr/lambda/tests.rs +++ b/vortex-array/src/expr/lambda/tests.rs @@ -158,8 +158,13 @@ fn ordinary_functions_cannot_take_lambdas() -> VortexResult<()> { .as_scalar() .cloned() .ok_or_else(|| vortex_error::vortex_err!("not is a scalar function"))?; + let expr = Expression::try_new_higher_order(not_fn, [lit(true)], [lambda(["x"], var("x"))?])?; + let error = expr + .bind(&row()) + .expect_err("a function without lambda hooks must reject lambdas"); assert!( - Expression::try_new_higher_order(not_fn, [lit(true)], [lambda(["x"], var("x"))?]).is_err() + error.to_string().contains("does not take lambda arguments"), + "{error}" ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index 3b7ca2ea88a..4cec11c1eaa 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -164,14 +164,8 @@ impl ScalarFnRef { self.0.reduce_array(node) } - /// Whether this function takes lambda arguments. - /// See [`ScalarFnVTable::higher_order`]. - pub fn is_higher_order(&self) -> bool { - self.0.is_higher_order() - } - /// The parameter dtypes of a lambda argument. - /// See [`HigherOrderVTable::lambda_params`](crate::scalar_fn::HigherOrderVTable::lambda_params). + /// See [`ScalarFnVTable::lambda_params`]. pub fn lambda_params( &self, lambda_index: usize, @@ -182,7 +176,7 @@ impl ScalarFnRef { } /// This scalar function with its options holding the given bound lambdas. - /// See [`HigherOrderVTable::with_lambdas`](crate::scalar_fn::HigherOrderVTable::with_lambdas). + /// See [`ScalarFnVTable::with_lambdas`]. pub fn with_lambdas(&self, lambdas: Vec) -> VortexResult { self.0.with_lambdas(lambdas) } diff --git a/vortex-array/src/scalar_fn/higher_order.rs b/vortex-array/src/scalar_fn/higher_order.rs deleted file mode 100644 index 4e21eb52103..00000000000 --- a/vortex-array/src/scalar_fn/higher_order.rs +++ /dev/null @@ -1,38 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_error::VortexResult; - -use crate::dtype::DType; -use crate::expr::BoundLambda; -use crate::scalar_fn::ScalarFnVTable; - -/// The lambda hooks of a scalar function that takes lambda arguments. -/// -/// A higher-order function implements this alongside [`ScalarFnVTable`] and returns itself from -/// [`ScalarFnVTable::higher_order`]. Binding an -/// [`Expression::HigherOrder`](crate::expr::Expression::HigherOrder) node asks the function for each -/// lambda's parameter dtypes, binds the lambdas, and stores them in the function's options. The -/// bound node is an ordinary scalar node whose children are the value arguments followed by every -/// lambda's captures. -pub trait HigherOrderVTable { - /// The parameter dtypes of the `lambda_index`th lambda, given the dtypes of the value - /// arguments. - /// - /// `param_count` is the number of parameters the lambda declares, so a function can offer an - /// optional parameter. - fn lambda_params( - &self, - options: &V::Options, - lambda_index: usize, - param_count: usize, - arg_dtypes: &[DType], - ) -> VortexResult>; - - /// These options holding the bound lambdas, in argument order. - fn with_lambdas( - &self, - options: &V::Options, - lambdas: Vec, - ) -> VortexResult; -} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 6e2b069e46d..6be34ce1f34 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -29,9 +29,6 @@ pub use plugin::*; mod foreign; pub use foreign::*; -mod higher_order; -pub use higher_order::*; - mod typed; pub use typed::*; diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index dd94ff8bec1..c2da218c5e9 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -19,7 +19,6 @@ use std::hash::Hasher; use std::sync::Arc; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use crate::ArrayRef; use crate::ExecutionCtx; @@ -96,8 +95,7 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { fn is_strict(&self) -> bool; fn is_infallible(&self) -> bool; - // Lambda arguments — options in self, see `ScalarFnVTable::higher_order` - fn is_higher_order(&self) -> bool; + // Lambda arguments — options in self, see `ScalarFnVTable::lambda_params` fn lambda_params( &self, lambda_index: usize, @@ -211,27 +209,23 @@ impl DynScalarFn for TypedScalarFnInstance { V::is_infallible(&self.vtable, &self.options) } - fn is_higher_order(&self) -> bool { - V::higher_order(&self.vtable).is_some() - } - fn lambda_params( &self, lambda_index: usize, param_count: usize, arg_dtypes: &[DType], ) -> VortexResult> { - let Some(higher_order) = V::higher_order(&self.vtable) else { - vortex_bail!("{} does not take lambda arguments", V::id(&self.vtable)); - }; - higher_order.lambda_params(&self.options, lambda_index, param_count, arg_dtypes) + V::lambda_params( + &self.vtable, + &self.options, + lambda_index, + param_count, + arg_dtypes, + ) } fn with_lambdas(&self, lambdas: Vec) -> VortexResult { - let Some(higher_order) = V::higher_order(&self.vtable) else { - vortex_bail!("{} does not take lambda arguments", V::id(&self.vtable)); - }; - let options = higher_order.with_lambdas(&self.options, lambdas)?; + let options = V::with_lambdas(&self.vtable, &self.options, lambdas)?; Ok(TypedScalarFnInstance::new(self.vtable.clone(), options).erased()) } diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 04669a80f17..78dbe4216fb 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -24,11 +24,11 @@ use crate::arrays::ScalarFn; use crate::arrays::ScalarFnArray; use crate::dtype::DType; use crate::expr::BoundExpression; +use crate::expr::BoundLambda; use crate::expr::Expression; use crate::expr::bound; use crate::expr::display::ExprDisplay; use crate::scalar::Scalar; -use crate::scalar_fn::HigherOrderVTable; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::TypedScalarFnInstance; @@ -279,13 +279,35 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { false } - /// The lambda hooks of this function, if it takes lambda arguments. + /// The parameter dtypes of the `lambda_index`th lambda, given the dtypes of the value + /// arguments, for a function applied to lambdas in an + /// [`Expression::HigherOrder`](crate::expr::Expression::HigherOrder) node. /// - /// A function that returns `Some` can be applied to lambdas in an - /// [`Expression::HigherOrder`](crate::expr::Expression::HigherOrder) node. The default is - /// `None`, so an ordinary function rejects lambdas when such a node is built. - fn higher_order(&self) -> Option<&dyn HigherOrderVTable> { - None + /// `param_count` is the number of parameters the lambda declares, so a function can offer an + /// optional parameter. The default rejects lambdas. + fn lambda_params( + &self, + options: &Self::Options, + lambda_index: usize, + param_count: usize, + arg_dtypes: &[DType], + ) -> VortexResult> { + _ = (options, lambda_index, param_count, arg_dtypes); + vortex_bail!("{} does not take lambda arguments", self.id()) + } + + /// These options holding the bound lambdas, in argument order. + /// + /// Binding a higher-order node stores its lambdas here. The bound node's children are the value + /// arguments followed by every lambda's captures. The default rejects lambdas. + fn with_lambdas( + &self, + options: &Self::Options, + lambdas: Vec, + ) -> VortexResult { + _ = options; + drop(lambdas); + vortex_bail!("{} does not take lambda arguments", self.id()) } }