Add lambda expressions

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<BoundExpression> {
+    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) {

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

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,19 @@ 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.
+    ///
+    /// 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,
+        /// The value arguments.
+        children: Arc<Vec<Expression>>,
+        /// The lambda arguments, which follow the value arguments.
+        lambdas: Arc<Vec<Lambda>>,
+    },
 }
 
 impl Expression {
@@ -68,6 +83,26 @@ 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<Item = Expression>,
+        lambdas: impl IntoIterator<Item = Lambda>,
+    ) -> VortexResult<Self> {
+        let children = Vec::from_iter(children);
+        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 +112,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 +157,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 +184,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 +211,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<DType> {
         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 +241,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)
+            }
         }
     }
 

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<str>) -> 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<Item = impl Into<Variable>>,
+    body: Expression,
+) -> VortexResult<Lambda> {
+    Lambda::try_new(params, body)
+}
+
 /// Return whether the expression is a root expression.
 pub fn is_root(expr: &Expression) -> bool {
     expr.is_root()