From aafd2d308b448f948f33460db6512340e2b63518 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 14 Oct 2026 16:47:18 -0400 Subject: [PATCH 1/2] lambdas for higher-order scalar functions Signed-off-by: Matt Katz --- vortex-array/src/arrays/scalar_fn/array.rs | 1 + vortex-array/src/expr/bound_expression.rs | 7 +- vortex-array/src/expr/display.rs | 2 + vortex-array/src/expr/expression.rs | 37 ++- vortex-array/src/expr/exprs.rs | 19 ++ vortex-array/src/expr/lambda/mod.rs | 291 +++++++++++++++++ vortex-array/src/expr/lambda/tests.rs | 343 +++++++++++++++++++++ vortex-array/src/expr/mod.rs | 3 + vortex-array/src/expr/proto.rs | 4 + vortex-array/src/expression.rs | 6 +- vortex-array/src/scalar_fn/erased.rs | 6 + vortex-array/src/scalar_fn/typed.rs | 6 + vortex-array/src/scalar_fn/vtable.rs | 18 ++ vortex-layout/src/layouts/row_idx/expr.rs | 23 ++ 14 files changed, 758 insertions(+), 8 deletions(-) create mode 100644 vortex-array/src/expr/lambda/mod.rs create mode 100644 vortex-array/src/expr/lambda/tests.rs diff --git a/vortex-array/src/arrays/scalar_fn/array.rs b/vortex-array/src/arrays/scalar_fn/array.rs index 1853cfca53e..391ea9a03cd 100644 --- a/vortex-array/src/arrays/scalar_fn/array.rs +++ b/vortex-array/src/arrays/scalar_fn/array.rs @@ -110,6 +110,7 @@ impl Array { Self::validate_children_len(&slots, len)?; let arg_dtypes = Self::arg_dtypes(&slots); + let scalar_fn = scalar_fn.bind_options(&arg_dtypes)?; let dtype = scalar_fn.return_dtype(&arg_dtypes)?; Ok(unsafe { Self::new_from_validated_slots(scalar_fn, dtype, slots, len) }) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 97666966b62..8c3ebb8eec0 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::display::DisplayTreeExpr; +use crate::expr::expression::unbound_variable; use crate::expr::scope::Scope; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; @@ -138,6 +139,7 @@ impl BoundExpression { .iter() .map(|child| child.dtype().clone()) .collect_vec(); + let scalar_fn = scalar_fn.bind_options(&arg_dtypes)?; let dtype = scalar_fn.return_dtype(&arg_dtypes)?; Ok(Self::Scalar { @@ -310,6 +312,9 @@ impl Expression { if self.is_root() { return Ok(BoundExpression::new_root(scope.root().clone())); } + if let Expression::Variable(name) = self { + return Err(unbound_variable(name)); + } let children: Vec<_> = self .children() @@ -318,7 +323,7 @@ impl Expression { .try_collect()?; let scalar_fn = self .as_scalar() - .vortex_expect("root was handled above, so this is a scalar node"); + .vortex_expect("leaves were handled above, so this is a scalar node"); BoundExpression::try_new(scalar_fn.clone(), children) } } diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index b22ae246522..8de2a7d112d 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -69,6 +69,7 @@ impl DisplayTreeNode for Expression { 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"), } } @@ -76,6 +77,7 @@ impl DisplayTreeNode for Expression { match self { Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), Expression::Root => write!(f, "{ROOT_DISPLAY}"), + Expression::Variable(name) => write!(f, "@{name}"), } } } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index d228f646363..ded8e2c87d7 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -10,11 +10,14 @@ use std::hash::Hash; use std::sync::Arc; use itertools::Itertools; +use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::dtype::DType; +use crate::dtype::FieldName; use crate::expr::display::DisplayTreeExpr; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; @@ -26,10 +29,12 @@ const NO_CHILDREN: &[Expression] = &[]; /// A node in a Vortex expression tree. /// -/// 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. A [`ScalarFnVTable`] can answer neither -/// of those, so `Root` is a variant instead. +/// Most nodes are a scalar function applied to child expressions. The other two are leaves that +/// name a value rather than compute one: [`Expression::Root`], the scope itself, and +/// [`Expression::Variable`], a parameter of an enclosing lambda. They are language primitives +/// rather than registered functions, because their dtype comes from where they appear rather than +/// from children and they are not executable. A [`ScalarFnVTable`] can answer neither of those, so +/// they are variants instead. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Expression { /// A scalar function applied to child expressions. @@ -41,6 +46,11 @@ pub enum Expression { }, /// The full scope of the expression evaluation. Root, + /// A parameter of the enclosing lambda, read by name. + /// + /// Binding a [`Lambda`](crate::expr::Lambda) turns its variables into reads of the struct of + /// its inputs. Anywhere else a variable fails to bind, and it is not serializable yet. + Variable(FieldName), } impl Expression { @@ -73,7 +83,7 @@ impl Expression { pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root => None, + Self::Root | Self::Variable(_) => None, } } @@ -101,7 +111,7 @@ impl Expression { pub fn children(&self) -> &[Expression] { match self { Self::Scalar { children, .. } => children.as_slice(), - Self::Root => NO_CHILDREN, + Self::Root | Self::Variable(_) => NO_CHILDREN, } } @@ -125,6 +135,14 @@ impl Expression { ); Ok(Self::Root) } + Self::Variable(name) => { + vortex_ensure!( + children.is_empty(), + "Expression arity mismatch: variable expects 0 children but got {}", + children.len() + ); + Ok(Self::Variable(name.clone())) + } Self::Scalar { scalar_fn, .. } => { vortex_ensure!( scalar_fn.signature().arity().matches(children.len()), @@ -144,6 +162,7 @@ impl Expression { pub fn return_dtype(&self, scope: &DType) -> VortexResult { match self { Self::Root => Ok(scope.clone()), + Self::Variable(name) => Err(unbound_variable(name)), Self::Scalar { scalar_fn, children, @@ -164,6 +183,7 @@ impl Expression { pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Self::Root => write!(f, "$"), + Self::Variable(name) => write!(f, "@{name}"), Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), } } @@ -249,6 +269,11 @@ impl Expression { } } +/// The error for a variable that no enclosing lambda declares. +pub(crate) fn unbound_variable(name: &FieldName) -> VortexError { + vortex_err!("variable @{name} is not a parameter of an enclosing lambda") +} + /// The default display implementation for expressions uses the 'SQL'-style format. impl Display for Expression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index b0bed8565f6..d77a3472d30 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -149,6 +149,25 @@ pub fn bound_get_item(field: impl Into, child: BoundExpression) -> Bo .vortex_expect("get-item expressions must reference a field in the child dtype") } +// ---- Variable ---- + +/// Creates an expression that reads the parameter `name` of the enclosing lambda. +/// +/// It is only meaningful in the body of a [`Lambda`](crate::expr::Lambda). Anywhere else it fails +/// to bind. +/// +/// ```rust +/// # use vortex_array::dtype::{DType, Nullability, PType}; +/// # use vortex_array::expr::{checked_add, lit, var}; +/// let expr = checked_add(var("x"), lit(1_i32)); +/// assert_eq!(expr.to_string(), "(@x + 1i32)"); +/// // Outside a lambda body, nothing gives `x` a value. +/// assert!(expr.bind(&DType::Primitive(PType::I32, Nullability::NonNullable)).is_err()); +/// ``` +pub fn var(name: impl Into) -> Expression { + Expression::Variable(name.into()) +} + // ---- VariantGet ---- /// Creates an expression that extracts a path from a Variant expression. diff --git a/vortex-array/src/expr/lambda/mod.rs b/vortex-array/src/expr/lambda/mod.rs new file mode 100644 index 00000000000..ea810dc428e --- /dev/null +++ b/vortex-array/src/expr/lambda/mod.rs @@ -0,0 +1,291 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lambdas, the function arguments of higher-order scalar functions such as `list_transform`. +//! +//! A [`Lambda`] is a list of parameter names and a body that reads them with +//! [`var`](crate::expr::var). A function that takes a lambda keeps it in its options. When the +//! function's node is bound, it binds the lambda to its parameters' dtypes, once, and keeps the +//! [`BoundLambda`] instead. +//! +//! Binding turns the lambda into an ordinary expression over a struct of its inputs, one field per +//! parameter: `@x` becomes `$.x`. To run it, the function puts the parameters' arrays in that +//! struct, such as one row per list element, and applies the bound body to it. The body reads +//! nothing but its parameters, so the function's children are everything it reads. + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +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::StructArray; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::expr::BoundExpression; +use crate::expr::Expression; +use crate::expr::get_item; +use crate::expr::root; +use crate::expr::traversal::NodeExt; +use crate::expr::traversal::Transformed; +use crate::expr::traversal::TraversalOrder; +use crate::expr::traversal::pre_order_visit_down; +use crate::scalar_fn::internal::row_count::RowCount; +use crate::validity::Validity; + +/// Parameter names and a body: the function argument of a higher-order scalar function, as +/// written. +/// +/// The body is an ordinary [`Expression`] that reads its parameters with +/// [`var`](crate::expr::var), and nothing else from outside the lambda. +/// +/// ```rust +/// # use vortex_array::expr::{Lambda, checked_add, lit, var}; +/// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap(); +/// assert_eq!(lambda.to_string(), "(x) -> (@x + 1i32)"); +/// ``` +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Lambda { + params: FieldNames, + body: Expression, +} + +impl Lambda { + /// Create a lambda with these parameters and body. + /// + /// Returns an error when a parameter repeats, or when the body reads from outside the lambda: + /// `$` or a column, another lambda's parameter, or the scan's `row_idx()` or `row_count()`. + pub fn new( + params: impl IntoIterator>, + body: Expression, + ) -> VortexResult { + let params: FieldNames = params.into_iter().map(Into::into).collect(); + if let Some(param) = params.iter().duplicates().next() { + vortex_bail!("lambda parameter {param} is declared more than once"); + } + pre_order_visit_down(&body, |node: &Expression| { + vortex_ensure!( + !reads_outside_lambda(node, params.as_ref()), + "lambda ({}) can only read its parameters, not {node}; reading anything else is \ + not supported yet", + params.iter().format(", ") + ); + Ok(TraversalOrder::Continue) + })?; + + Ok(Self { params, body }) + } + + /// The parameter names, in order. + pub fn params(&self) -> &[FieldName] { + self.params.as_ref() + } + + /// The body, as written. + pub fn body(&self) -> &Expression { + &self.body + } + + /// The dtype of the body, given the parameters' dtypes, without binding it. + pub fn return_dtype(&self, param_dtypes: &[DType]) -> VortexResult { + let inputs = inputs_dtype(self.param_fields(param_dtypes)?); + self.body_over_inputs().return_dtype(&inputs) + } + + /// Bind the body to these parameter dtypes: `@x` becomes `$.x`, a field of the struct of the + /// lambda's inputs. + /// + /// A function that takes a lambda calls this from its + /// [`bind_options`](crate::scalar_fn::ScalarFnVTable::bind_options), so the body is bound once, + /// with the node that holds it, and never while the function runs. + pub fn bind(&self, param_dtypes: &[DType]) -> VortexResult { + let params = self.param_fields(param_dtypes)?; + let body = self + .body_over_inputs() + .bind(&inputs_dtype(params.clone()))?; + Ok(BoundLambda { params, body }) + } + + /// Whether nothing in the body can fail, so that the body may run over values whose results + /// are never read, such as dictionary values that no row uses. + pub fn is_infallible(&self) -> bool { + let mut infallible = true; + pre_order_visit_down(&self.body, |node: &Expression| { + infallible &= node + .as_scalar() + .is_none_or(|scalar_fn| scalar_fn.signature().is_infallible()); + Ok(if infallible { + TraversalOrder::Continue + } else { + TraversalOrder::Stop + }) + }) + .vortex_expect("expression traversal cannot fail"); + infallible + } + + fn param_fields(&self, param_dtypes: &[DType]) -> VortexResult { + vortex_ensure!( + param_dtypes.len() == self.params.len(), + "lambda {self} needs one dtype per parameter, got {}", + param_dtypes.len() + ); + Ok(StructFields::new( + self.params.clone(), + param_dtypes.to_vec(), + )) + } + + /// The body over the struct of the lambda's inputs, where `@x` reads the field `x`. + /// + /// A lambda in the body is a function's options, which this leaves alone: its variables are its + /// own parameters. + fn body_over_inputs(&self) -> Expression { + self.body + .clone() + .transform_down(|node| { + Ok(match &node { + Expression::Variable(name) => Transformed { + value: get_item(name.clone(), root()), + changed: true, + order: TraversalOrder::Skip, + }, + _ => Transformed::no(node), + }) + }) + .vortex_expect("expression traversal cannot fail") + .into_inner() + } +} + +/// The struct of a lambda's inputs, one field per parameter. +fn inputs_dtype(params: StructFields) -> DType { + DType::Struct(params, Nullability::NonNullable) +} + +impl Display for Lambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "({}) -> {}", self.params.iter().format(", "), self.body) + } +} + +/// A [`Lambda`] bound to its parameters' dtypes, ready to run. +/// +/// The body is a [`BoundExpression`] over the struct of the lambda's inputs, one field per +/// parameter: `@x` was bound as `$.x`. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct BoundLambda { + params: StructFields, + body: BoundExpression, +} + +impl BoundLambda { + /// The parameters' names and dtypes, in order. + pub fn params(&self) -> &StructFields { + &self.params + } + + /// The bound body, over the struct of the lambda's inputs. + pub fn body(&self) -> &BoundExpression { + &self.body + } + + /// The dtype of the body. + pub fn dtype(&self) -> &DType { + self.body.dtype() + } + + /// Whether the lambda is bound to these parameter dtypes. + pub fn is_bound_to(&self, param_dtypes: &[DType]) -> bool { + self.params.fields().eq(param_dtypes.iter().cloned()) + } + + /// Whether nothing in the body can fail. See [`Lambda::is_infallible`]. + pub fn is_infallible(&self) -> bool { + let mut infallible = true; + pre_order_visit_down(&self.body, |node: &BoundExpression| { + 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 + } + + /// Evaluate the body over `len` rows of `args`, one array per parameter, with the dtypes the + /// lambda is bound to. They become the fields of the struct the body is applied to. + /// + /// The result is computed later, like any applied expression, so errors from running the body + /// come from executing it. + pub fn evaluate(&self, args: Vec, len: usize) -> VortexResult { + vortex_ensure!( + args.len() == self.params.nfields(), + "lambda {self} needs one array per parameter, got {}", + args.len() + ); + for ((name, dtype), arg) in self + .params + .names() + .iter() + .zip(self.params.fields()) + .zip(&args) + { + vortex_ensure!( + arg.dtype() == &dtype, + "lambda {self}: parameter {name} is bound to {dtype}, not {}", + arg.dtype() + ); + } + let inputs = StructArray::try_new( + self.params.names().clone(), + args, + len, + Validity::NonNullable, + )? + .into_array(); + inputs.apply_bound(&self.body) + } +} + +impl Display for BoundLambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "({}) -> {}", + self.params.names().iter().format(", "), + self.body + ) + } +} + +/// Whether `node`, in the body of a lambda with these parameters, reads a value from outside the +/// lambda, which lambdas can't do yet: `$`, another lambda's parameter, or the scan's `row_idx()` +/// or `row_count()`. The scan replaces those two before anything runs, but never looks inside a +/// lambda's body. +fn reads_outside_lambda(node: &Expression, params: &[FieldName]) -> bool { + match node { + Expression::Root => true, + Expression::Variable(name) => !params.contains(name), + // `row_idx` is defined in vortex-layout, so it is recognized by its id. + Expression::Scalar { scalar_fn, .. } => { + scalar_fn.is::() || scalar_fn.id().as_ref() == "vortex.row_idx" + } + } +} diff --git a/vortex-array/src/expr/lambda/tests.rs b/vortex-array/src/expr/lambda/tests.rs new file mode 100644 index 00000000000..d1f4f076fe9 --- /dev/null +++ b/vortex-array/src/expr/lambda/tests.rs @@ -0,0 +1,343 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::StructFields; +use crate::expr::BoundLambda; +use crate::expr::Expression; +use crate::expr::Lambda; +use crate::expr::checked_add; +use crate::expr::col; +use crate::expr::eq; +use crate::expr::get_item; +use crate::expr::lit; +use crate::expr::proto::ExprSerializeProtoExt; +use crate::expr::root; +use crate::expr::var; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::internal::row_count::RowCount; + +/// `apply(value, lambda)` gives the lambda's one parameter the value of `value`, row by row: the +/// smallest function that takes a lambda. +#[derive(Clone)] +struct Apply; + +/// `apply`'s lambda, as written or bound to the value's dtype. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +enum ApplyOptions { + Lambda(Lambda), + Bound(BoundLambda), +} + +impl Display for ApplyOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Lambda(lambda) => lambda.fmt(f), + Self::Bound(lambda) => lambda.fmt(f), + } + } +} + +impl ScalarFnVTable for Apply { + type Options = ApplyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.apply"); + *ID + } + + fn arity(&self, _options: &ApplyOptions) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &ApplyOptions, _child_idx: usize) -> ChildName { + ChildName::from("value") + } + + fn return_dtype(&self, options: &ApplyOptions, args: &[DType]) -> VortexResult { + Ok(match options { + ApplyOptions::Lambda(lambda) => lambda.return_dtype(args)?, + ApplyOptions::Bound(lambda) => lambda.dtype().clone(), + }) + } + + fn bind_options( + &self, + options: &ApplyOptions, + args: &[DType], + ) -> VortexResult> { + match options { + ApplyOptions::Lambda(lambda) => Ok(Some(ApplyOptions::Bound(lambda.bind(args)?))), + ApplyOptions::Bound(lambda) => { + vortex_ensure!( + lambda.is_bound_to(args), + "{lambda} is bound to other dtypes" + ); + Ok(None) + } + } + } + + fn execute( + &self, + options: &ApplyOptions, + args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + let ApplyOptions::Bound(lambda) = options else { + vortex_bail!("apply's lambda is not bound"); + }; + lambda.evaluate(vec![args.get(0)?], args.row_count()) + } + + fn is_infallible(&self, options: &ApplyOptions) -> bool { + match options { + ApplyOptions::Lambda(lambda) => lambda.is_infallible(), + ApplyOptions::Bound(lambda) => lambda.is_infallible(), + } + } +} + +fn apply(value: Expression, param: &str, body: Expression) -> VortexResult { + Apply.try_new_expr(ApplyOptions::Lambda(Lambda::new([param], body)?), [value]) +} + +fn i32_() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) +} + +/// Rows `{a: i32, s: {i: i32}}`. +fn rows() -> VortexResult { + let s = StructArray::from_fields(&[("i", buffer![10_i32, 20, 30].into_array())])?; + Ok(StructArray::from_fields(&[ + ("a", buffer![1_i32, 2, 3].into_array()), + ("s", s.into_array()), + ])? + .into_array()) +} + +fn assert_applies(expr: &Expression, expected: ArrayRef) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(rows()?.apply(expr)?, expected, &mut ctx); + Ok(()) +} + +#[test] +fn new_keeps_params_and_body() -> VortexResult<()> { + let lambda = Lambda::new( + ["x", "p"], + eq(checked_add(var("x"), lit(1_i32)), get_item("k", var("p"))), + )?; + assert_eq!(lambda.params(), ["x", "p"].map(FieldName::from)); + assert_eq!( + lambda.body(), + &eq(checked_add(var("x"), lit(1_i32)), get_item("k", var("p"))) + ); + assert_eq!(lambda.to_string(), "(x, p) -> ((@x + 1i32) = @p.k)"); + Ok(()) +} + +#[rstest] +#[case::root(checked_add(var("x"), root()), "not $;")] +#[case::column(checked_add(var("x"), col("a")), "not $;")] +#[case::other_variable(checked_add(var("x"), var("p")), "not @p;")] +#[case::field_of_another_variable(get_item("k", var("p")), "not @p;")] +#[case::function_of_the_scan(RowCount.new_expr(EmptyOptions, []), "not row_count();")] +fn a_body_reads_only_its_parameters(#[case] body: Expression, #[case] read: &str) { + let error = Lambda::new(["x"], body).expect_err("reads from outside"); + let message = error.to_string(); + assert!( + message.contains("lambda (x) can only read its parameters") && message.contains(read), + "{message}" + ); +} + +#[test] +fn duplicate_parameters_are_rejected() { + let error = Lambda::new(["x", "x"], var("x")).expect_err("x is declared twice"); + assert!( + error.to_string().contains("declared more than once"), + "{error}" + ); +} + +#[test] +fn apply_runs_the_body_per_row() -> VortexResult<()> { + let expr = apply(col("a"), "x", checked_add(var("x"), lit(1_i32)))?; + assert_eq!( + expr.to_string(), + "vortex.test.apply($.a, opts=(x) -> (@x + 1i32))" + ); + assert_applies(&expr, buffer![2_i32, 3, 4].into_array()) +} + +#[test] +fn a_parameter_named_like_a_column() -> VortexResult<()> { + // `var("a")` is the parameter, whatever the table's column `a` holds. + let expr = apply( + get_item("i", col("s")), + "a", + checked_add(var("a"), lit(1_i32)), + )?; + assert_applies(&expr, buffer![11_i32, 21, 31].into_array()) +} + +#[test] +fn nested_lambdas_read_their_own_parameters() -> VortexResult<()> { + // p -> apply(p.i, y -> y + 1): the inner value reads the outer parameter. + let inner = apply( + get_item("i", var("p")), + "y", + checked_add(var("y"), lit(1_i32)), + )?; + let expr = apply(col("s"), "p", inner)?; + assert_eq!( + expr.to_string(), + "vortex.test.apply($.s, opts=(p) -> vortex.test.apply(@p.i, opts=(y) -> (@y + 1i32)))" + ); + assert_applies(&expr, buffer![11_i32, 21, 31].into_array()) +} + +#[test] +fn an_inner_parameter_shadows_an_outer_one() -> VortexResult<()> { + // x -> apply(x + 1, x -> x + x) + let inner = apply( + checked_add(var("x"), lit(1_i32)), + "x", + checked_add(var("x"), var("x")), + )?; + let expr = apply(col("a"), "x", inner)?; + assert_applies(&expr, buffer![4_i32, 6, 8].into_array()) +} + +#[rstest] +#[case::bind(|expr: &Expression, rows: ArrayRef| expr.bind(rows.dtype()).map(drop))] +#[case::return_dtype(|expr: &Expression, rows: ArrayRef| expr.return_dtype(rows.dtype()).map(drop))] +#[case::apply(|expr: &Expression, rows: ArrayRef| rows.apply(expr).map(drop))] +#[case::serialize(|expr: &Expression, _rows: ArrayRef| expr.serialize_proto().map(drop))] +fn a_variable_outside_a_lambda_is_an_error( + #[case] use_expr: fn(&Expression, ArrayRef) -> VortexResult<()>, +) -> VortexResult<()> { + let error = + use_expr(&checked_add(var("y"), lit(1_i32)), rows()?).expect_err("y is not a parameter"); + assert!( + error + .to_string() + .contains("variable @y is not a parameter of an enclosing lambda"), + "{error}" + ); + Ok(()) +} + +#[test] +fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { + let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?; + assert_eq!(lambda.return_dtype(&[i32_()])?, i32_()); + let bound = lambda.bind(&[i32_()])?; + assert_eq!(bound.dtype(), &i32_()); + assert!(bound.is_bound_to(&[i32_()])); + // The body reads the struct of the lambda's inputs: `@x` became `$.x`. + assert_eq!(bound.to_string(), "(x) -> ($.x + 1i32)"); + let inputs = DType::Struct( + StructFields::new(["x"].into(), vec![i32_()]), + Nullability::NonNullable, + ); + assert_eq!( + bound.body(), + &checked_add(col("x"), lit(1_i32)).bind(&inputs)? + ); + assert_eq!( + lambda.bind(&[i32_().as_nullable()])?.dtype(), + &i32_().as_nullable() + ); + for error in [ + lambda.return_dtype(&[]).expect_err("x has no dtype"), + lambda.bind(&[]).expect_err("x has no dtype"), + ] { + assert!( + error + .to_string() + .contains("lambda (x) -> (@x + 1i32) needs one dtype per parameter, got 0"), + "{error}" + ); + } + Ok(()) +} + +#[test] +fn binding_an_expression_binds_its_lambdas() -> VortexResult<()> { + let expr = apply(col("a"), "x", checked_add(var("x"), lit(1_i32)))?; + assert!(matches!(expr.as_::(), ApplyOptions::Lambda(_))); + let bound = expr.bind(rows()?.dtype())?; + assert!(matches!( + bound.as_::(), + ApplyOptions::Bound(lambda) if lambda.is_bound_to(&[i32_()]) + )); + Ok(()) +} + +#[test] +fn evaluate_runs_the_bound_body() -> VortexResult<()> { + let lambda = + Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?.bind(&[i32_().as_nullable()])?; + let x = PrimitiveArray::from_option_iter([Some(1_i32), None]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + lambda.evaluate(vec![x], 2)?, + PrimitiveArray::from_option_iter([Some(2_i32), None]).into_array(), + &mut ctx + ); + + // The parameters must have the dtypes the lambda is bound to. + let error = lambda + .evaluate(vec![buffer![1_i32, 2].into_array()], 2) + .expect_err("bound to i32?, given i32"); + assert!( + error + .to_string() + .contains("parameter x is bound to i32?, not i32"), + "{error}" + ); + Ok(()) +} + +#[rstest] +#[case::comparison(eq(var("x"), lit(3_i32)), true)] +#[case::checked_arithmetic(checked_add(var("x"), lit(1_i32)), false)] +fn infallibility_comes_from_the_body( + #[case] body: Expression, + #[case] infallible: bool, +) -> VortexResult<()> { + let lambda = Lambda::new(["x"], body)?; + assert_eq!(lambda.is_infallible(), infallible); + assert_eq!(lambda.bind(&[i32_()])?.is_infallible(), infallible); + Ok(()) +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 8dbb3c92bcb..230dae1f25d 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; @@ -109,8 +110,10 @@ pub use exprs::replace_time_zone; pub use exprs::root; pub use exprs::select; 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::*; #[cfg(feature = "_test-harness")] diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index c09da2c2a2f..fae410aaaa7 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -8,6 +8,7 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use crate::expr::Expression; +use crate::expr::expression::unbound_variable; use crate::proto::expr as pb; use crate::scalar_fn::ForeignScalarFnVTable; use crate::scalar_fn::ScalarFnId; @@ -24,6 +25,9 @@ pub(crate) const ROOT_ID: &str = "vortex.root"; impl ExprSerializeProtoExt for Expression { fn serialize_proto(&self) -> VortexResult { + if let Expression::Variable(name) = self { + return Err(unbound_variable(name)); + } let Some(scalar_fn) = self.as_scalar() else { return Ok(pb::Expr { id: ROOT_ID.to_string(), diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index 3d2684f2dc7..4f0628a1166 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -11,6 +11,7 @@ use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::expression::unbound_variable; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::literal::Literal; @@ -47,6 +48,9 @@ impl ArrayRef { if expr.is_root() { return Ok(self); } + if let Expression::Variable(name) = expr { + return Err(unbound_variable(name)); + } // Manually convert literals to ConstantArray. if let Some(scalar) = expr.as_opt::() { @@ -64,7 +68,7 @@ impl ArrayRef { // And wrap the scalar function up in an array. let scalar_fn = expr .as_scalar() - .vortex_expect("root and literal were handled above, so this is a scalar node"); + .vortex_expect("leaves and literals were handled above, so this is a scalar node"); let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array(); // Optimize the resulting array's root. diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index d1b09ce5f1f..ae951376778 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -122,6 +122,12 @@ impl ScalarFnRef { self.0.return_dtype(arg_types) } + /// Bind whatever in the options depends on the argument types, or keep this function if + /// nothing does. See [`ScalarFnVTable::bind_options`]. + pub(crate) fn bind_options(self, arg_types: &[DType]) -> VortexResult { + Ok(self.0.bind_options(arg_types)?.unwrap_or(self)) + } + /// Symbolic validity of this node in an expression tree pub(crate) fn validity_expression<'a>( &self, diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index aa5ba2d36b6..ab6db6e098e 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -81,6 +81,7 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { // Bound methods — options accessed from self fn execute(&self, args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; fn return_dtype(&self, arg_types: &[DType]) -> VortexResult; + fn bind_options(&self, arg_types: &[DType]) -> VortexResult>; fn reduce_expression<'a>( &self, node: &ExpressionReduceNode<'a>, @@ -169,6 +170,11 @@ impl DynScalarFn for TypedScalarFnInstance { V::return_dtype(&self.vtable, &self.options, arg_dtypes) } + fn bind_options(&self, arg_dtypes: &[DType]) -> VortexResult> { + Ok(V::bind_options(&self.vtable, &self.options, arg_dtypes)? + .map(|options| TypedScalarFnInstance::new(self.vtable.clone(), options).erased())) + } + fn reduce_expression<'a>( &self, node: &ExpressionReduceNode<'a>, diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index cf13ba98bcf..40cb9066b35 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -158,6 +158,24 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// [`Expression::try_new`]: crate::expr::Expression::try_new fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult; + /// Bind whatever in the options depends on the argument dtypes, such as the body of a + /// [`Lambda`](crate::expr::Lambda), returning the bound options. + /// + /// This is called whenever a bound node is built and whenever a + /// [`ScalarFnArray`] is built from its children, before + /// [`return_dtype`](Self::return_dtype), so [`execute`](Self::execute) finds the options bound. + /// A rebuilt node calls it again, possibly with other dtypes: options that are already bound + /// must be checked against these dtypes, returning `None` if they fit and an error if they + /// don't. Return `None` to keep the options as they are, which the default does. + fn bind_options( + &self, + options: &Self::Options, + args: &[DType], + ) -> VortexResult> { + _ = (options, args); + Ok(None) + } + /// Execute the expression over the input arguments. /// /// Implementations are encouraged to check their inputs for constant arrays to perform diff --git a/vortex-layout/src/layouts/row_idx/expr.rs b/vortex-layout/src/layouts/row_idx/expr.rs index 31dbbd4d47e..659eb1f69d7 100644 --- a/vortex-layout/src/layouts/row_idx/expr.rs +++ b/vortex-layout/src/layouts/row_idx/expr.rs @@ -71,3 +71,26 @@ impl ScalarFnVTable for RowIdx { pub fn row_idx() -> Expression { RowIdx.new_expr(EmptyOptions, []) } + +#[cfg(test)] +mod tests { + use vortex_array::expr::Lambda; + use vortex_array::expr::checked_add; + use vortex_array::expr::var; + + use super::row_idx; + + /// The scan replaces `row_idx()` before anything runs, but never looks inside a lambda's body, + /// so a lambda can't read it. `Lambda::new` recognizes it by its id. + #[test] + fn a_lambda_cannot_read_row_idx() { + let error = Lambda::new(["x"], checked_add(var("x"), row_idx())) + .expect_err("row_idx() is not a parameter"); + assert!( + error + .to_string() + .contains("lambda (x) can only read its parameters, not #row_idx"), + "{error}" + ); + } +} From f6662bb39493cb774e147dabea179da44ba70806 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 22 Oct 2026 15:01:41 -0400 Subject: [PATCH 2/2] hide bind_options from docs Signed-off-by: Matt Katz --- vortex-array/src/expr/lambda/mod.rs | 5 ++--- vortex-array/src/scalar_fn/vtable.rs | 1 + 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/expr/lambda/mod.rs b/vortex-array/src/expr/lambda/mod.rs index ea810dc428e..198226622a3 100644 --- a/vortex-array/src/expr/lambda/mod.rs +++ b/vortex-array/src/expr/lambda/mod.rs @@ -107,9 +107,8 @@ impl Lambda { /// Bind the body to these parameter dtypes: `@x` becomes `$.x`, a field of the struct of the /// lambda's inputs. /// - /// A function that takes a lambda calls this from its - /// [`bind_options`](crate::scalar_fn::ScalarFnVTable::bind_options), so the body is bound once, - /// with the node that holds it, and never while the function runs. + /// A function that takes a lambda calls this from its `bind_options`, so the body is bound + /// once, with the node that holds it, and never while the function runs. pub fn bind(&self, param_dtypes: &[DType]) -> VortexResult { let params = self.param_fields(param_dtypes)?; let body = self diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 40cb9066b35..4cd658d4774 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -167,6 +167,7 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// A rebuilt node calls it again, possibly with other dtypes: options that are already bound /// must be checked against these dtypes, returning `None` if they fit and an error if they /// don't. Return `None` to keep the options as they are, which the default does. + #[doc(hidden)] fn bind_options( &self, options: &Self::Options,