feat(functions): lambdas with transform, filter, exists and forall

crates/krishiv-exec/src/planner/eval/mod.rs

@@ -27,6 +27,7 @@ mod cast;
 mod cond;
 mod field;
 mod func;
+mod hof;
 mod logic;
 mod nested_cast;
 mod strings;
@@ -43,7 +44,7 @@ use krishiv_common::config::{
 };
 use krishiv_kernels::KernelCtx;
 use krishiv_plan::{
-    AttrId, Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, ResolvedExpr,
+    AttrId, BoundColumn, Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, ResolvedExpr,
     ResolvedFunction,
 };
 
@@ -212,7 +213,9 @@ impl EvalEnv {
         }
     }
 
-    /// The environment a plan fixed, with this worker's kernels.
+    /// The environment a plan fixed, with this worker's kernels. Lambda
+    /// variables match by stable identity, so no name-resolution flag is
+    /// plan data.
     pub(crate) fn from_plan(env: &ExprEnv) -> Self {
         Self {
             ansi: env.ansi,
@@ -250,6 +253,23 @@ pub(crate) struct Input<'a> {
     pub(crate) subqueries: &'a [(AttrId, Arc<dyn PhysExpr>)],
 }
 
+/// A column of `input`'s batch: the layout position's field. A lambda
+/// variable reads such a column of the flattened batch the higher-order
+/// call compiled around the body (see `hof`).
+fn column(input: Input<'_>, c: &BoundColumn) -> Result<Arc<dyn PhysExpr>> {
+    let index = input
+        .layout
+        .iter()
+        .position(|a| *a == c.id)
+        .ok_or_else(|| internal(&format!("attribute #{} is not an input column", c.id.0)))?;
+    let field = input.schema.field(index);
+    Ok(Arc::new(Column {
+        index,
+        data_type: field.data_type().clone(),
+        nullable: field.is_nullable(),
+    }))
+}
+
 /// Compiles `e` over `input`.
 pub(crate) fn compile(e: &Expr, input: Input<'_>) -> Result<Arc<dyn PhysExpr>> {
     match &e.node {
@@ -258,20 +278,11 @@ pub(crate) fn compile(e: &Expr, input: Input<'_>) -> Result<Arc<dyn PhysExpr>> {
             if let Some((_, v)) = input.subqueries.iter().find(|(a, _)| *a == c.id) {
                 return Ok(Arc::clone(v));
             }
-            let index = input
-                .layout
-                .iter()
-                .position(|a| *a == c.id)
-                .ok_or_else(|| {
-                    internal(&format!("attribute #{} is not an input column", c.id.0))
-                })?;
-            let field = input.schema.field(index);
-            Ok(Arc::new(Column {
-                index,
-                data_type: field.data_type().clone(),
-                nullable: field.is_nullable(),
-            }))
+            column(input, c)
         }
+        // A lambda variable: a column of the flattened batch the
+        // higher-order call compiled around the body (see `hof`).
+        ExprNode::Resolved(ResolvedExpr::LambdaVar(c)) => column(input, c),
         ExprNode::Resolved(ResolvedExpr::Alias { expr, .. }) => compile(expr, input),
         ExprNode::Resolved(ResolvedExpr::Field {
             base,
@@ -287,6 +298,11 @@ pub(crate) fn compile(e: &Expr, input: Input<'_>) -> Result<Arc<dyn PhysExpr>> {
         ),
         ExprNode::Cast { expr, to, mode } => cast::compile(compile(expr, input)?, to, *mode, input),
         ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Scalar => {
+            if let FunctionId::Registry { name, .. } = &f.id
+                && hof::is_supported(name)
+            {
+                return hof::compile(name, f, input, e);
+            }
             let args = f
                 .args
                 .iter()

crates/krishiv-functions/src/python_wrappers/table.rs

@@ -43,6 +43,22 @@ pub(super) const NOT_EXPORTED: &[(&str, &str)] = &[
         "pow",
         "an alias of `power`; the exported `pow` is a row that sends `power`",
     ),
+    (
+        "transform",
+        "takes a lambda; the Python `Column` API cannot build lambdas yet",
+    ),
+    (
+        "filter",
+        "takes a lambda; the Python `Column` API cannot build lambdas yet",
+    ),
+    (
+        "exists",
+        "takes a lambda; the Python `Column` API cannot build lambdas yet",
+    ),
+    (
+        "forall",
+        "takes a lambda; the Python `Column` API cannot build lambdas yet",
+    ),
 ];
 
 /// Registry names whose wrapper is written by hand in `functions.py`.
@@ -264,7 +280,14 @@ pub(super) fn rows() -> Vec<Row> {
             "col: ColumnOrName, ignoreNulls: bool | Column = None",
         ));
     }
-    for name in ["coalesce", "hash", "named_struct", "struct", "xxhash64"] {
+    for name in [
+        "array",
+        "coalesce",
+        "hash",
+        "named_struct",
+        "struct",
+        "xxhash64",
+    ] {
         rows.push(Row::new(name, "*cols: ColumnOrName"));
     }
     for name in ["lpad", "rpad"] {

crates/krishiv-functions/src/registry/defs.rs

@@ -852,6 +852,7 @@ pub(super) fn specs() -> Vec<Spec> {
     out.extend(super::url::specs());
     out.extend(super::csv::specs());
     out.extend(super::text::specs());
+    out.extend(super::hof::specs());
     out.extend(super::r#struct::specs());
     out
 }

crates/krishiv-functions/src/registry/hof.rs

@@ -0,0 +1,177 @@
+// SPDX-License-Identifier: Apache-2.0
+//! The analysis signatures of the higher-order functions over arrays
+//! (`transform`, `filter`, `exists`, `forall`) and the `array`
+//! constructor they need (Spark's `ArrayHigherOrderFunction`s and
+//! `CreateArray`).
+//!
+//! The analyzer binds a lambda's variables before the generic path types
+//! the call (see `krishiv-plan`'s `hof`): at a lambda position the
+//! argument's type is its body's, so these rules read the body type from
+//! the argument. The executor compiles the bodies (see `krishiv-exec`'s
+//! `hof`); only `array` binds a kernel here.
+//!
+//! `array` is the minimal constructor this family needs (all arguments at
+//! their common type, as Spark's `CreateArray` types them). The collection
+//! functions own its completion (nested coercion edge cases).
+
+use krishiv_common::error::{DatatypeMismatchUnexpectedInputType, InternalError, KrishivError};
+use krishiv_types::{ArrayType, DataType};
+
+use super::Execution;
+use super::defs::Spec;
+use crate::signature::{ArgType, CallTypes, FunctionKind, Inputs};
+
+/// `array(...)`: an array of the arguments' common type (Spark's
+/// `CreateArray`; untyped `NULL` is `VOID`). The analyzer casts the
+/// arguments to their common type first.
+fn array_result(c: &CallTypes<'_>) -> Result<ArgType, KrishivError> {
+    let element = c.args().first().and_then(|a| a.data_type.clone());
+    let nullable = c.args().iter().any(|a| a.nullable);
+    Ok(ArgType::new(
+        DataType::Array(ArrayType::new(element.unwrap_or(DataType::Null), nullable).into()),
+        false,
+    ))
+}
+
+/// The call's SQL as Spark quotes it in errors: `name(arg, ...)`.
+fn call_sql(name: &str, c: &CallTypes<'_>) -> String {
+    let args: Vec<String> = (0..c.args().len())
+        .map(|i| c.sql(i).unwrap_or_default())
+        .collect();
+    format!("\"{name}({})\"", args.join(", "))
+}
+
+/// A type as Spark's message parameters quote it (`"VOID"` for untyped
+/// `NULL`).
+fn quoted(t: Option<&DataType>) -> String {
+    format!("\"{}\"", t.map_or_else(|| "VOID".to_owned(), DataType::sql))
+}
+
+/// The array argument's element type, `containsNull` and nullability, or
+/// the `UNEXPECTED_INPUT_TYPE` error quoting it (Spark's
+/// `checkInputDataTypes`). The arity check runs before typing, so a
+/// missing argument is an internal error.
+fn array_argument(name: &str, c: &CallTypes<'_>) -> Result<(DataType, bool, bool), KrishivError> {
+    let Some(arg) = c.args().first() else {
+        return Err(InternalError {
+            message: format!("{name} typed without its array argument"),
+        }
+        .into());
+    };
+    match &arg.data_type {
+        Some(DataType::Array(element)) => Ok((
+            element.element_type().clone(),
+            element.contains_null(),
+            arg.nullable,
+        )),
+        other => Err(DatatypeMismatchUnexpectedInputType {
+            sql_expr: call_sql(name, c),
+            param_index: "first".to_owned(),
+            required_type: "\"ARRAY\"".to_owned(),
+            input_sql: format!("\"{}\"", c.sql(0).unwrap_or_default()),
+            input_type: quoted(other.as_ref()),
+        }
+        .into()),
+    }
+}
+
+/// Whether the lambda body (the second argument's type) is boolean, or the
+/// `UNEXPECTED_INPUT_TYPE` error quoting it (Spark quotes the lambda,
+/// constant lambdas wrapped over their hidden variable by the analyzer).
+fn boolean_body(name: &str, c: &CallTypes<'_>) -> Result<bool, KrishivError> {
+    let body = c.args().get(1);
+    match body.and_then(|a| a.data_type.as_ref()) {
+        Some(DataType::Boolean) => Ok(body.is_some_and(|a| a.nullable)),
+        other => Err(DatatypeMismatchUnexpectedInputType {
+            sql_expr: call_sql(name, c),
+            param_index: "second".to_owned(),
+            required_type: "\"BOOLEAN\"".to_owned(),
+            input_sql: c.sql(1).unwrap_or_default(),
+            input_type: quoted(other),
+        }
+        .into()),
+    }
+}
+
+/// `transform(arr, f)`: an array of the body's type.
+fn transform_result(c: &CallTypes<'_>) -> Result<ArgType, KrishivError> {
+    let (_, _, nullable) = array_argument("transform", c)?;
+    let body = c.args().get(1);
+    let element = body
+        .and_then(|a| a.data_type.clone())
+        .unwrap_or(DataType::Null);
+    let contains_null = body.is_none_or(|a| a.nullable);
+    Ok(ArgType::new(
+        DataType::Array(ArrayType::new(element, contains_null).into()),
+        nullable,
+    ))
+}
+
+/// `filter(arr, p)`: the array itself, with a boolean body.
+fn filter_result(c: &CallTypes<'_>) -> Result<ArgType, KrishivError> {
+    let (element, contains_null, nullable) = array_argument("filter", c)?;
+    let _ = boolean_body("filter", c)?;
+    Ok(ArgType::new(
+        DataType::Array(ArrayType::new(element, contains_null).into()),
+        nullable,
+    ))
+}
+
+/// `exists(arr, p)`: a boolean, with a boolean body. The result is
+/// nullable when the array is or when the body is (a NULL mark decides
+/// nothing, so the fold stays NULL). Spark derives nullability from the
+/// array and body alone; `containsNull` is not part of the rule.
+fn exists_result(c: &CallTypes<'_>) -> Result<ArgType, KrishivError> {
+    let (_, _, nullable) = array_argument("exists", c)?;
+    let body_nullable = boolean_body("exists", c)?;
+    Ok(ArgType::new(DataType::Boolean, nullable || body_nullable))
+}
+
+/// `forall(arr, p)`: a boolean, with a boolean body (same nullability
+/// as `exists`: a NULL mark decides nothing either way).
+fn forall_result(c: &CallTypes<'_>) -> Result<ArgType, KrishivError> {
+    let (_, _, nullable) = array_argument("forall", c)?;
+    let body_nullable = boolean_body("forall", c)?;
+    Ok(ArgType::new(DataType::Boolean, nullable || body_nullable))
+}
+
+/// The signatures of this family.
+pub(super) fn specs() -> Vec<Spec> {
+    vec![
+        Spec {
+            name: "array",
+            kind: FunctionKind::Scalar,
+            inputs: Inputs::Common,
+            result: array_result,
+            execution: Execution::Scalar,
+        },
+        Spec {
+            name: "transform",
+            kind: FunctionKind::Scalar,
+            inputs: Inputs::Any,
+            result: transform_result,
+            execution: Execution::HigherOrder,
+        },
+        Spec {
+            name: "filter",
+            kind: FunctionKind::Scalar,
+            inputs: Inputs::Any,
+            result: filter_result,
+            execution: Execution::HigherOrder,
+        },
+        Spec {
+            name: "exists",
+            kind: FunctionKind::Scalar,
+            inputs: Inputs::Any,
+            result: exists_result,
+            execution: Execution::HigherOrder,
+        },
+        Spec {
+            name: "forall",
+            kind: FunctionKind::Scalar,
+            inputs: Inputs::Any,
+            result: forall_result,
+            execution: Execution::HigherOrder,
+        },
+    ]
+}

crates/krishiv-functions/src/registry/mod.rs

@@ -19,6 +19,7 @@ mod conditional;
 mod csv;
 mod defs;
 mod hashing;
+mod hof;
 mod math;
 mod misc;
 mod number;
@@ -51,6 +52,11 @@ pub enum Execution {
     /// By the executor's window operator, over a window (`row_number`,
     /// `rank`, ...); never called without `OVER`.
     Window,
+    /// By the executor's higher-order compiler, which evaluates the lambda
+    /// bodies: `transform`, `filter`, `exists`, `forall`
+    /// (`krishiv_exec::planner::eval::hof`, over `krishiv_kernels::list`).
+    /// A lambda is no Arrow value, so these never bind a kernel here.
+    HigherOrder,
 }
 
 /// The aggregates the executor implements.