lambdas for higher-order scalar functions

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

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;
@@ -310,6 +311,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 +322,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)
     }
 }

vortex-array/src/expr/display.rs

@@ -69,13 +69,15 @@ 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"),
         }
     }
 
     fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
             Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f),
             Expression::Root => write!(f, "{ROOT_DISPLAY}"),
+            Expression::Variable(name) => write!(f, "@{name}"),
         }
     }
 }

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,12 @@ pub enum Expression {
     },
     /// The full scope of the expression evaluation.
     Root,
+    /// A parameter of an enclosing lambda, read by name.
+    ///
+    /// Building the [`Lambda`](crate::expr::Lambda) replaces each variable of its parameters with a
+    /// field of the lambda's row, so a variable left over is not a parameter of any enclosing
+    /// lambda. It fails to bind and cannot be serialized.
+    Variable(FieldName),
 }
 
 impl Expression {
@@ -73,7 +84,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 +112,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 +136,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 +163,7 @@ impl Expression {
     pub fn return_dtype(&self, scope: &DType) -> VortexResult<DType> {
         match self {
             Self::Root => Ok(scope.clone()),
+            Self::Variable(name) => Err(unbound_variable(name)),
             Self::Scalar {
                 scalar_fn,
                 children,
@@ -164,6 +184,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 +270,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 {

vortex-array/src/expr/exprs.rs

@@ -146,6 +146,23 @@ pub fn bound_get_item(field: impl Into<FieldName>, 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 an enclosing lambda.
+///
+/// It is only meaningful inside a body passed to [`Lambda::new`](crate::expr::Lambda::new), which
+/// rewrites it into a field of the lambda's row. A variable that no enclosing lambda declares fails
+/// to bind.
+///
+/// ```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)");
+/// ```
+pub fn var(name: impl Into<FieldName>) -> Expression {
+    Expression::Variable(name.into())
+}
+
 // ---- VariantGet ----
 
 /// Creates an expression that extracts a path from a Variant expression.