From c3198531b1e64260d71782d41a3dc033e48b99b7 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 22 Oct 2026 15:27:53 -0400 Subject: [PATCH] capture values read from outside list_transform lambdas Signed-off-by: Matt Katz --- vortex-array/src/expr/exprs.rs | 16 +- vortex-array/src/expr/lambda/mod.rs | 183 +++++++++++++--- vortex-array/src/expr/lambda/tests.rs | 51 ++++- .../src/scalar_fn/fns/list_transform/mod.rs | 139 ++++++++---- .../src/scalar_fn/fns/list_transform/tests.rs | 200 ++++++++++++++++-- vortex-array/src/scalar_fn/vtable.rs | 4 +- vortex-layout/src/layouts/row_idx/expr.rs | 23 -- vortex-layout/src/layouts/row_idx/mod.rs | 66 ++++++ 8 files changed, 569 insertions(+), 113 deletions(-) diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 4904c3791c3..3a47936bca3 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; @@ -1397,16 +1398,27 @@ 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. Whatever else its body reads from outside, such as +/// a column or a parameter of an enclosing lambda, becomes another input of the call and another +/// parameter of the lambda (see [`Lambda::extract_captures`]): every element of a row's list sees +/// that row's value. /// /// ```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, $.a, (x, a) -> (@x + @a))"); /// ``` pub fn list_transform(list: Expression, lambda: Lambda) -> Expression { - ListTransform.new_expr(ListTransformOptions::Lambda(lambda), [list]) + let (lambda, captures) = lambda.extract_captures(); + ListTransform.new_expr( + ListTransformOptions::Lambda(lambda), + iter::once(list).chain(captures), + ) } /// 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..9e6a30a10d2 100644 --- a/vortex-array/src/expr/lambda/mod.rs +++ b/vortex-array/src/expr/lambda/mod.rs @@ -10,8 +10,13 @@ //! //! 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. +//! struct, such as one row per list element, and applies the bound body to it. +//! +//! A body may read values from outside the lambda, such as a column. Before a function keeps a +//! lambda, [`Lambda::extract_captures`] turns each of those values into a new parameter and returns +//! the expressions that compute them, which the function takes as extra children. A lambda in a +//! function's options therefore reads nothing but its parameters, and the function's children are +//! everything it reads. #[cfg(test)] mod tests; @@ -25,6 +30,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use crate::ArrayRef; use crate::IntoArray; @@ -42,6 +48,8 @@ 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::expr::var; +use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::internal::row_count::RowCount; use crate::validity::Validity; @@ -49,7 +57,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 values from outside the lambda, which +/// [`Lambda::extract_captures`] turns into parameters. /// /// ```rust /// # use vortex_array::expr::{Lambda, checked_add, lit, var}; @@ -65,8 +74,10 @@ 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()`. + /// The body may read from outside the lambda: `$` or a column, another lambda's parameter, or + /// the scan's `row_idx()` or `row_count()`. See [`Lambda::extract_captures`]. + /// + /// Returns an error when a parameter repeats. pub fn new( params: impl IntoIterator>, body: Expression, @@ -75,19 +86,61 @@ 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) - })?; - Ok(Self { params, body }) } + /// Turn every value the body reads from outside the lambda into a new parameter. + /// + /// Returns the lambda with those parameters after its own, and the expressions that compute + /// their values, in the same order. A function that takes a lambda keeps the returned lambda + /// and takes the expressions as extra children, so that a rewrite of the expression, such as + /// `replace` or field pruning, sees them. + /// + /// A value is read as a whole: `$.s.a` becomes one parameter, `a`, not a read of `$`. A value + /// read twice becomes one parameter. A new parameter is named after the field or variable it + /// reads, with a number added when that name is taken. + /// + /// ```rust + /// # use vortex_array::expr::{Lambda, checked_add, col, var}; + /// let lambda = Lambda::new(["x"], checked_add(var("x"), col("a"))).unwrap(); + /// let (lambda, captures) = lambda.extract_captures(); + /// assert_eq!(lambda.to_string(), "(x, a) -> (@x + @a)"); + /// assert_eq!(captures, [col("a")]); + /// ``` + pub fn extract_captures(self) -> (Lambda, Vec) { + let written = self.params; + let mut params = written.iter().cloned().collect_vec(); + let mut captures: Vec = Vec::new(); + let body = self + .body + .transform_down(|node| { + if !is_capture(&node, written.as_ref()) { + return Ok(Transformed::no(node)); + } + let name = match captures.iter().position(|capture| capture == &node) { + Some(index) => params[written.len() + index].clone(), + None => { + let name = unused_name(capture_name(&node), ¶ms); + params.push(name.clone()); + captures.push(node); + name + } + }; + Ok(Transformed { + value: var(name), + changed: true, + order: TraversalOrder::Skip, + }) + }) + .vortex_expect("expression traversal cannot fail") + .into_inner(); + let lambda = Lambda { + params: params.into(), + body, + }; + (lambda, captures) + } + /// The parameter names, in order. pub fn params(&self) -> &[FieldName] { self.params.as_ref() @@ -101,7 +154,7 @@ impl Lambda { /// 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)?); - self.body_over_inputs().return_dtype(&inputs) + self.body_over_inputs()?.return_dtype(&inputs) } /// Bind the body to these parameter dtypes: `@x` becomes `$.x`, a field of the struct of the @@ -109,10 +162,13 @@ impl Lambda { /// /// 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. + /// + /// Returns an error when the body reads from outside the lambda: the function must pass those + /// values in as parameters, with [`Lambda::extract_captures`]. pub fn bind(&self, param_dtypes: &[DType]) -> VortexResult { let params = self.param_fields(param_dtypes)?; let body = self - .body_over_inputs() + .body_over_inputs()? .bind(&inputs_dtype(params.clone()))?; Ok(BoundLambda { params, body }) } @@ -149,12 +205,20 @@ impl Lambda { /// The body over the struct of the lambda's inputs, where `@x` reads the field `x`. /// + /// Returns an error when the body reads from outside the lambda, which the struct can't hold. /// A lambda in the body is a function's options, which this leaves alone: its variables are its /// own parameters. - fn body_over_inputs(&self) -> Expression { - self.body + fn body_over_inputs(&self) -> VortexResult { + Ok(self + .body .clone() .transform_down(|node| { + vortex_ensure!( + !reads_outside_lambda(&node, self.params.as_ref()), + "lambda ({}) can only read its parameters, not {node}; \ + Lambda::extract_captures turns what it reads from outside into parameters", + self.params.iter().format(", ") + ); Ok(match &node { Expression::Variable(name) => Transformed { value: get_item(name.clone(), root()), @@ -163,9 +227,8 @@ impl Lambda { }, _ => Transformed::no(node), }) - }) - .vortex_expect("expression traversal cannot fail") - .into_inner() + })? + .into_inner()) } } @@ -211,6 +274,19 @@ impl BoundLambda { self.params.fields().eq(param_dtypes.iter().cloned()) } + /// The same lambda bound to other parameter dtypes, such as after a rewrite of the function's + /// node made a value it passes in non-nullable. + /// + /// Every node of the body is rebuilt over the new struct of inputs, so the body's dtype follows + /// the parameters'. The result is the same as binding the written lambda to these dtypes. + pub fn rebind(&self, param_dtypes: &[DType]) -> VortexResult { + vortex_ensure_eq!(param_dtypes.len(), self.params.nfields(), "lambda {self}"); + let params = StructFields::new(self.params.names().clone(), param_dtypes.to_vec()); + let root = BoundExpression::new_root(inputs_dtype(params.clone())); + let body = with_root(&self.body, &root)?; + Ok(BoundLambda { params, body }) + } + /// Whether nothing in the body can fail. See [`Lambda::is_infallible`]. pub fn is_infallible(&self) -> bool { let mut infallible = true; @@ -274,10 +350,23 @@ impl Display for BoundLambda { } } +/// `expr` rebuilt over another root, node by node, so that every node's dtype follows the root's. +fn with_root(expr: &BoundExpression, root: &BoundExpression) -> VortexResult { + if expr.as_scalar().is_none() { + return Ok(root.clone()); + } + let children: Vec<_> = expr + .children() + .iter() + .map(|child| with_root(child, root)) + .try_collect()?; + expr.clone().with_children(children) +} + /// 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 +/// must be passed in as parameters. fn reads_outside_lambda(node: &Expression, params: &[FieldName]) -> bool { match node { Expression::Root => true, @@ -288,3 +377,47 @@ fn reads_outside_lambda(node: &Expression, params: &[FieldName]) -> bool { } } } + +/// Whether `node` is a value from outside a lambda with these parameters, read as a whole: a read +/// from outside, or a field of one, such as `$.s.a`. +fn is_capture(node: &Expression, params: &[FieldName]) -> bool { + if node + .as_scalar() + .is_some_and(|scalar_fn| scalar_fn.is::()) + { + return is_capture(&node.children()[0], params); + } + reads_outside_lambda(node, params) +} + +/// The name of the parameter for a captured value: the field or variable it reads, such as `a` +/// for `$.s.a`. +fn capture_name(capture: &Expression) -> FieldName { + match capture { + Expression::Root => FieldName::from("root"), + Expression::Variable(name) => name.clone(), + Expression::Scalar { scalar_fn, .. } => match scalar_fn.as_opt::() { + Some(field) => field.clone(), + // `row_idx()` or `row_count()`, named after their ids. + None => { + let id = scalar_fn.id(); + let (_, name) = id + .as_ref() + .rsplit_once('.') + .unwrap_or_else(|| ("", id.as_ref())); + FieldName::from(name) + } + }, + } +} + +/// `name`, or `name_2`, `name_3`, ... if it is taken. +fn unused_name(name: FieldName, taken: &[FieldName]) -> FieldName { + if !taken.contains(&name) { + return name; + } + (2..) + .map(|n| FieldName::from(format!("{name}_{n}"))) + .find(|candidate| !taken.contains(candidate)) + .vortex_expect("some numbered name is free") +} diff --git a/vortex-array/src/expr/lambda/tests.rs b/vortex-array/src/expr/lambda/tests.rs index d09592bdf11..f5940bb6806 100644 --- a/vortex-array/src/expr/lambda/tests.rs +++ b/vortex-array/src/expr/lambda/tests.rs @@ -168,19 +168,57 @@ fn new_keeps_params_and_body() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::column(checked_add(var("x"), col("a")), "(x, a) -> (@x + @a)", vec![col("a")])] +#[case::field_of_a_column( + checked_add(var("x"), get_item("i", col("s"))), + "(x, i) -> (@x + @i)", + vec![get_item("i", col("s"))] +)] +#[case::read_twice(checked_add(col("a"), col("a")), "(x, a) -> (@a + @a)", vec![col("a")])] +#[case::name_taken(checked_add(var("x"), col("x")), "(x, x_2) -> (@x + @x_2)", vec![col("x")])] +#[case::root(eq(var("x"), root()), "(x, root) -> (@x = @root)", vec![root()])] +#[case::other_variable( + checked_add(var("x"), get_item("k", var("p"))), + "(x, k) -> (@x + @k)", + vec![get_item("k", var("p"))] +)] +#[case::function_of_the_scan( + checked_add(var("x"), RowCount.new_expr(EmptyOptions, [])), + "(x, row_count) -> (@x + @row_count)", + vec![RowCount.new_expr(EmptyOptions, [])] +)] +#[case::field_of_a_parameter(get_item("k", var("x")), "(x) -> @x.k", vec![])] +fn extract_captures_turns_reads_from_outside_into_parameters( + #[case] body: Expression, + #[case] lambda: &str, + #[case] captures: Vec, +) -> VortexResult<()> { + let (extracted, extracted_captures) = Lambda::new(["x"], body)?.extract_captures(); + assert_eq!(extracted.to_string(), lambda); + assert_eq!(extracted_captures, captures); + Ok(()) +} + #[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"); +fn binding_a_body_that_reads_from_outside_is_an_error( + #[case] body: Expression, + #[case] read: &str, +) -> VortexResult<()> { + let error = Lambda::new(["x"], body)? + .bind(&[i32_()]) + .expect_err("reads from outside"); let message = error.to_string(); assert!( message.contains("lambda (x) can only read its parameters") && message.contains(read), "{message}" ); + Ok(()) } #[test] @@ -295,6 +333,15 @@ fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { Ok(()) } +#[test] +fn rebinding_is_binding_to_the_new_dtypes() -> VortexResult<()> { + let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?; + let rebound = lambda.bind(&[i32_()])?.rebind(&[i32_().as_nullable()])?; + assert_eq!(rebound, lambda.bind(&[i32_().as_nullable()])?); + assert_eq!(rebound.dtype(), &i32_().as_nullable()); + Ok(()) +} + #[test] fn binding_an_expression_binds_its_lambdas() -> VortexResult<()> { let expr = apply(col("a"), "x", checked_add(var("x"), lit(1_i32)))?; 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..24bdd358f87 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,19 @@ use crate::IntoArray; use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; 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::dtype::DType; use crate::dtype::FieldName; 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,9 +46,10 @@ 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 first child is the list, of `List` dtype in any encoding, and the result is a `List`. The +/// lambda's first parameter is the element. Each other child is a value the lambda reads from +/// outside, one per row, and is the lambda's next parameter: every element of a row's list sees +/// that row's value. A lambda that can fail runs only on the elements of valid lists. #[derive(Clone)] pub struct ListTransform; @@ -56,6 +63,16 @@ pub enum ListTransformOptions { Bound(BoundLambda), } +impl ListTransformOptions { + /// The lambda's parameter names: the element, then one per value it reads from outside. + fn params(&self) -> &[FieldName] { + match self { + Self::Lambda(lambda) => lambda.params(), + Self::Bound(lambda) => lambda.params().names().as_ref(), + } + } +} + impl Display for ListTransformOptions { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { @@ -74,13 +91,13 @@ impl ScalarFnVTable for ListTransform { } fn arity(&self, _options: &ListTransformOptions) -> Arity { - Arity::Exact(1) + Arity::Variadic { min: 1, max: None } } - fn child_name(&self, _options: &ListTransformOptions, child_idx: usize) -> ChildName { + fn child_name(&self, options: &ListTransformOptions, child_idx: usize) -> ChildName { match child_idx { 0 => ChildName::from("list"), - _ => unreachable!("Invalid child index {child_idx} for list_transform()"), + _ => ChildName::from(Arc::clone(options.params()[child_idx].inner())), } } @@ -90,18 +107,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 params = param_dtypes(options.params(), args)?; let element = match options { - ListTransformOptions::Lambda(lambda) => { - lambda.return_dtype(&element_dtype(&args[0], lambda.params())?)? - } - ListTransformOptions::Bound(lambda) => { - check_bound_to(lambda, &args[0])?; + ListTransformOptions::Lambda(lambda) => lambda.return_dtype(¶ms)?, + ListTransformOptions::Bound(lambda) if lambda.is_bound_to(¶ms) => { lambda.dtype().clone() } + ListTransformOptions::Bound(lambda) => lambda.rebind(¶ms)?.dtype().clone(), }; Ok(DType::List(Arc::new(element), args[0].nullability())) } @@ -111,16 +131,18 @@ impl ScalarFnVTable for ListTransform { options: &ListTransformOptions, args: &[DType], ) -> VortexResult> { - match options { + let params = param_dtypes(options.params(), args)?; + Ok(match options { ListTransformOptions::Lambda(lambda) => { - let element = element_dtype(&args[0], lambda.params())?; - Ok(Some(ListTransformOptions::Bound(lambda.bind(&element)?))) + Some(ListTransformOptions::Bound(lambda.bind(¶ms)?)) } + ListTransformOptions::Bound(lambda) if lambda.is_bound_to(¶ms) => None, + // A rewrite changed the dtype of the list or of a value the lambda reads, such as by + // making it non-nullable. ListTransformOptions::Bound(lambda) => { - check_bound_to(lambda, &args[0])?; - Ok(None) + Some(ListTransformOptions::Bound(lambda.rebind(¶ms)?)) } - } + }) } fn execute( @@ -132,14 +154,25 @@ impl ScalarFnVTable for ListTransform { let ListTransformOptions::Bound(lambda) = options else { vortex_bail!("list_transform()'s lambda is not bound"); }; + let captures = (1..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>>()?; match args.get(0)?.try_downcast::() { - Ok(list) => transform_list(list, lambda, ctx), + Ok(list) => transform_list(list, lambda, &captures, ctx), + // Overlapping views share elements between rows, but each element must see the value + // its own row reads from outside, so the views are copied apart into a list first. + Err(list) if !captures.is_empty() => { + let list = list_from_list_view(list.execute::(ctx)?, ctx)?; + transform_list(list, lambda, &captures, ctx) + } Err(list) => transform_list_view(list.execute::(ctx)?, lambda, ctx), } } - fn is_strict(&self, _options: &ListTransformOptions) -> bool { - true + /// Strict only when the lambda reads nothing from outside: a null value read from outside + /// doesn't make a valid list null. + fn is_strict(&self, options: &ListTransformOptions) -> bool { + options.params().len() == 1 } fn is_infallible(&self, options: &ListTransformOptions) -> bool { @@ -150,47 +183,69 @@ 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())?; - vortex_ensure!( - lambda.is_bound_to(&element), - "list_transform()'s lambda {lambda} is bound to element dtype {}, not {}", - lambda.params().fields().format(", "), - element[0] - ); - 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 parameters for these children: the list's element, then the values +/// it reads from outside. +fn param_dtypes(params: &[FieldName], args: &[DType]) -> VortexResult> { + let DType::List(element, _) = &args[0] else { + vortex_bail!("list_transform() takes a List, got {}", args[0]); }; + // The parameters for values read from outside come last, so the rest are the written ones. + let written = ¶ms[..params.len().saturating_sub(args.len() - 1)]; vortex_ensure!( - params.len() == 1, + params.len() == args.len(), "list_transform() takes a lambda of one parameter, the element, not ({})", - params.iter().format(", ") + written.iter().format(", ") ); - Ok([element.as_ref().clone()]) + Ok(iter::once(element.as_ref().clone()) + .chain(args[1..].iter().cloned()) + .collect()) +} + +/// 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(len); + 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. +/// `captures` are the values the lambda reads from outside, one per row. fn transform_list( list: ListArray, lambda: &BoundLambda, + captures: &[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 !captures.is_empty() { + // Each value read from outside is repeated for every element of its row's list. + let lists = list_of_each_element(list.offsets(), elements.len(), ctx)?; + for capture in captures { + inputs.push(capture.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..54ace9c32a4 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; @@ -72,6 +73,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 +110,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(); @@ -217,7 +235,7 @@ fn the_lambda_runs_on_a_slice_without_its_leftovers(#[case] body: Expression) -> 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, &[], &mut ctx)? .try_downcast::() .map_err(|_| vortex_err!("not a list"))?; assert_eq!(result.elements().len(), 2); @@ -535,7 +553,7 @@ fn applying_the_bound_node() -> VortexResult<()> { } #[test] -fn a_lambda_bound_to_another_element_dtype_is_an_error() -> VortexResult<()> { +fn a_bound_lambda_rebinds_when_the_element_dtype_changes() -> VortexResult<()> { let list = |elements| { DType::List( Arc::new(DType::Primitive(PType::I32, elements)), @@ -543,22 +561,170 @@ fn a_lambda_bound_to_another_element_dtype_is_an_error() -> VortexResult<()> { ) }; let bound = list_transform(root(), plus_one()?).bind(&list(Nullability::NonNullable))?; - let error = bound + let rebuilt = bound .clone() - .with_children([BoundExpression::new_root(list(Nullability::Nullable))]) - .expect_err("bound to i32, rebuilt over i32?"); - let options = bound.as_::(); - let return_dtype_error = ListTransform - .return_dtype(options, &[list(Nullability::Nullable)]) - .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}" - ); - } + .with_children([BoundExpression::new_root(list(Nullability::Nullable))])?; + assert_eq!( + rebuilt, + list_transform(root(), plus_one()?).bind(&list(Nullability::Nullable))? + ); + assert_eq!( + ListTransform.return_dtype(bound.as_::(), &[list(Nullability::Nullable)])?, + list(Nullability::Nullable) + ); + 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, $.a, (x, a) -> (@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` as an input. + 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, $.a, (xs, a) -> vortex.list.transform(@xs, @a, (y, a) -> (@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, @p.k, (v, k) -> (@v + @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_value_read_from_outside_that_became_non_nullable() -> VortexResult<()> { + // Bound where `a` is nullable, applied where it isn't, as after a rewrite: the lambda is + // rebound to the new dtype. + let a = PrimitiveArray::from_option_iter([Some(10_i32), Some(20), Some(30)]).into_array(); + let bound = list_transform(col("l"), plus_a()?).bind(rows_with_a(lists()?, a)?.dtype())?; + 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(), + )?; + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(input.apply_bound(&bound)?, expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 487404c420a..e47405a1b5a 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -165,8 +165,8 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// [`ScalarFnArray`] is built from its children, before /// [`return_dtype`](Self::return_dtype), so [`execute`](Self::execute) finds the options bound. /// A rebuilt node calls it again, possibly with other dtypes: options that are already bound - /// must be checked against these dtypes, returning `None` if they fit and an error if they - /// don't. Return `None` to keep the options as they are, which the default does. + /// must be checked against these dtypes, returning `None` if they fit, and rebound or an error + /// if they don't. Return `None` to keep the options as they are, which the default does. #[doc(hidden)] fn bind_options( &self, 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 f9641cbd6cb..752b8bd5190 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -354,15 +354,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; @@ -524,6 +530,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()` becomes an input of the call, which the reader fills in like any other. + 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 {