From b660eee03bfb77809f690cdc1874f00a6eb7af0d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 27 Oct 2026 11:47:01 +0000 Subject: [PATCH] pass what a list_transform lambda reads from outside as its environment Signed-off-by: Matt Katz --- vortex-array/src/expr/exprs.rs | 18 +- vortex-array/src/expr/lambda/mod.rs | 240 ++++++++++++++---- vortex-array/src/expr/lambda/tests.rs | 98 +++++-- .../src/scalar_fn/fns/list_transform/mod.rs | 135 +++++++--- .../src/scalar_fn/fns/list_transform/tests.rs | 200 ++++++++++++++- vortex-expr/src/transform/replace.rs | 15 ++ vortex-layout/src/layouts/row_idx/expr.rs | 23 -- vortex-layout/src/layouts/row_idx/mod.rs | 66 +++++ 8 files changed, 665 insertions(+), 130 deletions(-) diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 96f67b582dc..9b6d6032e9e 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -3,6 +3,7 @@ //! Factory functions for creating [`Expression`]s from scalar function vtables. +use std::iter; use std::sync::Arc; use vortex_error::VortexExpect; @@ -1427,16 +1428,29 @@ pub fn bound_list_sum_opts( /// Creates an expression that applies a lambda to every element of each list. /// -/// The lambda takes one parameter, the element. +/// The lambda takes one parameter, the element. When its body reads values from outside it, such as +/// a column, they are the call's second input, the lambda's [environment](Lambda::environment): +/// every element of a row's list sees that row's values. /// /// ```rust /// # use vortex_array::expr::{Lambda, checked_add, col, list_transform, lit, var}; /// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap(); /// let expr = list_transform(col("l"), lambda); /// assert_eq!(expr.to_string(), "vortex.list.transform($.l, (x) -> (@x + 1i32))"); +/// +/// let lambda = Lambda::new(["x"], checked_add(var("x"), col("a"))).unwrap(); +/// let expr = list_transform(col("l"), lambda); +/// assert_eq!( +/// expr.to_string(), +/// "vortex.list.transform($.l, pack($: pack(a: $.a)), (x) -> (@x + $.a))" +/// ); /// ``` pub fn list_transform(list: Expression, lambda: Lambda) -> Expression { - ListTransform.new_expr(ListTransformOptions::Lambda(lambda), [list]) + let environment = lambda.environment(); + ListTransform.new_expr( + ListTransformOptions::Lambda(lambda), + iter::once(list).chain(environment), + ) } /// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit. diff --git a/vortex-array/src/expr/lambda/mod.rs b/vortex-array/src/expr/lambda/mod.rs index 198226622a3..a923e43c987 100644 --- a/vortex-array/src/expr/lambda/mod.rs +++ b/vortex-array/src/expr/lambda/mod.rs @@ -8,10 +8,14 @@ //! function's node is bound, it binds the lambda to its parameters' dtypes, once, and keeps the //! [`BoundLambda`] instead. //! -//! Binding turns the lambda into an ordinary expression over a struct of its inputs, one field per -//! parameter: `@x` becomes `$.x`. To run it, the function puts the parameters' arrays in that -//! struct, such as one row per list element, and applies the bound body to it. The body reads -//! nothing but its parameters, so the function's children are everything it reads. +//! Binding turns the lambda into an ordinary expression over a struct of its inputs: `@x` becomes +//! `$.x`. To run it, the function puts the inputs' arrays in that struct, such as one row per list +//! element, and applies the bound body to it. +//! +//! The body may also read values from outside the lambda, such as a column. The function takes +//! them as one more child, the lambda's [environment](Lambda::environment), so that a rewrite of +//! the expression sees everything it reads. The environment's fields are the lambda's other +//! inputs, named the way the body reads them, so the body itself stays as written. #[cfg(test)] mod tests; @@ -37,11 +41,13 @@ use crate::dtype::StructFields; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::get_item; +use crate::expr::pack; use crate::expr::root; use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; +use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::internal::row_count::RowCount; use crate::validity::Validity; @@ -49,7 +55,8 @@ use crate::validity::Validity; /// written. /// /// The body is an ordinary [`Expression`] that reads its parameters with -/// [`var`](crate::expr::var), and nothing else from outside the lambda. +/// [`var`](crate::expr::var). It may also read from outside the lambda: see +/// [`Lambda::environment`]. /// /// ```rust /// # use vortex_array::expr::{Lambda, checked_add, lit, var}; @@ -65,8 +72,8 @@ pub struct Lambda { impl Lambda { /// Create a lambda with these parameters and body. /// - /// Returns an error when a parameter repeats, or when the body reads from outside the lambda: - /// `$` or a column, another lambda's parameter, or the scan's `row_idx()` or `row_count()`. + /// Returns an error when a parameter repeats, or is named `$`, the name of the row in the + /// lambda's environment. pub fn new( params: impl IntoIterator>, body: Expression, @@ -75,19 +82,39 @@ impl Lambda { if let Some(param) = params.iter().duplicates().next() { vortex_bail!("lambda parameter {param} is declared more than once"); } - pre_order_visit_down(&body, |node: &Expression| { - vortex_ensure!( - !reads_outside_lambda(node, params.as_ref()), - "lambda ({}) can only read its parameters, not {node}; reading anything else is \ - not supported yet", - params.iter().format(", ") - ); - Ok(TraversalOrder::Continue) - })?; - + vortex_ensure!( + params.iter().all(|param| param.as_ref() != ROW), + "a lambda parameter can't be named {ROW}, the row the lambda reads from outside" + ); Ok(Self { params, body }) } + /// The values the body reads from outside the lambda, as one struct, or `None` when it reads + /// nothing but its parameters. + /// + /// A function that takes the lambda takes this as one more child, so that a rewrite of the + /// expression, such as `replace` or field pruning, sees everything the lambda reads. Each field + /// is named the way the body reads it: `$` for the row, another lambda's parameter by its name, + /// and the scan's `row_idx()` or `row_count()` by its id. A value is cut down to the fields the + /// body reads below it. + /// + /// ```rust + /// # use vortex_array::expr::{Lambda, checked_add, col, var}; + /// let lambda = Lambda::new(["x"], checked_add(var("x"), col("a"))).unwrap(); + /// assert_eq!(lambda.environment().unwrap().to_string(), "pack($: pack(a: $.a))"); + /// ``` + pub fn environment(&self) -> Option { + let reads = self.outside_reads(); + (!reads.is_empty()).then(|| { + pack( + reads + .into_iter() + .map(|read| (read.name, project(read.value, &read.paths))), + Nullability::NonNullable, + ) + }) + } + /// The parameter names, in order. pub fn params(&self) -> &[FieldName] { self.params.as_ref() @@ -98,19 +125,30 @@ impl Lambda { &self.body } - /// The dtype of the body, given the parameters' dtypes, without binding it. - pub fn return_dtype(&self, param_dtypes: &[DType]) -> VortexResult { - let inputs = inputs_dtype(self.param_fields(param_dtypes)?); + /// The dtype of the body, given the parameters' dtypes and the fields of the environment, + /// without binding it. + pub fn return_dtype( + &self, + param_dtypes: &[DType], + environment: &StructFields, + ) -> VortexResult { + let inputs = inputs_dtype(self.inputs(param_dtypes, environment)?); self.body_over_inputs().return_dtype(&inputs) } - /// Bind the body to these parameter dtypes: `@x` becomes `$.x`, a field of the struct of the - /// lambda's inputs. + /// Bind the body to these parameter dtypes and fields of the [environment](Self::environment) + /// (empty without one): `@x` becomes `$.x`, a field of the struct of the lambda's inputs, which + /// are its parameters followed by the environment's fields. `$` becomes `$.$`, the row in the + /// environment. /// /// A function that takes a lambda calls this from its `bind_options`, so the body is bound /// once, with the node that holds it, and never while the function runs. - pub fn bind(&self, param_dtypes: &[DType]) -> VortexResult { - let params = self.param_fields(param_dtypes)?; + pub fn bind( + &self, + param_dtypes: &[DType], + environment: &StructFields, + ) -> VortexResult { + let params = self.inputs(param_dtypes, environment)?; let body = self .body_over_inputs() .bind(&inputs_dtype(params.clone()))?; @@ -135,32 +173,80 @@ impl Lambda { infallible } - fn param_fields(&self, param_dtypes: &[DType]) -> VortexResult { + /// The lambda's inputs: its parameters, with these dtypes, then the fields of its environment. + fn inputs( + &self, + param_dtypes: &[DType], + environment: &StructFields, + ) -> VortexResult { vortex_ensure!( param_dtypes.len() == self.params.len(), "lambda {self} needs one dtype per parameter, got {}", param_dtypes.len() ); - Ok(StructFields::new( - self.params.clone(), - param_dtypes.to_vec(), - )) + for read in self.outside_reads() { + vortex_ensure!( + environment.find(&read.name).is_some(), + "lambda {self} reads {} from outside, which its function must pass in as the \ + lambda's environment", + read.value + ); + } + let names: FieldNames = self + .params + .iter() + .chain(environment.names().iter()) + .cloned() + .collect(); + let dtypes = param_dtypes + .iter() + .cloned() + .chain(environment.fields()) + .collect(); + Ok(StructFields::new(names, dtypes)) } - /// The body over the struct of the lambda's inputs, where `@x` reads the field `x`. + /// The values the body reads from outside the lambda, in the order it first reads them. + fn outside_reads(&self) -> Vec { + let mut reads: Vec = Vec::new(); + pre_order_visit_down(&self.body, |node: &Expression| { + let Some((value, path)) = outside_read(node, self.params.as_ref()) else { + return Ok(TraversalOrder::Continue); + }; + let name = environment_name(&value); + match reads.iter_mut().find(|read| read.name == name) { + Some(read) => read.paths.push(path), + None => reads.push(OutsideRead { + name, + value, + paths: vec![path], + }), + } + Ok(TraversalOrder::Skip) + }) + .vortex_expect("expression traversal cannot fail"); + reads + } + + /// The body over the struct of the lambda's inputs, where `@x` reads the field `x` and a value + /// from outside the lambda reads its field of the environment: `$` reads `$.$`. /// - /// A lambda in the body is a function's options, which this leaves alone: its variables are its - /// own parameters. + /// A lambda in the body is a function's options, which this leaves alone: it reads its own + /// inputs. fn body_over_inputs(&self) -> Expression { + let input = |name: FieldName| Transformed { + value: get_item(name, root()), + changed: true, + order: TraversalOrder::Skip, + }; self.body .clone() .transform_down(|node| { Ok(match &node { - Expression::Variable(name) => Transformed { - value: get_item(name.clone(), root()), - changed: true, - order: TraversalOrder::Skip, - }, + Expression::Variable(name) => input(name.clone()), + _ if reads_outside_lambda(&node, self.params.as_ref()) => { + input(environment_name(&node)) + } _ => Transformed::no(node), }) }) @@ -169,6 +255,66 @@ impl Lambda { } } +/// The field of a lambda's environment that holds the row. +const ROW: &str = "$"; + +/// A value the body of a lambda reads from outside it. +struct OutsideRead { + /// The field of the lambda's environment that holds the value. + name: FieldName, + value: Expression, + /// The fields the body reads below the value, in the order it reads them. An empty path reads + /// the whole value. + paths: Vec>, +} + +/// The value from outside a lambda with these parameters that `node` reads, and the fields it reads +/// below it: `$` and `[s, a]` for `$.s.a`. +fn outside_read(node: &Expression, params: &[FieldName]) -> Option<(Expression, Vec)> { + let mut path = Vec::new(); + let mut value = node; + while let Some(field) = value.as_opt::() { + path.push(field.clone()); + value = &value.children()[0]; + } + path.reverse(); + reads_outside_lambda(value, params).then(|| (value.clone(), path)) +} + +/// The field of a lambda's environment that holds this value from outside the lambda. +fn environment_name(value: &Expression) -> FieldName { + match value { + Expression::Root => FieldName::from(ROW), + Expression::Variable(name) => name.clone(), + Expression::Scalar { scalar_fn, .. } => FieldName::from(scalar_fn.id().as_ref()), + } +} + +/// `value` cut down to the fields read at `paths`, keeping their nesting: `pack(a: $.a, s: pack(i: +/// $.s.i))` for `$` read at `a` and `s.i`. An empty path reads the whole value. +fn project(value: Expression, paths: &[Vec]) -> Expression { + if paths.iter().any(Vec::is_empty) { + return value; + } + let fields = paths + .iter() + .map(|path| path[0].clone()) + .unique() + .collect_vec(); + pack( + fields.into_iter().map(|field| { + let below = paths + .iter() + .filter(|path| path[0] == field) + .map(|path| path[1..].to_vec()) + .collect_vec(); + let value = project(get_item(field.clone(), value.clone()), &below); + (field, value) + }), + Nullability::NonNullable, + ) +} + /// The struct of a lambda's inputs, one field per parameter. fn inputs_dtype(params: StructFields) -> DType { DType::Struct(params, Nullability::NonNullable) @@ -180,10 +326,10 @@ impl Display for Lambda { } } -/// A [`Lambda`] bound to its parameters' dtypes, ready to run. +/// A [`Lambda`] bound to its inputs' dtypes, ready to run. /// -/// The body is a [`BoundExpression`] over the struct of the lambda's inputs, one field per -/// parameter: `@x` was bound as `$.x`. +/// The body is a [`BoundExpression`] over the struct of the lambda's inputs, its parameters then +/// its environment's fields: `@x` was bound as `$.x`. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct BoundLambda { params: StructFields, @@ -191,7 +337,7 @@ pub struct BoundLambda { } impl BoundLambda { - /// The parameters' names and dtypes, in order. + /// The inputs' names and dtypes, in order: the parameters, then the environment's fields. pub fn params(&self) -> &StructFields { &self.params } @@ -206,7 +352,7 @@ impl BoundLambda { self.body.dtype() } - /// Whether the lambda is bound to these parameter dtypes. + /// Whether the lambda is bound to these input dtypes. pub fn is_bound_to(&self, param_dtypes: &[DType]) -> bool { self.params.fields().eq(param_dtypes.iter().cloned()) } @@ -228,8 +374,8 @@ impl BoundLambda { infallible } - /// Evaluate the body over `len` rows of `args`, one array per parameter, with the dtypes the - /// lambda is bound to. They become the fields of the struct the body is applied to. + /// Evaluate the body over `len` rows of `args`, one array per input, with the dtypes the lambda + /// is bound to. They become the fields of the struct the body is applied to. /// /// The result is computed later, like any applied expression, so errors from running the body /// come from executing it. @@ -275,9 +421,9 @@ impl Display for BoundLambda { } /// Whether `node`, in the body of a lambda with these parameters, reads a value from outside the -/// lambda, which lambdas can't do yet: `$`, another lambda's parameter, or the scan's `row_idx()` -/// or `row_count()`. The scan replaces those two before anything runs, but never looks inside a -/// lambda's body. +/// lambda: `$`, another lambda's parameter, or the scan's `row_idx()` or `row_count()`. The scan +/// replaces those two before anything runs, but never looks inside a lambda's body, so they too +/// come in through the environment. fn reads_outside_lambda(node: &Expression, params: &[FieldName]) -> bool { match node { Expression::Root => true, diff --git a/vortex-array/src/expr/lambda/tests.rs b/vortex-array/src/expr/lambda/tests.rs index d09592bdf11..0ba9c22c2d4 100644 --- a/vortex-array/src/expr/lambda/tests.rs +++ b/vortex-array/src/expr/lambda/tests.rs @@ -87,7 +87,7 @@ impl ScalarFnVTable for Apply { fn return_dtype(&self, options: &ApplyOptions, args: &[DType]) -> VortexResult { Ok(match options { - ApplyOptions::Lambda(lambda) => lambda.return_dtype(args)?, + ApplyOptions::Lambda(lambda) => lambda.return_dtype(args, &StructFields::empty())?, ApplyOptions::Bound(lambda) => lambda.dtype().clone(), }) } @@ -98,7 +98,9 @@ impl ScalarFnVTable for Apply { args: &[DType], ) -> VortexResult> { match options { - ApplyOptions::Lambda(lambda) => Ok(Some(ApplyOptions::Bound(lambda.bind(args)?))), + ApplyOptions::Lambda(lambda) => Ok(Some(ApplyOptions::Bound( + lambda.bind(args, &StructFields::empty())?, + ))), ApplyOptions::Bound(lambda) => { vortex_ensure!( lambda.is_bound_to(args), @@ -169,18 +171,58 @@ fn new_keeps_params_and_body() -> VortexResult<()> { } #[rstest] -#[case::root(checked_add(var("x"), root()), "not $;")] -#[case::column(checked_add(var("x"), col("a")), "not $;")] -#[case::other_variable(checked_add(var("x"), var("p")), "not @p;")] -#[case::field_of_another_variable(get_item("k", var("p")), "not @p;")] -#[case::function_of_the_scan(RowCount.new_expr(EmptyOptions, []), "not row_count();")] -fn a_body_reads_only_its_parameters(#[case] body: Expression, #[case] read: &str) { - let error = Lambda::new(["x"], body).expect_err("reads from outside"); - let message = error.to_string(); - assert!( - message.contains("lambda (x) can only read its parameters") && message.contains(read), - "{message}" +#[case::only_its_parameter(checked_add(var("x"), lit(1_i32)), None)] +#[case::field_of_its_parameter(get_item("k", var("x")), None)] +#[case::column(checked_add(var("x"), col("a")), Some("pack($: pack(a: $.a))"))] +#[case::read_twice(checked_add(col("a"), col("a")), Some("pack($: pack(a: $.a))"))] +#[case::nested_field( + checked_add(var("x"), get_item("i", col("s"))), + Some("pack($: pack(s: pack(i: $.s.i)))") +)] +#[case::field_and_its_struct(eq(get_item("i", col("s")), col("s")), Some("pack($: pack(s: $.s))"))] +#[case::root(eq(var("x"), root()), Some("pack($: $)"))] +#[case::other_variable( + checked_add(var("x"), get_item("k", var("p"))), + Some("pack(p: pack(k: @p.k))") +)] +#[case::function_of_the_scan( + checked_add(var("x"), RowCount.new_expr(EmptyOptions, [])), + Some("pack(vortex.row_count: row_count())") +)] +fn the_environment_holds_what_the_body_reads_from_outside( + #[case] body: Expression, + #[case] environment: Option<&str>, +) -> VortexResult<()> { + let lambda = Lambda::new(["x"], body.clone())?; + assert_eq!( + lambda.environment().map(|e| e.to_string()).as_deref(), + environment ); + assert_eq!(lambda.body(), &body); + Ok(()) +} + +#[rstest] +#[case::root(checked_add(var("x"), root()), "reads $ from outside")] +#[case::column(checked_add(var("x"), col("a")), "reads $ from outside")] +#[case::other_variable(checked_add(var("x"), var("p")), "reads @p from outside")] +#[case::field_of_another_variable(get_item("k", var("p")), "reads @p from outside")] +#[case::function_of_the_scan(RowCount.new_expr(EmptyOptions, []), "reads row_count() from outside")] +fn binding_without_the_environment_is_an_error( + #[case] body: Expression, + #[case] read: &str, +) -> VortexResult<()> { + let error = Lambda::new(["x"], body)? + .bind(&[i32_()], &StructFields::empty()) + .expect_err("reads from outside"); + assert!(error.to_string().contains(read), "{error}"); + Ok(()) +} + +#[test] +fn a_parameter_cannot_be_named_like_the_row() { + let error = Lambda::new(["$"], var("$")).expect_err("$ is the row"); + assert!(error.to_string().contains("can't be named $"), "{error}"); } #[test] @@ -263,8 +305,11 @@ fn a_variable_outside_a_lambda_is_an_error( #[test] fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?; - assert_eq!(lambda.return_dtype(&[i32_()])?, i32_()); - let bound = lambda.bind(&[i32_()])?; + assert_eq!( + lambda.return_dtype(&[i32_()], &StructFields::empty())?, + i32_() + ); + let bound = lambda.bind(&[i32_()], &StructFields::empty())?; assert_eq!(bound.dtype(), &i32_()); assert!(bound.is_bound_to(&[i32_()])); // The body reads the struct of the lambda's inputs: `@x` became `$.x`. @@ -278,12 +323,18 @@ fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { &checked_add(col("x"), lit(1_i32)).bind(&inputs)? ); assert_eq!( - lambda.bind(&[i32_().as_nullable()])?.dtype(), + lambda + .bind(&[i32_().as_nullable()], &StructFields::empty())? + .dtype(), &i32_().as_nullable() ); for error in [ - lambda.return_dtype(&[]).expect_err("x has no dtype"), - lambda.bind(&[]).expect_err("x has no dtype"), + lambda + .return_dtype(&[], &StructFields::empty()) + .expect_err("x has no dtype"), + lambda + .bind(&[], &StructFields::empty()) + .expect_err("x has no dtype"), ] { assert!( error @@ -309,8 +360,8 @@ fn binding_an_expression_binds_its_lambdas() -> VortexResult<()> { #[test] fn evaluate_runs_the_bound_body() -> VortexResult<()> { - let lambda = - Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?.bind(&[i32_().as_nullable()])?; + let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))? + .bind(&[i32_().as_nullable()], &StructFields::empty())?; let x = PrimitiveArray::from_option_iter([Some(1_i32), None]).into_array(); let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( @@ -341,6 +392,11 @@ fn infallibility_comes_from_the_body( ) -> VortexResult<()> { let lambda = Lambda::new(["x"], body)?; assert_eq!(lambda.is_infallible(), infallible); - assert_eq!(lambda.bind(&[i32_()])?.is_infallible(), infallible); + assert_eq!( + lambda + .bind(&[i32_()], &StructFields::empty())? + .is_infallible(), + infallible + ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs index c7dc2e77ae3..c42ac209efc 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -7,9 +7,12 @@ mod tests; use std::fmt; use std::fmt::Display; use std::fmt::Formatter; +use std::iter; use std::sync::Arc; use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -22,16 +25,21 @@ use crate::IntoArray; use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; use crate::arrays::list::ListArrayExt; use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; +use crate::arrays::listview::list_from_list_view; +use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; -use crate::dtype::FieldName; +use crate::dtype::StructFields; use crate::expr::BoundLambda; use crate::expr::Lambda; use crate::expr::display::ExprDisplay; +use crate::match_each_integer_ptype; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; @@ -40,18 +48,19 @@ use crate::scalar_fn::ScalarFnVTable; /// Apply a lambda to every element of each list, keeping each list's length and validity. /// -/// The lambda takes one parameter, the element. The only child is the list, of `List` dtype in any -/// encoding, and the result is a `List` of the same shape. A lambda that can fail runs only on the -/// elements of valid lists. +/// The lambda takes one parameter, the element. The first child is the list, of `List` dtype in any +/// encoding, and the result is a `List`. When the lambda reads values from outside it, the second +/// child is its [environment](Lambda::environment), one struct per row: every element of a row's +/// list sees that row's values. A lambda that can fail runs only on the elements of valid lists. #[derive(Clone)] pub struct ListTransform; -/// The lambda of a [`ListTransform`], as written or bound to the element's dtype. +/// The lambda of a [`ListTransform`], as written or bound to the dtypes of its inputs. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum ListTransformOptions { /// The lambda as written, in an [`Expression`](crate::expr::Expression). Lambda(Lambda), - /// The lambda bound to the element's dtype, in a + /// The lambda bound to the dtypes of its inputs, in a /// [`BoundExpression`](crate::expr::BoundExpression) or an array. Bound(BoundLambda), } @@ -74,12 +83,16 @@ impl ScalarFnVTable for ListTransform { } fn arity(&self, _options: &ListTransformOptions) -> Arity { - Arity::Exact(1) + Arity::Variadic { + min: 1, + max: Some(2), + } } fn child_name(&self, _options: &ListTransformOptions, child_idx: usize) -> ChildName { match child_idx { 0 => ChildName::from("list"), + 1 => ChildName::from("environment"), _ => unreachable!("Invalid child index {child_idx} for list_transform()"), } } @@ -90,16 +103,21 @@ impl ScalarFnVTable for ListTransform { expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> fmt::Result { - write!(f, "{}({}, {options})", self.id(), expr.display_child(0)) + write!(f, "{}(", self.id())?; + for index in 0..expr.display_children_count() { + write!(f, "{}, ", expr.display_child(index))?; + } + write!(f, "{options})") } fn return_dtype(&self, options: &ListTransformOptions, args: &[DType]) -> VortexResult { + let (element, environment) = input_dtypes(args)?; let element = match options { ListTransformOptions::Lambda(lambda) => { - lambda.return_dtype(&element_dtype(&args[0], lambda.params())?)? + lambda.return_dtype(&element_param(lambda, element)?, &environment)? } ListTransformOptions::Bound(lambda) => { - check_bound_to(lambda, &args[0])?; + check_bound_to(lambda, element, &environment)?; lambda.dtype().clone() } }; @@ -111,13 +129,13 @@ impl ScalarFnVTable for ListTransform { options: &ListTransformOptions, args: &[DType], ) -> VortexResult> { + let (element, environment) = input_dtypes(args)?; match options { - ListTransformOptions::Lambda(lambda) => { - let element = element_dtype(&args[0], lambda.params())?; - Ok(Some(ListTransformOptions::Bound(lambda.bind(&element)?))) - } + ListTransformOptions::Lambda(lambda) => Ok(Some(ListTransformOptions::Bound( + lambda.bind(&element_param(lambda, element)?, &environment)?, + ))), ListTransformOptions::Bound(lambda) => { - check_bound_to(lambda, &args[0])?; + check_bound_to(lambda, element, &environment)?; Ok(None) } } @@ -132,8 +150,15 @@ impl ScalarFnVTable for ListTransform { let ListTransformOptions::Bound(lambda) = options else { vortex_bail!("list_transform()'s lambda is not bound"); }; + let environment = (args.num_inputs() > 1).then(|| args.get(1)).transpose()?; match args.get(0)?.try_downcast::() { - Ok(list) => transform_list(list, lambda, ctx), + Ok(list) => transform_list(list, lambda, environment.as_ref(), ctx), + // Overlapping views share elements between rows, but each element must see its own + // row's environment, so the views are copied apart into a list first. + Err(list) if environment.is_some() => { + let list = list_from_list_view(list.execute::(ctx)?, ctx)?; + transform_list(list, lambda, environment.as_ref(), ctx) + } Err(list) => transform_list_view(list.execute::(ctx)?, lambda, ctx), } } @@ -150,47 +175,95 @@ impl ScalarFnVTable for ListTransform { } } -/// Check that a bound lambda's parameter has the element dtype of a list of this dtype. -fn check_bound_to(lambda: &BoundLambda, list: &DType) -> VortexResult<()> { - let element = element_dtype(list, lambda.params().names().as_ref())?; +/// Check that a bound lambda's inputs have the dtypes of the element and of the environment's +/// fields. +fn check_bound_to( + lambda: &BoundLambda, + element: DType, + environment: &StructFields, +) -> VortexResult<()> { + let inputs = iter::once(element) + .chain(environment.fields()) + .collect_vec(); vortex_ensure!( - lambda.is_bound_to(&element), - "list_transform()'s lambda {lambda} is bound to element dtype {}, not {}", + lambda.is_bound_to(&inputs), + "list_transform()'s lambda {lambda} is bound to {}, not {}", lambda.params().fields().format(", "), - element[0] + inputs.iter().format(", ") ); Ok(()) } -/// The dtype of the lambda's one parameter, the element, for a list of this dtype. -fn element_dtype(list: &DType, params: &[FieldName]) -> VortexResult<[DType; 1]> { - let DType::List(element, _) = list else { - vortex_bail!("list_transform() takes a List, got {list}"); +/// The dtypes of the lambda's inputs for these children: the list's element, and the fields of the +/// environment, empty without one. +fn input_dtypes(args: &[DType]) -> VortexResult<(DType, StructFields)> { + let DType::List(element, _) = &args[0] else { + vortex_bail!("list_transform() takes a List, got {}", args[0]); }; + let environment = match args.get(1) { + None => StructFields::empty(), + Some(DType::Struct(fields, _)) => fields.clone(), + Some(other) => vortex_bail!("list_transform()'s environment must be a struct, got {other}"), + }; + Ok((element.as_ref().clone(), environment)) +} + +/// The dtype of the lambda's one parameter, the element. +fn element_param(lambda: &Lambda, element: DType) -> VortexResult<[DType; 1]> { vortex_ensure!( - params.len() == 1, + lambda.params().len() == 1, "list_transform() takes a lambda of one parameter, the element, not ({})", - params.iter().format(", ") + lambda.params().iter().format(", ") ); - Ok([element.as_ref().clone()]) + Ok([element]) +} + +/// For each element of lists whose offsets start at 0 and end at the last element, the index of +/// the list that holds it: `[0, 0, 2]` for the offsets `[0, 2, 2, 3]`. +fn list_of_each_element( + offsets: &ArrayRef, + len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let offsets = offsets.clone().execute::(ctx)?; + let mut lists = BufferMut::::with_capacity_in(len, ctx.allocator().clone()); + match_each_integer_ptype!(offsets.ptype(), |O| { + for (list, bounds) in offsets.as_slice::().windows(2).enumerate() { + lists.push_n(list as u64, (bounds[1] - bounds[0]).as_()); + } + }); + Ok(lists.freeze().into_array()) } /// The lambda over a list's elements, inside the list's offsets and validity. /// /// A lambda that can fail runs only on the elements valid lists reference. One that can't runs on /// every element from the first list to the last: results under null lists are never read. +/// `environment` is the lambda's environment, one struct per row, when it reads from outside. fn transform_list( list: ListArray, lambda: &BoundLambda, + environment: Option<&ArrayRef>, ctx: &mut ExecutionCtx, ) -> VortexResult { + // Either way, the offsets start at 0 and end at the last element. let list = if lambda.is_infallible() { trim_ends(list, ctx)? } else { list.compact(ctx)? }; - let values = lambda.evaluate(vec![list.elements().clone()], list.elements().len())?; - vortex_ensure_eq!(values.len(), list.elements().len()); + let elements = list.elements(); + let mut inputs = vec![elements.clone()]; + if let Some(environment) = environment { + // Each row's environment is repeated for every element of its list. + let lists = list_of_each_element(list.offsets(), elements.len(), ctx)?; + let environment = environment.clone().execute::(ctx)?; + for field in environment.iter_unmasked_fields() { + inputs.push(field.take(lists.clone())?); + } + } + let values = lambda.evaluate(inputs, elements.len())?; + vortex_ensure_eq!(values.len(), elements.len()); // SAFETY: the offsets and validity are the list's, over one value per element. Ok( unsafe { ListArray::new_unchecked(values, list.offsets().clone(), list.list_validity()) } diff --git a/vortex-array/src/scalar_fn/fns/list_transform/tests.rs b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs index a93bc64edeb..16671c898d5 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -18,6 +18,7 @@ use crate::IntoArray; use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::FixedSizeListArray; @@ -33,6 +34,7 @@ use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::dtype::StructFields; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::Lambda; @@ -72,6 +74,18 @@ fn rows(l: ArrayRef) -> VortexResult { ) } +/// Rows `{l, a}`. +fn rows_with_a(l: ArrayRef, a: ArrayRef) -> VortexResult { + let len = l.len(); + Ok(StructArray::try_new( + FieldNames::from(["l", "a"]), + vec![l, a], + len, + Validity::NonNullable, + )? + .into_array()) +} + fn list_of(values: ArrayRef, offsets: &[u32], validity: Validity) -> VortexResult { let mut ctx = SESSION.create_execution_ctx(); Ok(ListArray::try_new( @@ -97,6 +111,11 @@ fn plus_one() -> VortexResult { Lambda::new(["x"], checked_add(var("x"), lit(1_i32))) } +/// `x -> x + $.a`, which reads `a` from outside the lambda. +fn plus_a() -> VortexResult { + Lambda::new(["x"], checked_add(var("x"), col("a"))) +} + fn assert_applies(input: ArrayRef, expr: &Expression, expected: ArrayRef) -> VortexResult<()> { let actual = input.apply(expr)?; let mut ctx = SESSION.create_execution_ctx(); @@ -212,12 +231,14 @@ fn a_lambda_that_can_fail_sees_only_the_lists_elements( fn the_lambda_runs_on_a_slice_without_its_leftovers(#[case] body: Expression) -> VortexResult<()> { // Canonicalizing the result drops the leftovers too, so look at the transformed list itself. let mut ctx = SESSION.create_execution_ctx(); - let lambda = Lambda::new(["x"], body)? - .bind(&[DType::Primitive(PType::I32, Nullability::NonNullable)])?; + let lambda = Lambda::new(["x"], body)?.bind( + &[DType::Primitive(PType::I32, Nullability::NonNullable)], + &StructFields::empty(), + )?; let list = sliced_list_with_leftovers()? .try_downcast::() .map_err(|_| vortex_err!("not a list"))?; - let result = transform_list(list, &lambda, &mut ctx)? + let result = transform_list(list, &lambda, None, &mut ctx)? .try_downcast::() .map_err(|_| vortex_err!("not a list"))?; assert_eq!(result.elements().len(), 2); @@ -553,9 +574,7 @@ fn a_lambda_bound_to_another_element_dtype_is_an_error() -> VortexResult<()> { .expect_err("bound to i32, given i32?"); for error in [error, return_dtype_error] { assert!( - error - .to_string() - .contains("is bound to element dtype i32, not i32?"), + error.to_string().contains("is bound to i32, not i32?"), "{error}" ); } @@ -631,3 +650,172 @@ fn lambdas_are_not_serializable_yet() -> VortexResult<()> { assert!(error.to_string().contains("is not serializable"), "{error}"); Ok(()) } + +#[test] +fn a_body_reads_a_column() -> VortexResult<()> { + let expr = list_transform(col("l"), plus_a()?); + assert_eq!( + expr.to_string(), + "vortex.list.transform($.l, pack($: pack(a: $.a)), (x) -> (@x + $.a))" + ); + // Every element of a list sees its row's `a`. The null list's `i32::MAX + 20` would overflow, + // so it must not reach the lambda. + let input = rows_with_a(lists()?, buffer![10_i32, 20, 30].into_array())?; + let expected = transformed( + PrimitiveArray::from_option_iter([Some(11_i32), Some(12), Some(33), None]).into_array(), + )?; + assert_applies(input, &expr, expected) +} + +#[test] +fn a_null_value_read_from_outside_keeps_its_list_valid() -> VortexResult<()> { + let input = rows_with_a( + lists()?, + PrimitiveArray::from_option_iter([None, Some(20_i32), Some(30)]).into_array(), + )?; + let expected = transformed( + PrimitiveArray::from_option_iter([None, None, Some(33_i32), None]).into_array(), + )?; + assert_applies(input, &list_transform(col("l"), plus_a()?), expected) +} + +#[test] +fn a_constant_value_read_from_outside() -> VortexResult<()> { + let input = rows_with_a(lists()?, ConstantArray::new(10_i32, 3).into_array())?; + let expected = transformed( + PrimitiveArray::from_option_iter([Some(11_i32), Some(12), Some(13), None]).into_array(), + )?; + assert_applies(input, &list_transform(col("l"), plus_a()?), expected) +} + +#[test] +fn overlapping_list_views_with_a_value_read_from_outside() -> VortexResult<()> { + // `[[3, 4], [1], [2, 3]]`, sharing element 3, which must see each of its two rows' `a`. + let mut ctx = SESSION.create_execution_ctx(); + let view = ListViewArray::try_new( + buffer![1_i32, 2, 3, 4].into_array(), + buffer![2_u32, 0, 1].into_array(), + buffer![2_u32, 1, 2].into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array(); + let input = rows_with_a(view, buffer![10_i32, 20, 30].into_array())?; + let expected = list_of( + buffer![13_i32, 14, 21, 32, 33].into_array(), + &[0, 2, 3, 5], + Validity::NonNullable, + )?; + assert_applies(input, &list_transform(col("l"), plus_a()?), expected) +} + +#[test] +fn a_nested_lambda_reads_a_column() -> VortexResult<()> { + // `[[[1], [2]], [[3]]]` with `a = [10, 20]` under `xs -> list_transform(xs, y -> y + $.a)`: + // both calls take `$.a` in their environments. + let l = list_of( + list_of( + buffer![1_i32, 2, 3].into_array(), + &[0, 1, 2, 3], + Validity::NonNullable, + )?, + &[0, 2, 3], + Validity::NonNullable, + )?; + let inner = Lambda::new(["y"], checked_add(var("y"), col("a")))?; + let expr = list_transform( + col("l"), + Lambda::new(["xs"], list_transform(var("xs"), inner))?, + ); + assert_eq!( + expr.to_string(), + "vortex.list.transform($.l, pack($: pack(a: $.a)), (xs) -> vortex.list.transform(@xs, \ + pack($: pack(a: $.a)), (y) -> (@y + $.a)))" + ); + let expected = list_of( + list_of( + buffer![11_i32, 12, 23].into_array(), + &[0, 1, 2, 3], + Validity::NonNullable, + )?, + &[0, 2, 3], + Validity::NonNullable, + )?; + assert_applies( + rows_with_a(l, buffer![10_i32, 20].into_array())?, + &expr, + expected, + ) +} + +#[test] +fn a_nested_lambda_reads_the_outer_parameter() -> VortexResult<()> { + // `[[{k: 10, vs: [1, 2]}, {k: 20, vs: [3]}]]` under `p -> list_transform(p.vs, v -> v + p.k)` + let points = StructArray::try_new( + FieldNames::from(["k", "vs"]), + vec![ + buffer![10_i32, 20].into_array(), + list_of( + buffer![1_i32, 2, 3].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?, + ], + 2, + Validity::NonNullable, + )? + .into_array(); + let inner = Lambda::new(["v"], checked_add(var("v"), get_item("k", var("p"))))?; + let expr = list_transform( + root(), + Lambda::new(["p"], list_transform(get_item("vs", var("p")), inner))?, + ); + assert_eq!( + expr.to_string(), + "vortex.list.transform($, (p) -> vortex.list.transform(@p.vs, pack(p: pack(k: @p.k)), \ + (v) -> (@v + @p.k)))" + ); + let expected = list_of( + list_of( + buffer![11_i32, 12, 23].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?, + &[0, 2], + Validity::NonNullable, + )?; + assert_applies( + list_of(points, &[0, 2], Validity::NonNullable)?, + &expr, + expected, + ) +} + +#[test] +fn a_body_reads_a_field_of_a_nullable_struct() -> VortexResult<()> { + // Rows `{l, s}` where the last row's `s` is null, under `x -> x + $.s.i`. + let s = StructArray::try_new( + FieldNames::from(["i"]), + vec![buffer![10_i32, 20, 30].into_array()], + 3, + Validity::from_iter([true, true, false]), + )? + .into_array(); + let input = StructArray::try_new( + FieldNames::from(["l", "s"]), + vec![lists()?, s], + 3, + Validity::NonNullable, + )? + .into_array(); + let lambda = Lambda::new(["x"], checked_add(var("x"), get_item("i", col("s"))))?; + let expr = list_transform(col("l"), lambda); + assert_eq!( + expr.to_string(), + "vortex.list.transform($.l, pack($: pack(s: pack(i: $.s.i))), (x) -> (@x + $.s.i))" + ); + let expected = transformed( + PrimitiveArray::from_option_iter([Some(11_i32), Some(12), None, None]).into_array(), + )?; + assert_applies(input, &expr, expected) +} diff --git a/vortex-expr/src/transform/replace.rs b/vortex-expr/src/transform/replace.rs index 12e408981d2..fc785966047 100644 --- a/vortex-expr/src/transform/replace.rs +++ b/vortex-expr/src/transform/replace.rs @@ -70,4 +70,19 @@ mod test { ); Ok(()) } + + #[test] + fn test_replace_reaches_what_a_lambda_reads_from_outside() -> VortexResult<()> { + // `$.a` is in the call's environment, among its children, so it moves with `$`. The body + // stays as written: its `$.a` reads the environment's row. + let expr = list_transform( + col("l"), + Lambda::new(["x"], checked_add(var("x"), col("a")))?, + ); + assert_eq!( + replace(expr, &root(), col("s")).to_string(), + "vortex.list.transform($.s.l, pack($: pack(a: $.s.a)), (x) -> (@x + $.a))" + ); + Ok(()) + } } diff --git a/vortex-layout/src/layouts/row_idx/expr.rs b/vortex-layout/src/layouts/row_idx/expr.rs index 659eb1f69d7..31dbbd4d47e 100644 --- a/vortex-layout/src/layouts/row_idx/expr.rs +++ b/vortex-layout/src/layouts/row_idx/expr.rs @@ -71,26 +71,3 @@ impl ScalarFnVTable for RowIdx { pub fn row_idx() -> Expression { RowIdx.new_expr(EmptyOptions, []) } - -#[cfg(test)] -mod tests { - use vortex_array::expr::Lambda; - use vortex_array::expr::checked_add; - use vortex_array::expr::var; - - use super::row_idx; - - /// The scan replaces `row_idx()` before anything runs, but never looks inside a lambda's body, - /// so a lambda can't read it. `Lambda::new` recognizes it by its id. - #[test] - fn a_lambda_cannot_read_row_idx() { - let error = Lambda::new(["x"], checked_add(var("x"), row_idx())) - .expect_err("row_idx() is not a parameter"); - assert!( - error - .to_string() - .contains("lambda (x) can only read its parameters, not #row_idx"), - "{error}" - ); - } -} diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 9ed98e3b31a..1b9b5a49b07 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -355,15 +355,21 @@ mod tests { use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ListArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::expr::Lambda; + use vortex_array::expr::checked_add; use vortex_array::expr::eq; use vortex_array::expr::gt; + use vortex_array::expr::list_transform; use vortex_array::expr::lit; use vortex_array::expr::or; use vortex_array::expr::root; + use vortex_array::expr::var; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; @@ -525,6 +531,66 @@ mod tests { }) } + #[test] + fn a_lambda_reads_row_idx() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + // `[[1, 2], [3]]` + let array = ListArray::try_new( + buffer![1u64, 2, 3].into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .unwrap() + .into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&segments), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await + .unwrap(); + + // `row_idx()` is in the call's environment, which the reader fills in like any child. + let lambda = Lambda::new(["x"], checked_add(var("x"), row_idx())).unwrap(); + let expr = list_transform(root(), lambda); + let reader = RowIdxLayoutReader::new( + 0, + layout + .new_reader("".into(), segments, &session, &Default::default()) + .unwrap(), + session.clone(), + ); + let expr = expr.bind(reader.dtype()).unwrap(); + let result = reader + .projection_evaluation( + &(0..layout.row_count()), + &expr, + MaskFuture::new_true(layout.row_count().try_into().unwrap()), + ) + .unwrap() + .await + .unwrap(); + + let expected = ListArray::try_new( + buffer![1u64, 2, 4].into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .unwrap(); + assert_arrays_eq!(result, expected, &mut ctx); + }) + } + #[test] fn row_idx_array_all_true_keeps_sequence() { block_on(|_| async {