From 9534b7d965878b780ee97278eb39e9ca49f60ac3 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 14 Oct 2026 16:57:32 -0400 Subject: [PATCH 1/3] 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 ed77d494daf84bf84aca8bb76da991a5faece4de Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 22 Oct 2026 15:59:54 -0400 Subject: [PATCH 2/3] skip the list rebuild when lists are already exact Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/list_transform/mod.rs | 77 +++++++++++++++---- 1 file changed, 62 insertions(+), 15 deletions(-) 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..6e3693ba416 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -10,15 +10,19 @@ use std::fmt::Formatter; 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; use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -133,10 +137,7 @@ impl ScalarFnVTable for ListTransform { 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 list = exact_lists(args.get(0)?, ctx)?; let elements = list.elements(); let values = lambda.evaluate(vec![elements.clone()], elements.len())?; vortex_ensure_eq!(values.len(), elements.len()); @@ -184,28 +185,74 @@ fn element_dtype(list: &DType, params: &[FieldName]) -> VortexResult<[DType; 1]> Ok([element.as_ref().clone()]) } +/// The lists laid out so that 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. +fn exact_lists(list: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // A list array usually is exact already, and then needs no rebuild. + let list = match list.try_downcast::() { + Ok(list) if is_exact(&list, ctx)? => return Ok(list), + Ok(list) => list.into_array(), + Err(list) => list, + }; + let list = list.execute::(ctx)?; + list_from_list_view(empty_null_lists(list, ctx)?, ctx) +} + +/// Whether nothing comes before the first list or after the last, and no null list holds elements. +fn is_exact(list: &ListArray, ctx: &mut ExecutionCtx) -> VortexResult { + let offsets = list.offsets().clone().execute::(ctx)?; + let nulls = null_runs(&list.list_validity(), list.len(), ctx)?; + Ok(match_each_integer_ptype!(offsets.ptype(), |O| { + let offsets = offsets.as_slice::(); + let last: usize = offsets[offsets.len() - 1].as_(); + // The offsets are sorted, so a run of null lists is empty when it starts and ends at the + // same offset. + offsets[0] == O::default() + && last == list.elements().len() + && nulls + .iter() + .all(|&(start, end)| offsets[start] == offsets[end]) + })) +} + +/// The runs of null rows, as `(start, end)` ranges. +fn null_runs( + validity: &Validity, + len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if matches!(validity, Validity::NonNullable | Validity::AllValid) { + return Ok(Vec::new()); + } + Ok(match (!&validity.execute_mask(len, ctx)?).slices() { + AllOr::All => vec![(0, len)], + AllOr::None => Vec::new(), + AllOr::Some(runs) => runs.to_vec(), + }) +} + /// 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) { + let nulls = null_runs(&validity, list.len(), ctx)?; + if nulls.is_empty() { 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()) { + if nulls + .iter() + .all(|&(start, end)| sizes[start..end].iter().all(|size| *size == T::default())) + { return Ok(list); } - PrimitiveArray::from_iter((0..sizes.len()).map(|index| { - if null(index) { - T::default() - } else { - sizes[index] - } - })) + let mut sizes = BufferMut::copy_from(sizes); + for &(start, end) in &nulls { + sizes[start..end].fill(T::default()); + } + PrimitiveArray::new(sizes.freeze(), Validity::NonNullable) }); // 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, From 5ac5458a82245e815b0d331de100f13abb143f83 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 23 Oct 2026 11:35:56 -0400 Subject: [PATCH 3/3] keep the lists' shape unless a lambda that can fail would see hidden elements Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/list_transform/mod.rs | 157 +++++++++++++----- .../src/scalar_fn/fns/list_transform/tests.rs | 82 +++++++++ 2 files changed, 202 insertions(+), 37 deletions(-) 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 6e3693ba416..390c9e3c7de 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -11,6 +11,7 @@ use std::sync::Arc; use itertools::Itertools; use num_traits::AsPrimitive; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -30,6 +31,7 @@ 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; @@ -47,7 +49,8 @@ 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. +/// encoding, and the result is a `List` of the same shape. A lambda that can fail runs only on the +/// elements of valid lists. #[derive(Clone)] pub struct ListTransform; @@ -137,15 +140,12 @@ impl ScalarFnVTable for ListTransform { let ListTransformOptions::Bound(lambda) = options else { vortex_bail!("list_transform()'s lambda is not bound"); }; - let list = exact_lists(args.get(0)?, 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()) + let list = if lambda.is_infallible() { + Lists::as_they_are(args.get(0)?, ctx)? + } else { + Lists::visible_elements_only(args.get(0)?, ctx)? + }; + list.with_elements(|elements| lambda.evaluate(vec![elements.clone()], elements.len())) } fn is_strict(&self, _options: &ListTransformOptions) -> bool { @@ -185,34 +185,113 @@ fn element_dtype(list: &DType, params: &[FieldName]) -> VortexResult<[DType; 1]> Ok([element.as_ref().clone()]) } -/// The lists laid out so that 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. -fn exact_lists(list: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // A list array usually is exact already, and then needs no rebuild. - let list = match list.try_downcast::() { - Ok(list) if is_exact(&list, ctx)? => return Ok(list), - Ok(list) => list.into_array(), - Err(list) => list, - }; - let list = list.execute::(ctx)?; - list_from_list_view(empty_null_lists(list, ctx)?, ctx) +/// The lists the lambda runs over: their elements buffer, and the offsets (and sizes) that put the +/// results back into lists of the same shape. +enum Lists { + List(ListArray), + View(ListViewArray), } -/// Whether nothing comes before the first list or after the last, and no null list holds elements. -fn is_exact(list: &ListArray, ctx: &mut ExecutionCtx) -> VortexResult { +impl Lists { + /// The lists as they are. For a lambda that can't fail: it runs over the whole elements buffer, + /// and results for elements no valid list shows are never read. + fn as_they_are(list: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(match list.try_downcast::() { + Ok(list) => Self::List(list), + Err(list) => Self::View(list.execute::(ctx)?), + }) + } + + /// The lists with nothing in their elements buffer but the elements of valid lists. For a lambda + /// that can fail, which must not fail on a value nobody sees. + /// + /// The lists keep their shape, trimmed of what comes before the first list and after the last, + /// unless a null list holds elements or a list view's lists aren't laid out back to back. Only + /// then are the elements copied. + fn visible_elements_only(list: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let list = match list.try_downcast::() { + Ok(list) => match trim_list(&list, ctx)? { + Some(trimmed) => return Ok(Self::List(trimmed)), + None => list.into_array(), + }, + Err(list) => list, + }; + let view = list.execute::(ctx)?; + match empty_null_lists(&view, ctx)? { + None if view.is_zero_copy_to_list() => Ok(Self::View( + view.rebuild(ListViewRebuildMode::TrimElements, ctx)?, + )), + None => Ok(Self::List(list_from_list_view(view, ctx)?)), + Some(emptied) => Ok(Self::List(list_from_list_view(emptied, ctx)?)), + } + } + + /// The same lists, with `f` of their elements buffer as their elements. + fn with_elements( + self, + f: impl FnOnce(&ArrayRef) -> VortexResult, + ) -> VortexResult { + match self { + Self::List(list) => { + let values = f(list.elements())?; + vortex_ensure_eq!(values.len(), list.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()) + } + .into_array()) + } + Self::View(view) => { + let values = f(view.elements())?; + vortex_ensure_eq!(values.len(), view.elements().len()); + // SAFETY: the offsets, sizes and validity are the view's, over one value per + // element, so the views keep their bounds and their layout. + Ok(unsafe { + ListViewArray::new_unchecked( + values, + view.offsets().clone(), + view.sizes().clone(), + view.listview_validity(), + ) + .with_zero_copy_to_list(view.is_zero_copy_to_list()) + } + .into_array()) + } + } + } +} + +/// The list without what comes before its first list and after its last, or `None` when a null +/// list holds elements, which only a rebuild can drop. +fn trim_list(list: &ListArray, ctx: &mut ExecutionCtx) -> VortexResult> { let offsets = list.offsets().clone().execute::(ctx)?; let nulls = null_runs(&list.list_validity(), list.len(), ctx)?; - Ok(match_each_integer_ptype!(offsets.ptype(), |O| { + match_each_integer_ptype!(offsets.ptype(), |O| { let offsets = offsets.as_slice::(); - let last: usize = offsets[offsets.len() - 1].as_(); // The offsets are sorted, so a run of null lists is empty when it starts and ends at the // same offset. - offsets[0] == O::default() - && last == list.elements().len() - && nulls - .iter() - .all(|&(start, end)| offsets[start] == offsets[end]) - })) + if nulls + .iter() + .any(|&(start, end)| offsets[start] != offsets[end]) + { + return Ok(None); + } + let first = offsets[0]; + let (start, end): (usize, usize) = (first.as_(), offsets[offsets.len() - 1].as_()); + if start == 0 && end == list.elements().len() { + return Ok(Some(list.clone())); + } + let shifted: Buffer = offsets.iter().map(|&offset| offset - first).collect(); + // SAFETY: the elements are cut to the lists' range and the offsets shifted by its start, + // so every list keeps its elements. + Ok(Some(unsafe { + ListArray::new_unchecked( + list.elements().slice(start..end)?, + shifted.into_array(), + list.list_validity(), + ) + })) + }) } /// The runs of null rows, as `(start, end)` ranges. @@ -231,12 +310,16 @@ fn null_runs( }) } -/// 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 { +/// The lists with every null list emptied, so that nothing under a null list reaches the lambda, or +/// `None` when no null list holds elements. +fn empty_null_lists( + list: &ListViewArray, + ctx: &mut ExecutionCtx, +) -> VortexResult> { let validity = list.listview_validity(); let nulls = null_runs(&validity, list.len(), ctx)?; if nulls.is_empty() { - return Ok(list); + return Ok(None); } let sizes = list.sizes().clone().execute::(ctx)?; let sizes = match_each_integer_ptype!(sizes.ptype(), |T| { @@ -246,7 +329,7 @@ fn empty_null_lists(list: ListViewArray, ctx: &mut ExecutionCtx) -> VortexResult .iter() .all(|&(start, end)| sizes[start..end].iter().all(|size| *size == T::default())) { - return Ok(list); + return Ok(None); } let mut sizes = BufferMut::copy_from(sizes); for &(start, end) in &nulls { @@ -257,12 +340,12 @@ fn empty_null_lists(list: ListViewArray, ctx: &mut ExecutionCtx) -> VortexResult // 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 { + Ok(Some(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 index 2f7b7268704..7b6783acb53 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -3,11 +3,13 @@ use std::sync::Arc; +use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexResult; use super::ListTransform; use super::ListTransformOptions; +use super::Lists; use crate::ArrayRef; use crate::IntoArray; use crate::RecursiveCanonical; @@ -21,6 +23,7 @@ use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::FieldNames; @@ -146,6 +149,85 @@ fn a_sliced_list_skips_a_null_lists_elements() -> VortexResult<()> { assert_applies(lists, &list_transform(root(), plus_one()?), expected) } +/// `[[1], [2]]` over `[MAX, 1, 2, MAX]`. +fn sliced_list_with_leftovers() -> VortexResult { + list_of( + buffer![i32::MAX, 1_i32, 2, i32::MAX].into_array(), + &[1, 2, 3], + Validity::NonNullable, + ) +} + +/// `[[1], [2]]` over `[MAX, 1, 2]`: the lists are back to back, after a leftover. +fn list_view_with_a_leftover_in_front() -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + let view = ListViewArray::try_new( + buffer![i32::MAX, 1_i32, 2].into_array(), + buffer![1_u32, 2].into_array(), + buffer![1_u32, 1].into_array(), + Validity::NonNullable, + &mut ctx, + )?; + // SAFETY: the lists are sorted, with nothing between them. + Ok(unsafe { view.with_zero_copy_to_list(true) }.into_array()) +} + +/// `[[1], [2]]` over `[1, MAX, 2]`: the lists skip an element. +fn list_view_with_a_gap() -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + Ok(ListViewArray::try_new( + buffer![1_i32, i32::MAX, 2].into_array(), + buffer![0_u32, 2].into_array(), + buffer![1_u32, 1].into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array()) +} + +/// `x + 1` would overflow on `i32::MAX`, which no list shows. +#[rstest] +#[case::sliced_list(sliced_list_with_leftovers)] +#[case::list_view_with_a_leftover_in_front(list_view_with_a_leftover_in_front)] +#[case::list_view_with_a_gap(list_view_with_a_gap)] +fn a_lambda_that_can_fail_sees_only_the_lists_elements( + #[case] input: fn() -> VortexResult, +) -> VortexResult<()> { + let expected = list_of( + buffer![2_i32, 3].into_array(), + &[0, 1, 2], + Validity::NonNullable, + )?; + assert_applies(input()?, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn a_lambda_that_can_fail_runs_on_a_slice_without_its_leftovers() -> VortexResult<()> { + // Canonicalizing the result drops the leftovers too, so look at what the lambda runs on. + let mut ctx = array_session().create_execution_ctx(); + let lists = Lists::visible_elements_only(sliced_list_with_leftovers()?, &mut ctx)?; + lists.with_elements(|elements| { + assert_arrays_eq!(elements.clone(), buffer![1_i32, 2].into_array(), &mut ctx); + Ok(elements.clone()) + })?; + Ok(()) +} + +#[test] +fn a_lambda_that_cannot_fail_runs_on_the_elements_in_place() -> VortexResult<()> { + // `x = 3` cannot fail, so it runs over the whole elements buffer, the null list's `i32::MAX` + // included, and nothing is copied. + let mut ctx = array_session().create_execution_ctx(); + let result = lists()? + .apply(&list_transform( + root(), + Lambda::new(["x"], eq(var("x"), lit(3_i32)))?, + ))? + .execute::(&mut ctx)?; + assert_eq!(result.elements().len(), 5); + Ok(()) +} + #[test] fn every_list_null() -> VortexResult<()> { // `[null (hiding [i32::MAX]), null (hiding [i32::MAX])]`