Add ListTransform as a (higher order) scalar function

vortex-array/src/expr/bound_expression.rs

@@ -101,6 +101,7 @@ pub struct BoundLambda {
     params: Box<[Variable]>,
     param_dtypes: Box<[DType]>,
     param_refs: Box<[VariableRef]>,
+    captures: Box<[BoundVariable]>,
     parameter_frame: usize,
     body: Arc<BoundExpression>,
 }
@@ -135,6 +136,9 @@ impl BoundLambda {
             "lambda parameters must be bound in the innermost lexical frame"
         );
 
+        let body = lambda.body().bind_scope(scope)?;
+        let captures = collect_captures(&body, parameter_frame);
+
         Ok(Self {
             params: lambda.params().into(),
             param_dtypes: parameter_bindings
@@ -145,8 +149,9 @@ impl BoundLambda {
                 .into_iter()
                 .map(|(_, variable_ref)| variable_ref)
                 .collect(),
+            captures,
             parameter_frame,
-            body: Arc::new(lambda.body().bind_scope(scope)?),
+            body: Arc::new(body),
         })
     }
 
@@ -165,6 +170,11 @@ impl BoundLambda {
         &self.param_refs
     }
 
+    /// The outer lexical bindings read by this lambda body.
+    pub fn captures(&self) -> &[BoundVariable] {
+        &self.captures
+    }
+
     /// The lexical frame containing the parameters.
     pub fn parameter_frame(&self) -> usize {
         self.parameter_frame
@@ -182,33 +192,10 @@ impl BoundLambda {
 
     /// The outer lexical bindings read by this lambda body.
     pub fn free_variables(&self) -> Vec<VariableRef> {
-        fn collect(
-            expression: &BoundExpression,
-            parameter_frame: usize,
-            variables: &mut Vec<VariableRef>,
-        ) {
-            match expression {
-                BoundExpression::Variable(variable)
-                    if variable.variable_ref().frame() < parameter_frame
-                        && !variables.contains(&variable.variable_ref()) =>
-                {
-                    variables.push(variable.variable_ref());
-                }
-                BoundExpression::Scalar { children, .. } => {
-                    for child in children.iter() {
-                        collect(child, parameter_frame, variables);
-                    }
-                }
-                BoundExpression::Lambda(_)
-                | BoundExpression::Root { .. }
-                | BoundExpression::Variable(_) => {}
-            }
-        }
-
-        let mut variables = Vec::new();
-        collect(&self.body, self.parameter_frame, &mut variables);
-        variables.sort_by_key(|variable_ref| (variable_ref.frame(), variable_ref.slot()));
-        variables
+        self.captures
+            .iter()
+            .map(BoundVariable::variable_ref)
+            .collect()
     }
 
     fn take_body(&mut self) -> Option<BoundExpression> {
@@ -220,6 +207,41 @@ impl BoundLambda {
     }
 }
 
+fn collect_captures(expression: &BoundExpression, parameter_frame: usize) -> Box<[BoundVariable]> {
+    fn collect(
+        expression: &BoundExpression,
+        parameter_frame: usize,
+        captures: &mut Vec<BoundVariable>,
+    ) {
+        match expression {
+            BoundExpression::Variable(variable)
+                if variable.variable_ref().frame() < parameter_frame
+                    && !captures
+                        .iter()
+                        .any(|capture| capture.variable_ref() == variable.variable_ref()) =>
+            {
+                captures.push(variable.clone());
+            }
+            BoundExpression::Scalar { children, .. } => {
+                for child in children.iter() {
+                    collect(child, parameter_frame, captures);
+                }
+            }
+            BoundExpression::Lambda(_)
+            | BoundExpression::Root { .. }
+            | BoundExpression::Variable(_) => {}
+        }
+    }
+
+    let mut captures = Vec::new();
+    collect(expression, parameter_frame, &mut captures);
+    captures.sort_by_key(|capture| {
+        let variable_ref = capture.variable_ref();
+        (variable_ref.frame(), variable_ref.slot())
+    });
+    captures.into_boxed_slice()
+}
+
 impl Display for BoundLambda {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         write!(f, "({}) -> {}", self.params.iter().join(", "), self.body)

vortex-array/src/expression.rs

@@ -4,34 +4,34 @@
 use itertools::Itertools;
 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::dtype::FieldName;
+use crate::dtype::Nullability;
 use crate::expr::BoundExpression;
+use crate::expr::BoundLambda;
 use crate::expr::Expression;
+use crate::expr::VariableRef;
 use crate::optimizer::ArrayOptimizer;
 use crate::scalar_fn::ScalarFnRef;
+use crate::scalar_fn::ScalarFnVTableExt;
+use crate::scalar_fn::fns::get_item::GetItem;
 use crate::scalar_fn::fns::literal::Literal;
+use crate::scalar_fn::fns::pack::Pack;
+use crate::scalar_fn::fns::pack::PackOptions;
 
 impl ArrayRef {
     /// Apply a bound expression to this array, producing a new array in constant time.
     pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult<ArrayRef> {
-        match expr {
-            BoundExpression::Root { .. } => Ok(self),
-            BoundExpression::Lambda(_) => {
-                vortex_bail!("cannot apply a lambda outside a higher-order function")
-            }
-            BoundExpression::Variable(variable) => {
-                vortex_bail!("cannot apply variable '{variable}' without a provided value")
-            }
-            BoundExpression::Scalar {
-                scalar_fn,
-                children,
-                ..
-            } => apply_bound_scalar_fn(self, scalar_fn, children),
+        BoundApplyCtx {
+            root: &self,
+            bindings: None,
         }
+        .apply(expr)
     }
 
     /// Apply the expression to this array, producing a new array in constant time.
@@ -52,21 +52,158 @@ impl ArrayRef {
     }
 }
 
+impl BoundLambda {
+    /// Apply this lambda to arrays in a common invocation row domain.
+    ///
+    /// Parameters and captures are packed into a lazy non-nullable struct. Variable projections
+    /// reduce through that pack to the original arrays, leaving only the lambda body's lazy scalar
+    /// function array tree.
+    pub fn apply(
+        &self,
+        root: ArrayRef,
+        parameters: &[ArrayRef],
+        captures: &[ArrayRef],
+    ) -> VortexResult<ArrayRef> {
+        vortex_ensure!(
+            parameters.len() == self.param_dtypes().len(),
+            "lambda takes {} parameters but was applied with {} arguments",
+            self.param_dtypes().len(),
+            parameters.len()
+        );
+        vortex_ensure!(
+            captures.len() == self.captures().len(),
+            "lambda requires {} captures but was applied with {}",
+            self.captures().len(),
+            captures.len()
+        );
+        vortex_ensure!(
+            self.body().is_root_bound_to(root.dtype()),
+            "lambda root expects a different dtype than {}",
+            root.dtype()
+        );
+
+        for (index, (expected_dtype, parameter)) in
+            self.param_dtypes().iter().zip(parameters).enumerate()
+        {
+            vortex_ensure!(
+                parameter.dtype() == expected_dtype,
+                "lambda parameter {index} expects dtype {expected_dtype}, got {}",
+                parameter.dtype()
+            );
+            vortex_ensure!(
+                parameter.len() == root.len(),
+                "lambda parameter {index} has length {}, expected {}",
+                parameter.len(),
+                root.len()
+            );
+        }
+        for (index, (capture, array)) in self.captures().iter().zip(captures).enumerate() {
+            vortex_ensure!(
+                array.dtype() == capture.dtype(),
+                "lambda capture {index} expects dtype {}, got {}",
+                capture.dtype(),
+                array.dtype()
+            );
+            vortex_ensure!(
+                array.len() == root.len(),
+                "lambda capture {index} has length {}, expected {}",
+                array.len(),
+                root.len()
+            );
+        }
+
+        let names = self
+            .param_refs()
+            .iter()
+            .copied()
+            .chain(self.captures().iter().map(|capture| capture.variable_ref()))
+            .map(binding_name)
+            .collect::<Vec<_>>()
+            .into();
+        let fields = parameters.iter().chain(captures).cloned().collect();
+        let bindings = ScalarFnArray::try_new_with_len(
+            Pack.bind(PackOptions {
+                names,
+                nullability: Nullability::NonNullable,
+            }),
+            fields,
+            root.len(),
+        )?
+        .into_array();
+
+        let result = BoundApplyCtx {
+            root: &root,
+            bindings: Some(&bindings),
+        }
+        .apply(self.body())?;
+        vortex_ensure!(
+            result.dtype() == self.body_dtype(),
+            "lambda produced dtype {}, expected {}",
+            result.dtype(),
+            self.body_dtype()
+        );
+        vortex_ensure!(
+            result.len() == root.len(),
+            "lambda produced {} rows, expected {}",
+            result.len(),
+            root.len()
+        );
+        Ok(result)
+    }
+}
+
+struct BoundApplyCtx<'a> {
+    root: &'a ArrayRef,
+    bindings: Option<&'a ArrayRef>,
+}
+
+impl BoundApplyCtx<'_> {
+    fn apply(&self, expr: &BoundExpression) -> VortexResult<ArrayRef> {
+        match expr {
+            BoundExpression::Root { .. } => Ok(self.root.clone()),
+            BoundExpression::Lambda(_) => {
+                vortex_bail!("cannot apply a lambda outside a higher-order function")
+            }
+            BoundExpression::Variable(variable) => {
+                let Some(bindings) = self.bindings else {
+                    vortex_bail!("cannot apply variable '{variable}' without a provided value");
+                };
+                GetItem::try_new(bindings.clone(), binding_name(variable.variable_ref()))?
+                    .into_array()
+                    .optimize()
+            }
+            BoundExpression::Scalar {
+                scalar_fn,
+                children,
+                ..
+            } => apply_bound_scalar_fn(self, scalar_fn, children),
+        }
+    }
+}
+
+fn binding_name(variable_ref: VariableRef) -> FieldName {
+    FieldName::from(format!(
+        "frame[{}].slot[{}]",
+        variable_ref.frame(),
+        variable_ref.slot()
+    ))
+}
+
 fn apply_bound_scalar_fn(
-    root: ArrayRef,
+    ctx: &BoundApplyCtx<'_>,
     scalar_fn: &ScalarFnRef,
     children: &[BoundExpression],
 ) -> VortexResult<ArrayRef> {
     if let Some(scalar) = scalar_fn.as_opt::<Literal>() {
-        return Ok(ConstantArray::new(scalar.clone(), root.len()).into_array());
+        return Ok(ConstantArray::new(scalar.clone(), ctx.root.len()).into_array());
     }
 
     let children: Vec<_> = children
         .iter()
-        .map(|child| root.clone().apply_bound(child))
+        .map(|child| ctx.apply(child))
         .try_collect()?;
     let array =
-        ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, root.len())?.into_array();
+        ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, ctx.root.len())?.into_array();
     array.optimize()
 }
 
@@ -92,12 +229,20 @@ fn apply_scalar_fn(
 mod tests {
     use vortex_buffer::buffer;
     use vortex_error::VortexResult;
+    use vortex_error::vortex_bail;
 
+    use crate::ArrayRef;
     use crate::IntoArray;
+    use crate::arrays::ScalarFn;
+    use crate::arrays::scalar_fn::ScalarFnArrayExt;
+    use crate::expr::Lambda;
     use crate::expr::Scope;
     use crate::expr::Variable;
+    use crate::expr::binary;
     use crate::expr::lambda;
     use crate::expr::var;
+    use crate::scalar_fn::fns::binary::Binary;
+    use crate::scalar_fn::fns::operators::Operator;
 
     #[test]
     fn variable_application_requires_a_runtime_binding() -> VortexResult<()> {
@@ -120,4 +265,34 @@ mod tests {
         assert!(root.apply(&expression).is_err());
         Ok(())
     }
+
+    #[test]
+    fn bound_lambda_pack_is_eliminated() -> VortexResult<()> {
+        let parameter = buffer![1_i32, 2, 3].into_array();
+        let capture = buffer![10_i32, 20, 30].into_array();
+        let scope = Scope::new(parameter.dtype().clone())
+            .with_bindings([(Variable::new("capture"), capture.dtype().clone())])?
+            .with_bindings([(Variable::new("x"), parameter.dtype().clone())])?;
+
+        let identity = crate::expr::BoundLambda::bind(&Lambda::try_new(["x"], var("x"))?, &scope)?
+            .apply(parameter.clone(), std::slice::from_ref(&parameter), &[])?;
+        assert!(ArrayRef::ptr_eq(&identity, &parameter));
+
+        let lambda = crate::expr::BoundLambda::bind(
+            &Lambda::try_new(["x"], binary(Operator::Add, var("x"), var("capture")))?,
+            &scope,
+        )?;
+        let result = lambda.apply(
+            parameter.clone(),
+            std::slice::from_ref(&parameter),
+            std::slice::from_ref(&capture),
+        )?;
+        let Some(result) = result.as_opt::<ScalarFn>() else {
+            vortex_bail!("bound lambda did not produce a ScalarFnArray");
+        };
+        assert!(result.scalar_fn().is::<Binary>());
+        assert!(ArrayRef::ptr_eq(result.child_at(0), &parameter));
+        assert!(ArrayRef::ptr_eq(result.child_at(1), &capture));
+        Ok(())
+    }
 }