From a4e9fe24109e8b114fb5ec84c064d93169770a19 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 14 Oct 2026 16:57:32 -0400 Subject: [PATCH 1/2] list_transform scalar fn Signed-off-by: Matt Katz --- vortex-array/src/expr/exprs.rs | 22 +- vortex-array/src/expr/mod.rs | 1 + .../src/scalar_fn/fns/list_transform/mod.rs | 221 +++++++ .../src/scalar_fn/fns/list_transform/tests.rs | 543 ++++++++++++++++++ vortex-array/src/scalar_fn/fns/mod.rs | 1 + 5 files changed, 786 insertions(+), 2 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/list_transform/mod.rs create mode 100644 vortex-array/src/scalar_fn/fns/list_transform/tests.rs diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index d77a3472d30..3203d89725e 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -18,6 +18,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::Lambda; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::scalar_fn::EmptyOptions; @@ -45,6 +46,8 @@ use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_contains::ListContainsOptions; use crate::scalar_fn::fns::list_length::ListLength; use crate::scalar_fn::fns::list_sum::ListSum; +use crate::scalar_fn::fns::list_transform::ListTransform; +use crate::scalar_fn::fns::list_transform::ListTransformOptions; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::DuplicateHandling; @@ -153,8 +156,7 @@ pub fn bound_get_item(field: impl Into, child: BoundExpression) -> Bo /// Creates an expression that reads the parameter `name` of the enclosing lambda. /// -/// It is only meaningful in the body of a [`Lambda`](crate::expr::Lambda). Anywhere else it fails -/// to bind. +/// It is only meaningful in the body of a [`Lambda`]. Anywhere else it fails to bind. /// /// ```rust /// # use vortex_array::dtype::{DType, Nullability, PType}; @@ -1382,6 +1384,22 @@ pub fn bound_list_sum_opts( .vortex_expect("list-sum expressions require a numeric list child") } +// ---- ListTransform ---- + +/// Creates an expression that applies a lambda to every element of each list. +/// +/// The lambda takes one parameter, the element. +/// +/// ```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))"); +/// ``` +pub fn list_transform(list: Expression, lambda: Lambda) -> Expression { + ListTransform.new_expr(ListTransformOptions::Lambda(lambda), [list]) +} + /// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit. /// /// `ambiguous` is a UTF-8 expression containing `raise`, `earliest`, `latest`, or `null`. diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 230dae1f25d..d848c166344 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -92,6 +92,7 @@ pub use exprs::list_contains_opts; pub use exprs::list_length; pub use exprs::list_sum; pub use exprs::list_sum_opts; +pub use exprs::list_transform; pub use exprs::lit; pub use exprs::lt; pub use exprs::lt_eq; diff --git a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs new file mode 100644 index 00000000000..ff50d3f5701 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -0,0 +1,221 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +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::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; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::validity::Validity; + +/// 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`. The lambda 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. +#[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 + /// [`BoundExpression`](crate::expr::BoundExpression) or an array. + Bound(BoundLambda), +} + +impl Display for ListTransformOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Lambda(lambda) => lambda.fmt(f), + Self::Bound(lambda) => lambda.fmt(f), + } + } +} + +impl ScalarFnVTable for ListTransform { + type Options = ListTransformOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.transform"); + *ID + } + + fn arity(&self, _options: &ListTransformOptions) -> Arity { + Arity::Exact(1) + } + + 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()"), + } + } + + fn fmt_sql( + &self, + options: &ListTransformOptions, + expr: &dyn ExprDisplay, + f: &mut Formatter<'_>, + ) -> fmt::Result { + write!(f, "{}({}, {options})", self.id(), expr.display_child(0)) + } + + fn return_dtype(&self, options: &ListTransformOptions, args: &[DType]) -> VortexResult { + 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])?; + lambda.dtype().clone() + } + }; + Ok(DType::List(Arc::new(element), args[0].nullability())) + } + + fn bind_options( + &self, + options: &ListTransformOptions, + args: &[DType], + ) -> VortexResult> { + match options { + ListTransformOptions::Lambda(lambda) => { + let element = element_dtype(&args[0], lambda.params())?; + Ok(Some(ListTransformOptions::Bound(lambda.bind(&element)?))) + } + ListTransformOptions::Bound(lambda) => { + check_bound_to(lambda, &args[0])?; + Ok(None) + } + } + } + + fn execute( + &self, + options: &ListTransformOptions, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let ListTransformOptions::Bound(lambda) = options else { + vortex_bail!("list_transform()'s lambda is not bound"); + }; + // Every element belongs to exactly one valid list, in order: nothing comes before the + // first list or after the last, and a null list has no elements. + let list = args.get(0)?.execute::(ctx)?; + let list = list_from_list_view(empty_null_lists(list, ctx)?, ctx)?; + let elements = list.elements(); + let values = lambda.evaluate(vec![elements.clone()], elements.len())?; + vortex_ensure_eq!(values.len(), elements.len()); + // SAFETY: the offsets are the list's, over one value per element. + Ok(unsafe { + ListArray::new_unchecked(values, list.offsets().clone(), list.list_validity()) + } + .into_array()) + } + + fn is_strict(&self, _options: &ListTransformOptions) -> bool { + true + } + + fn is_infallible(&self, options: &ListTransformOptions) -> bool { + match options { + ListTransformOptions::Lambda(lambda) => lambda.is_infallible(), + ListTransformOptions::Bound(lambda) => lambda.is_infallible(), + } + } +} + +/// 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}"); + }; + vortex_ensure!( + params.len() == 1, + "list_transform() takes a lambda of one parameter, the element, not ({})", + params.iter().format(", ") + ); + Ok([element.as_ref().clone()]) +} + +/// The lists with every null list emptied, so that nothing under a null list reaches the lambda. +fn empty_null_lists(list: ListViewArray, ctx: &mut ExecutionCtx) -> VortexResult { + let validity = list.listview_validity(); + if matches!(validity, Validity::NonNullable | Validity::AllValid) { + return Ok(list); + } + let valid = validity.execute_mask(list.len(), ctx)?; + let sizes = list.sizes().clone().execute::(ctx)?; + let sizes = match_each_integer_ptype!(sizes.ptype(), |T| { + let sizes = sizes.as_slice::(); + let null = |index: usize| !valid.value(index); + // Writers usually leave null lists empty already. + if !(0..sizes.len()).any(|index| null(index) && sizes[index] != T::default()) { + return Ok(list); + } + PrimitiveArray::from_iter((0..sizes.len()).map(|index| { + if null(index) { + T::default() + } else { + sizes[index] + } + })) + }); + // SAFETY: only sizes changed, and only to 0, so every view stays within the elements. The new + // array is not marked zero-copy to a list: an emptied list leaves its elements behind as a gap, + // which the rebuild must drop. + Ok(unsafe { + ListViewArray::new_unchecked( + list.elements().clone(), + list.offsets().clone(), + sizes.into_array(), + 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 new file mode 100644 index 00000000000..2f7b7268704 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -0,0 +1,543 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::ListTransform; +use super::ListTransformOptions; +use crate::ArrayRef; +use crate::IntoArray; +use crate::RecursiveCanonical; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::Dict; +use crate::arrays::DictArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::BoundExpression; +use crate::expr::Expression; +use crate::expr::Lambda; +use crate::expr::checked_add; +use crate::expr::col; +use crate::expr::eq; +use crate::expr::get_item; +use crate::expr::is_not_null; +use crate::expr::list_contains; +use crate::expr::list_transform; +use crate::expr::lit; +use crate::expr::proto::ExprSerializeProtoExt; +use crate::expr::root; +use crate::expr::transform::replace; +use crate::expr::var; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTable; +use crate::validity::Validity; + +/// `[[1, 2], null (hiding [i32::MAX]), [3, null]]` +fn lists() -> VortexResult { + list_of( + PrimitiveArray::from_option_iter([Some(1_i32), Some(2), Some(i32::MAX), Some(3), None]) + .into_array(), + &[0, 2, 3, 5], + Validity::from_iter([true, false, true]), + ) +} + +/// Rows `{l}`. +fn rows(l: ArrayRef) -> VortexResult { + let len = l.len(); + Ok( + StructArray::try_new(FieldNames::from(["l"]), vec![l], len, Validity::NonNullable)? + .into_array(), + ) +} + +fn list_of(values: ArrayRef, offsets: &[u32], validity: Validity) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + Ok(ListArray::try_new( + values, + PrimitiveArray::from_iter(offsets.iter().copied()).into_array(), + validity, + &mut ctx, + )? + .into_array()) +} + +/// The lists of [`lists`] with new elements: the null list keeps no elements. +fn transformed(values: ArrayRef) -> VortexResult { + list_of( + values, + &[0, 2, 2, 4], + Validity::from_iter([true, false, true]), + ) +} + +/// `x -> x + 1`, which fails on `i32::MAX`. +fn plus_one() -> VortexResult { + Lambda::new(["x"], checked_add(var("x"), lit(1_i32))) +} + +fn assert_applies(input: ArrayRef, expr: &Expression, expected: ArrayRef) -> VortexResult<()> { + let actual = input.apply(expr)?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn the_lambda_runs_on_every_element() -> VortexResult<()> { + let expr = list_transform(col("l"), plus_one()?); + assert_eq!( + expr.to_string(), + "vortex.list.transform($.l, (x) -> (@x + 1i32))" + ); + // The null list's `i32::MAX` would overflow, so it must not reach the lambda. + let expected = transformed( + PrimitiveArray::from_option_iter([Some(2_i32), Some(3), Some(4), None]).into_array(), + )?; + assert_applies(rows(lists()?)?, &expr, expected) +} + +#[test] +fn a_body_of_another_dtype() -> VortexResult<()> { + let expr = list_transform(root(), Lambda::new(["x"], eq(var("x"), lit(3_i32)))?); + let expected = transformed( + BoolArray::from_iter([Some(false), Some(false), Some(true), None]).into_array(), + )?; + assert_applies(lists()?, &expr, expected) +} + +#[test] +fn a_null_list_without_elements() -> VortexResult<()> { + // `[[1, 2], null, [3]]` with nothing under the null list. + let validity = || Validity::from_iter([true, false, true]); + let lists = list_of(buffer![1_i32, 2, 3].into_array(), &[0, 2, 2, 3], validity())?; + let expected = list_of(buffer![2_i32, 3, 4].into_array(), &[0, 2, 2, 3], validity())?; + assert_applies(lists, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn a_sliced_list_skips_a_null_lists_elements() -> VortexResult<()> { + // `[[1, 2], null (hiding [i32::MAX]), [3], [4]]` sliced to `[null, [3], [4]]`, whose offsets + // no longer start at 0. + let lists = list_of( + buffer![1_i32, 2, i32::MAX, 3, 4].into_array(), + &[0, 2, 3, 4, 5], + Validity::from_iter([true, false, true, true]), + )? + .slice(1..4)?; + let expected = list_of( + buffer![4_i32, 5].into_array(), + &[0, 0, 1, 2], + Validity::from_iter([false, true, true]), + )?; + assert_applies(lists, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn every_list_null() -> VortexResult<()> { + // `[null (hiding [i32::MAX]), null (hiding [i32::MAX])]` + let lists = list_of( + buffer![i32::MAX, i32::MAX].into_array(), + &[0, 1, 2], + Validity::from_iter([false, false]), + )?; + let expected = list_of( + PrimitiveArray::from_iter::<[i32; 0]>([]).into_array(), + &[0, 0, 0], + Validity::from_iter([false, false]), + )?; + assert_applies(lists, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn a_failure_in_a_valid_list_is_an_error() -> VortexResult<()> { + let lists = list_of( + buffer![1_i32, i32::MAX].into_array(), + &[0, 1, 2], + Validity::from_iter([false, true]), + )?; + let mut ctx = array_session().create_execution_ctx(); + let error = lists + .apply(&list_transform(root(), plus_one()?))? + .execute::(&mut ctx) + .map(drop) + .expect_err("i32::MAX + 1 overflows in a valid list"); + assert!(error.to_string().contains("overflow"), "{error}"); + Ok(()) +} + +#[test] +fn a_lambda_takes_one_parameter() -> VortexResult<()> { + let expr = list_transform(root(), Lambda::new(["x", "i"], var("x"))?); + let error = expr.bind(lists()?.dtype()).expect_err("two parameters"); + assert!( + error + .to_string() + .contains("takes a lambda of one parameter"), + "{error}" + ); + Ok(()) +} + +#[test] +fn struct_elements_read_their_fields() -> VortexResult<()> { + // `[[{x: 1}, {x: 2}], [{x: 3}]]` under `p -> p.x + 1` + let points = StructArray::try_new( + FieldNames::from(["x"]), + vec![buffer![1_i32, 2, 3].into_array()], + 3, + Validity::NonNullable, + )? + .into_array(); + let expr = list_transform( + root(), + Lambda::new(["p"], checked_add(get_item("x", var("p")), lit(1_i32)))?, + ); + let expected = list_of( + buffer![2_i32, 3, 4].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + assert_applies( + list_of(points, &[0, 2, 3], Validity::NonNullable)?, + &expr, + expected, + ) +} + +#[test] +fn a_body_can_test_membership_in_a_literal_set() -> VortexResult<()> { + // `[[1, 2], [3]]` under `x -> list_contains([2, 3], x)` + let set = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![2_i32.into(), 3_i32.into()], + Nullability::NonNullable, + ); + let expr = list_transform( + root(), + Lambda::new(["x"], list_contains(lit(set), var("x")))?, + ); + let lists = list_of( + buffer![1_i32, 2, 3].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + let expected = list_of( + BoolArray::from_iter([false, true, true]).into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + assert_applies(lists, &expr, expected) +} + +#[test] +fn a_nested_list_result_over_null_lists() -> VortexResult<()> { + // `[[[1], [2]], null, [[3]]]` under `xs -> list_transform(xs, y -> y + 1)`: the body returns + // a list, and a null list must not reach the inner body. + let inner = list_of( + buffer![1_i32, 2, i32::MAX, 3].into_array(), + &[0, 1, 2, 3, 4], + Validity::NonNullable, + )?; + let outer = list_of( + inner, + &[0, 2, 3, 4], + Validity::from_iter([true, false, true]), + )?; + let expr = list_transform( + root(), + Lambda::new(["xs"], list_transform(var("xs"), plus_one()?))?, + ); + assert_eq!( + expr.to_string(), + "vortex.list.transform($, (xs) -> vortex.list.transform(@xs, (x) -> (@x + 1i32)))" + ); + let inner_expected = list_of( + buffer![2_i32, 3, 4].into_array(), + &[0, 1, 2, 3], + Validity::NonNullable, + )?; + let expected = list_of( + inner_expected, + &[0, 2, 2, 3], + Validity::from_iter([true, false, true]), + )?; + assert_applies(outer, &expr, expected) +} + +#[test] +fn three_nested_lambdas() -> VortexResult<()> { + // `[[[[1, 2], [3]]], [[[4]]]]` under + // `x -> list_transform(x, y -> list_transform(y, z -> z + 1))` + let l = list_of( + list_of( + list_of( + buffer![1_i32, 2, 3, 4].into_array(), + &[0, 2, 3, 4], + Validity::NonNullable, + )?, + &[0, 2, 3], + Validity::NonNullable, + )?, + &[0, 1, 2], + Validity::NonNullable, + )?; + let innermost = Lambda::new(["z"], checked_add(var("z"), lit(1_i32)))?; + let middle = Lambda::new(["y"], list_transform(var("y"), innermost))?; + let expr = list_transform( + root(), + Lambda::new(["x"], list_transform(var("x"), middle))?, + ); + let expected = list_of( + list_of( + list_of( + buffer![2_i32, 3, 4, 5].into_array(), + &[0, 2, 3, 4], + Validity::NonNullable, + )?, + &[0, 2, 3], + Validity::NonNullable, + )?, + &[0, 1, 2], + Validity::NonNullable, + )?; + assert_applies(l, &expr, expected) +} + +#[test] +fn sliced_list_views() -> VortexResult<()> { + // `[[1, 2], [3], [4, 5]]` as a list view, sliced to `[[3], [4, 5]]`, which leaves unused + // elements. + let mut ctx = array_session().create_execution_ctx(); + let view = list_of( + buffer![1_i32, 2, 3, 4, 5].into_array(), + &[0, 2, 3, 5], + Validity::NonNullable, + )? + .execute::(&mut ctx)? + .into_array() + .slice(1..3)?; + let expected = list_of( + buffer![4_i32, 5, 6].into_array(), + &[0, 1, 3], + Validity::NonNullable, + )?; + assert_applies(view, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn overlapping_list_views() -> VortexResult<()> { + // `[[3, 4], [1], [2, 3]]`, sharing element 3. + let mut ctx = array_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 expr = list_transform(root(), Lambda::new(["x"], eq(var("x"), lit(3_i32)))?); + let expected = list_of( + BoolArray::from_iter([true, false, false, false, true]).into_array(), + &[0, 2, 3, 5], + Validity::NonNullable, + )?; + assert_applies(view, &expr, expected) +} + +#[test] +fn fixed_size_lists_are_not_supported_yet() -> VortexResult<()> { + let fsl = FixedSizeListArray::try_new( + buffer![1_i32, 2, 3, 4].into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let error = list_transform(root(), plus_one()?) + .bind(fsl.dtype()) + .expect_err("a fixed-size list is not a list"); + assert!( + error.to_string().contains("list_transform() takes a List"), + "{error}" + ); + Ok(()) +} + +#[test] +fn a_constant_body() -> VortexResult<()> { + let expr = list_transform(root(), Lambda::new(["x"], lit(7_i32))?); + assert_applies( + lists()?, + &expr, + transformed(buffer![7_i32, 7, 7, 7].into_array())?, + ) +} + +#[test] +fn no_lists() -> VortexResult<()> { + let empty = list_of( + PrimitiveArray::from_iter::<[i32; 0]>([]).into_array(), + &[0], + Validity::NonNullable, + )?; + let expected = list_of( + PrimitiveArray::from_iter::<[i32; 0]>([]).into_array(), + &[0], + Validity::NonNullable, + )?; + assert_applies(empty, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn the_bound_node_is_an_ordinary_call() -> VortexResult<()> { + let dtype = rows(lists()?)?.dtype().clone(); + let expr = list_transform(col("l"), plus_one()?); + let bound = expr.bind(&dtype)?; + // The bound lambda reads the struct of its inputs: `@x` became `$.x`. + assert_eq!( + bound.to_string(), + "vortex.list.transform($.l, (x) -> ($.x + 1i32))" + ); + assert_eq!(bound.children(), [col("l").bind(&dtype)?]); + let nullable_i32 = DType::Primitive(PType::I32, Nullability::Nullable); + assert_eq!( + bound.dtype(), + &DType::List(Arc::new(nullable_i32.clone()), Nullability::Nullable) + ); + assert_eq!(&expr.return_dtype(&dtype)?, bound.dtype()); + assert_eq!(bound.validity()?, is_not_null(col("l")).bind(&dtype)?); + // Binding the node bound its lambda to the element's dtype. + assert!(matches!( + bound.as_::(), + ListTransformOptions::Bound(lambda) if lambda.is_bound_to(&[nullable_i32]) + )); + Ok(()) +} + +#[test] +fn applying_the_bound_node() -> VortexResult<()> { + // A scan binds the expression once and applies the bound form to every batch. + let input = rows(lists()?)?; + let bound = list_transform(col("l"), plus_one()?).bind(input.dtype())?; + let expected = transformed( + PrimitiveArray::from_option_iter([Some(2_i32), Some(3), Some(4), None]).into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(input.apply_bound(&bound)?, expected, &mut ctx); + Ok(()) +} + +#[test] +fn a_lambda_bound_to_another_element_dtype_is_an_error() -> VortexResult<()> { + let list = |elements| { + DType::List( + Arc::new(DType::Primitive(PType::I32, elements)), + Nullability::NonNullable, + ) + }; + let bound = list_transform(root(), plus_one()?).bind(&list(Nullability::NonNullable))?; + let error = 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}" + ); + } + Ok(()) +} + +#[test] +fn replacing_the_root_rereads_the_list() -> VortexResult<()> { + // Rows `{s: {l}}`: re-rooting the transform at `$.s` changes where the list is read. The body + // is in the options, which rewrites of the expression leave alone. + let input = StructArray::try_new( + FieldNames::from(["s"]), + vec![rows(lists()?)?], + 3, + Validity::NonNullable, + )? + .into_array(); + let expr = list_transform(col("l"), plus_one()?); + assert_eq!(replace(expr.clone(), &var("x"), lit(0_i32)), expr); + let expr = replace(expr, &root(), col("s")); + assert_eq!( + expr.to_string(), + "vortex.list.transform($.s.l, (x) -> (@x + 1i32))" + ); + let expected = transformed( + PrimitiveArray::from_option_iter([Some(2_i32), Some(3), Some(4), None]).into_array(), + )?; + assert_applies(input, &expr, expected) +} + +#[test] +fn an_infallible_transform_pushes_into_dictionary_values() -> VortexResult<()> { + let values = list_of( + buffer![1_i32, 2, 3].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + let dict = DictArray::try_new(buffer![0_u8, 1, 0, 1].into_array(), values)?.into_array(); + let expr = list_transform(root(), Lambda::new(["x"], eq(var("x"), lit(1_i32)))?); + + let actual = dict.apply(&expr)?; + assert!(actual.is::(), "{actual}"); + let expected = list_of( + BoolArray::from_iter([true, false, false, true, false, false]).into_array(), + &[0, 2, 3, 5, 6], + Validity::NonNullable, + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn a_fallible_transform_skips_unused_dictionary_values() -> VortexResult<()> { + // Dictionary values `[[1], [i32::MAX]]`, of which the codes use only the first. + let values = list_of( + buffer![1_i32, i32::MAX].into_array(), + &[0, 1, 2], + Validity::NonNullable, + )?; + let dict = DictArray::try_new(buffer![0_u8, 0].into_array(), values)?.into_array(); + let expected = list_of( + buffer![2_i32, 2].into_array(), + &[0, 1, 2], + Validity::NonNullable, + )?; + assert_applies(dict, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn lambdas_are_not_serializable_yet() -> VortexResult<()> { + let expr = list_transform(col("l"), Lambda::new(["x"], var("x"))?); + let error = expr + .serialize_proto() + .expect_err("lambdas are not serializable"); + assert!(error.to_string().contains("is not serializable"), "{error}"); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index 3cb30209dd9..67dcbdb8976 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -17,6 +17,7 @@ pub mod like; pub mod list_contains; pub mod list_length; pub mod list_sum; +pub mod list_transform; pub mod literal; pub mod mask; pub mod merge; From 12459aee324a0c22e559f5e9ce0400412d304b14 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 22 Oct 2026 15:27:53 -0400 Subject: [PATCH 2/2] 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 | 177 +++++++++++++++--- vortex-array/src/expr/lambda/tests.rs | 51 ++++- .../src/scalar_fn/fns/list_transform/mod.rs | 118 ++++++++---- .../src/scalar_fn/fns/list_transform/tests.rs | 177 ++++++++++++++++-- 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, 523 insertions(+), 109 deletions(-) diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 3203d89725e..e96c2079cd9 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; @@ -1388,16 +1389,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..cca043dda4d 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,41 @@ 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(("", 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 d1f4f076fe9..e4ab50a62e8 100644 --- a/vortex-array/src/expr/lambda/tests.rs +++ b/vortex-array/src/expr/lambda/tests.rs @@ -165,19 +165,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] @@ -292,6 +330,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 ff50d3f5701..5484c3f0b59 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; @@ -42,8 +45,10 @@ use crate::validity::Validity; /// 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`. The lambda 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. The lambda runs only on the elements of valid lists. #[derive(Clone)] pub struct ListTransform; @@ -57,6 +62,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 { @@ -75,13 +90,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())), } } @@ -91,18 +106,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())) } @@ -112,16 +130,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( @@ -138,7 +158,15 @@ impl ScalarFnVTable for ListTransform { let list = args.get(0)?.execute::(ctx)?; let list = list_from_list_view(empty_null_lists(list, ctx)?, ctx)?; let elements = list.elements(); - let values = lambda.evaluate(vec![elements.clone()], elements.len())?; + let mut inputs = vec![elements.clone()]; + if args.num_inputs() > 1 { + // 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 index in 1..args.num_inputs() { + inputs.push(args.get(index)?.take(lists.clone())?); + } + } + let values = lambda.evaluate(inputs, elements.len())?; vortex_ensure_eq!(values.len(), elements.len()); // SAFETY: the offsets are the list's, over one value per element. Ok(unsafe { @@ -147,8 +175,10 @@ impl ScalarFnVTable for ListTransform { .into_array()) } - 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 { @@ -159,29 +189,39 @@ 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 exact lists, 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 lists with every null list emptied, so that nothing under a null list reaches the lambda. 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 2f7b7268704..c66474301f9 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -14,6 +14,7 @@ use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::FixedSizeListArray; @@ -64,6 +65,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 = array_session().create_execution_ctx(); Ok(ListArray::try_new( @@ -89,6 +102,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 = array_session().create_execution_ctx(); @@ -442,7 +460,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)), @@ -450,22 +468,149 @@ 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 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 = array_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 4cd658d4774..25ac1a7b957 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 e02b3659b73..c1646aa9684 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 {