lambdas for higher-order scalar functions

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

Author: mhk197

vortex-array/src/arrays/scalar_fn/array.rs

@@ -110,6 +110,7 @@ impl Array<ScalarFn> {
         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) })

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)
     }
 }

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,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<DType> {
         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 {

vortex-array/src/expr/exprs.rs

@@ -149,6 +149,25 @@ 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 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<FieldName>) -> Expression {
+    Expression::Variable(name.into())
+}
+
 // ---- VariantGet ----
 
 /// Creates an expression that extracts a path from a Variant expression.