From e831e18fb2aedd2f7b0617cc74b251bd9603fd4b Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 12 Oct 2026 15:56:47 -0400 Subject: [PATCH 1/2] list_transform scalar fn Signed-off-by: Matt Katz --- vortex-array/src/expr/exprs.rs | 13 + vortex-array/src/expr/mod.rs | 1 + .../src/scalar_fn/fns/list_transform/mod.rs | 444 ++++++++++++++++++ .../src/scalar_fn/fns/list_transform/tests.rs | 400 ++++++++++++++++ vortex-array/src/scalar_fn/fns/mod.rs | 1 + 5 files changed, 859 insertions(+) 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 7da0d3ead4b..586b5adfbf2 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -45,6 +45,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::ListLambdaOptions; +use crate::scalar_fn::fns::list_transform::ListTransform; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::DuplicateHandling; @@ -1262,6 +1264,17 @@ pub fn list_sum(input: Expression) -> Expression { ListSum.new_expr(NumericalAggregateOpts::default(), [input]) } +/// Creates an expression that transforms every element of a list with `lambda`. +/// +/// `lambda` takes the element and, optionally, its `u64` index. +pub fn list_transform(list: Expression, lambda: Lambda) -> VortexResult { + Expression::try_new_higher_order( + ListTransform.bind(ListLambdaOptions::default()), + [list], + [lambda], + ) +} + /// Creates a bound expression that sums the elements of each list. pub fn bound_list_sum(input: BoundExpression) -> BoundExpression { ListSum diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index a3074af5299..c28158ca6a7 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -93,6 +93,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..ea9979d5129 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -0,0 +1,444 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::sync::Arc; + +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::FixedSizeListArray; +use crate::arrays::InterleaveArray; +use crate::arrays::List; +use crate::arrays::ListArray; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +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::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::BoundLambda; +use crate::matcher::Matcher; +use crate::scalar::Scalar; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::HigherOrderVTable; +use crate::scalar_fn::ReduceNode; +use crate::scalar_fn::ReduceNodeValidity; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::is_not_null_node; +use crate::validity::Validity; + +/// The options of [`ListTransform`]: its lambda, once bound. +/// +/// The bound node's children are the list followed by the lambda's captures. +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] +pub struct ListLambdaOptions(Option>); + +impl ListLambdaOptions { + /// The bound lambda, or `None` in an unbound node. + pub fn lambda(&self) -> Option<&BoundLambda> { + self.0.as_deref() + } + + fn bound(&self) -> VortexResult<&BoundLambda> { + match self.lambda() { + Some(lambda) => Ok(lambda), + None => vortex_bail!("list_transform() is unbound; bind it first"), + } + } +} + +impl From for ListLambdaOptions { + fn from(lambda: BoundLambda) -> Self { + Self(Some(Arc::new(lambda))) + } +} + +impl Display for ListLambdaOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self.lambda() { + None => Ok(()), + Some(lambda) => Display::fmt(lambda, f), + } + } +} + +/// Transform every element of a list with a lambda, preserving the list structure. +/// +/// The lambda's first parameter is the element; an optional second parameter is its zero-based +/// `u64` index within its list. A null list stays null and the body never observably runs on its +/// elements. +#[derive(Clone)] +pub struct ListTransform; + +fn element_dtype(list: &DType) -> VortexResult<&DType> { + match list { + DType::List(element, _) | DType::FixedSizeList(element, ..) => Ok(element), + other => vortex_bail!("list_transform() requires List or FixedSizeList, got {other}"), + } +} + +fn index_dtype() -> DType { + DType::Primitive(PType::U64, Nullability::NonNullable) +} + +impl ScalarFnVTable for ListTransform { + type Options = ListLambdaOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.transform"); + *ID + } + + fn arity(&self, options: &Self::Options) -> Arity { + Arity::Exact( + 1 + options + .lambda() + .map_or(0, |lambda| lambda.capture_dtypes().len()), + ) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("list"), + index => ChildName::from(Arc::from(format!("capture[{}]", index - 1))), + } + } + + fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let lambda = options.bound()?; + let Some((list, captures)) = arg_dtypes.split_first() else { + vortex_bail!("list_transform() requires a list argument"); + }; + let params = self.lambda_params( + options, + 0, + lambda.params().len(), + std::slice::from_ref(list), + )?; + lambda.check_inputs(params.iter().chain(captures))?; + let body = Arc::new(lambda.body_dtype().clone()); + Ok(match list { + DType::List(_, nullability) => DType::List(body, *nullability), + DType::FixedSizeList(_, size, nullability) => { + DType::FixedSizeList(body, *size, *nullability) + } + _ => unreachable!("lambda_params checked the list dtype"), + }) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let lambda = options.bound()?; + let input = args.get(0)?; + let captures = (1..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>>()?; + let input = input.execute_until::(ctx)?; + if let Some(list) = input.as_opt::() { + execute_list(lambda, list.into_owned(), captures, ctx) + } else if let Some(list) = input.as_opt::() { + execute_fixed_size_list(lambda, list.into_owned(), captures, ctx) + } else if let Some(list) = input.as_opt::() { + execute_list_view(lambda, list.into_owned(), captures, ctx) + } else { + unreachable!("AnyList matcher returned a non-list array") + } + } + + fn validity( + &self, + _options: &Self::Options, + node: &T, + ) -> VortexResult> { + Ok(ReduceNodeValidity::Reduced(is_not_null_node( + &node.child(0), + )?)) + } + + fn is_infallible(&self, options: &Self::Options) -> bool { + options.lambda().is_some_and(BoundLambda::is_infallible) + } + + fn higher_order(&self) -> Option<&dyn HigherOrderVTable> { + Some(self) + } +} + +impl HigherOrderVTable for ListTransform { + fn lambda_params( + &self, + _options: &ListLambdaOptions, + lambda_index: usize, + param_count: usize, + arg_dtypes: &[DType], + ) -> VortexResult> { + vortex_ensure!(lambda_index == 0, "list_transform() takes one lambda"); + let Some(list) = arg_dtypes.first() else { + vortex_bail!("list_transform() requires a list argument"); + }; + let element = element_dtype(list)?.clone(); + match param_count { + 1 => Ok(vec![element]), + 2 => Ok(vec![element, index_dtype()]), + n => vortex_bail!("list_transform() lambda must take one or two parameters, got {n}"), + } + } + + fn with_lambdas( + &self, + _options: &ListLambdaOptions, + mut lambdas: Vec, + ) -> VortexResult { + vortex_ensure!(lambdas.len() == 1, "list_transform() takes one lambda"); + Ok(lambdas.pop().vortex_expect("checked length").into()) + } +} + +fn execute_list( + lambda: &BoundLambda, + list: ListArray, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let list = list.reset_offsets(false, ctx)?; + let offsets = list.offsets(); + let sizes = offsets + .slice(1..offsets.len())? + .binary(offsets.slice(0..list.len())?, Operator::Sub)?; + let transformed = transform_elements( + lambda, + list.elements().clone(), + sizes, + list.list_validity(), + captures, + lambda.body_dtype(), + ctx, + )?; + + Ok(ListArray::try_new(transformed, list.offsets().clone(), list.list_validity())?.into_array()) +} + +fn execute_fixed_size_list( + lambda: &BoundLambda, + list: FixedSizeListArray, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let list_size = list.list_size(); + let sizes = ConstantArray::new(u64::from(list_size), list.len()).into_array(); + let transformed = transform_elements( + lambda, + list.elements().clone(), + sizes, + list.fixed_size_list_validity(), + captures, + lambda.body_dtype(), + ctx, + )?; + + Ok(FixedSizeListArray::try_new( + transformed, + list_size, + list.fixed_size_list_validity(), + list.len(), + )? + .into_array()) +} + +fn execute_list_view( + lambda: &BoundLambda, + list: ListViewArray, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // Logical flattening removes overlaps and gives null lists empty element ranges, so each + // invocation has one unambiguous parent row. + let list = list.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?; + let transformed = transform_elements( + lambda, + list.elements().clone(), + list.sizes().clone(), + list.listview_validity(), + captures, + lambda.body_dtype(), + ctx, + )?; + + // SAFETY: the rebuilt view is zero-copyable to List and `transform_elements` preserves its + // element domain, so its offsets, sizes, and validity remain valid and zero-copyable. + Ok(unsafe { + ListViewArray::new_unchecked( + transformed, + list.offsets().clone(), + list.sizes().clone(), + list.listview_validity(), + ) + .with_zero_copy_to_list(true) + } + .into_array()) +} + +fn transform_elements( + lambda: &BoundLambda, + elements: ArrayRef, + sizes: ArrayRef, + validity: Validity, + captures: Vec, + body_dtype: &DType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let invocation_count = elements.len(); + + // An infallible body can run over element slots hidden under null lists: the results stay + // hidden behind the list validity. Only a fallible body needs those slots filtered out. + let may_hide = !matches!(validity, Validity::NonNullable | Validity::AllValid); + let needs_parents = !captures.is_empty() || (may_hide && !lambda.is_infallible()); + let parent_indices = needs_parents + .then(|| parent_indices(sizes.clone(), invocation_count)) + .transpose()?; + + let invocation_mask = match &parent_indices { + Some(parent_indices) if may_hide && !lambda.is_infallible() => { + let invocation_validity = validity.take(parent_indices)?; + let invocation_mask = invocation_validity.execute_mask(invocation_count, ctx)?; + (invocation_mask.true_count() != invocation_count).then_some(invocation_mask) + } + _ => None, + }; + + let filter = |array: ArrayRef| -> VortexResult { + match &invocation_mask { + Some(mask) => array.filter(mask.clone()), + None => Ok(array), + } + }; + + let mut env = vec![filter(elements)?]; + if lambda.params().len() > 1 { + env.push(filter(local_indices(sizes, invocation_count)?)?); + } + if let Some(parent_indices) = parent_indices { + let capture_indices = filter(parent_indices)?; + for capture in captures { + env.push(capture.take(capture_indices.clone())?); + } + } + let transformed = lambda.evaluate(&env)?; + + match invocation_mask { + Some(mask) => scatter_valid_invocations(transformed, &mask, body_dtype), + None => Ok(transformed), + } +} + +fn parent_indices(sizes: ArrayRef, element_count: usize) -> VortexResult { + let parent_count = sizes.len(); + let sizes_dtype = DType::Primitive( + sizes.dtype().as_ptype().to_unsigned(), + Nullability::NonNullable, + ); + let sizes = sizes.cast(sizes_dtype)?; + let starts = + PrimitiveArray::from_iter((0..parent_count).map(|index| index as u64)).into_array(); + let multipliers = ConstantArray::new(0_u64, parent_count).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, sizes, multipliers, element_count)?.into_array()) +} + +fn local_indices(sizes: ArrayRef, element_count: usize) -> VortexResult { + let parent_count = sizes.len(); + let sizes_dtype = DType::Primitive( + sizes.dtype().as_ptype().to_unsigned(), + Nullability::NonNullable, + ); + let sizes = sizes.cast(sizes_dtype)?; + let starts = ConstantArray::new(0_u64, parent_count).into_array(); + let multipliers = ConstantArray::new(1_u64, parent_count).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, sizes, multipliers, element_count)?.into_array()) +} + +fn scatter_valid_invocations( + transformed: ArrayRef, + mask: &Mask, + body_dtype: &DType, +) -> VortexResult { + vortex_ensure!( + transformed.len() == mask.true_count(), + "list_transform() produced {} valid invocations, expected {}", + transformed.len(), + mask.true_count() + ); + + let mut array_indices = BufferMut::::with_capacity(mask.len()); + let mut row_indices = BufferMut::::with_capacity(mask.len()); + let mut valid_index = 0_u64; + for valid in mask.iter() { + if valid { + array_indices.push(0); + row_indices.push(valid_index); + valid_index += 1; + } else { + array_indices.push(1); + row_indices.push(0); + } + } + + let placeholder = if body_dtype.is_nullable() { + Scalar::null(body_dtype.clone()) + } else { + Scalar::zero_value(body_dtype) + }; + let default = ConstantArray::new(placeholder, 1).into_array(); + Ok(InterleaveArray::try_new( + vec![transformed, default], + array_indices.into_array(), + row_indices.into_array(), + )? + .into_array()) +} + +/// Matches a `List`, `ListView`, or `FixedSizeList` physical array. +struct AnyList; + +impl Matcher for AnyList { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option> { + (array.as_opt::().is_some() + || array.as_opt::().is_some() + || array.as_opt::().is_some()) + .then_some(()) + } +} + +#[cfg(test)] +mod tests; 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..cacc780bb87 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -0,0 +1,400 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::Canonical; +use crate::IntoArray; +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::Expression; +use crate::expr::cast; +use crate::expr::checked_add; +use crate::expr::col; +use crate::expr::eq; +use crate::expr::is_not_null; +use crate::expr::lambda; +use crate::expr::list_length; +use crate::expr::list_sum; +use crate::expr::list_transform; +use crate::expr::lit; +use crate::expr::root; +use crate::expr::var; +use crate::validity::Validity; + +/// `[[1, 2], null (hiding [i32::MAX]), [3, null]]` +fn lists() -> VortexResult { + let elements = + PrimitiveArray::from_option_iter([Some(1_i32), Some(2), Some(i32::MAX), Some(3), None]); + Ok(ListArray::try_new( + elements.into_array(), + buffer![0_u32, 2, 3, 5].into_array(), + Validity::from_iter([true, false, true]), + )? + .into_array()) +} + +/// `[0, 10, 20]` +fn a() -> ArrayRef { + buffer![0_i32, 10, 20].into_array() +} + +/// Rows `{l, a}`. +fn rows(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 { + Ok(ListArray::try_new( + values, + PrimitiveArray::from_iter(offsets.iter().copied()).into_array(), + validity, + )? + .into_array()) +} + +/// A transform of [`lists`]: same offsets and validity, new elements. +fn transformed(values: ArrayRef) -> VortexResult { + list_of( + values, + &[0, 2, 3, 5], + Validity::from_iter([true, false, true]), + ) +} + +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 fallible_body_skips_elements_under_null_lists() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x"], checked_add(var("x"), lit(1_i32)))?)?; + // The hidden slot holds a placeholder, not `i32::MAX + 1`. + let expected = transformed( + PrimitiveArray::from_option_iter([Some(2_i32), Some(3), Some(0), Some(4), None]) + .into_array(), + )?; + assert_applies(rows(lists()?, a())?, &expr, expected) +} + +#[test] +fn fallible_captures_are_not_evaluated_for_rows_without_elements() -> VortexResult<()> { + // `[[1], [], null (hiding [5])]` with `a = [1, MAX, MAX]`. `a + 1` overflows on the last two + // rows, but neither has an element to evaluate, so the body never computes it there. + let l = list_of( + buffer![1_i32, 5].into_array(), + &[0, 1, 1, 2], + Validity::from_iter([true, true, false]), + )?; + let a = buffer![1_i32, i32::MAX, i32::MAX].into_array(); + let expr = list_transform( + col("l"), + lambda( + ["x"], + checked_add(var("x"), checked_add(col("a"), lit(1_i32))), + )?, + )?; + let expected = list_of( + buffer![3_i32, 0].into_array(), + &[0, 1, 1, 2], + Validity::from_iter([true, true, false]), + )?; + assert_applies(rows(l, a)?, &expr, expected) +} + +#[test] +fn captures_are_spread_per_element() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x"], checked_add(var("x"), col("a")))?)?; + let input = rows(lists()?, a())?; + assert_eq!(expr.bind(input.dtype())?.children().len(), 2); + + let expected = transformed( + PrimitiveArray::from_option_iter([Some(1_i32), Some(2), Some(0), Some(23), None]) + .into_array(), + )?; + assert_applies(input, &expr, expected) +} + +#[test] +fn a_null_capture_nulls_the_elements_of_its_row() -> VortexResult<()> { + let a = PrimitiveArray::from_option_iter([None, Some(10_i32), Some(20)]).into_array(); + let expr = list_transform(col("l"), lambda(["x"], checked_add(var("x"), col("a")))?)?; + let expected = transformed( + PrimitiveArray::from_option_iter([None, None, None, Some(23_i32), None]).into_array(), + )?; + assert_applies(rows(lists()?, a)?, &expr, expected) +} + +#[test] +fn index_parameter_is_local_to_each_list() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x", "i"], eq(var("i"), lit(1_u64)))?)?; + let expected = + transformed(BoolArray::from_iter([false, true, false, false, true]).into_array())?; + assert_applies(rows(lists()?, a())?, &expr, expected) +} + +#[test] +fn an_unused_index_parameter_is_accepted() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x", "i"], eq(var("x"), lit(1_i32)))?)?; + let expected = transformed( + BoolArray::from_iter([Some(true), Some(false), Some(false), Some(false), None]) + .into_array(), + )?; + assert_applies(rows(lists()?, a())?, &expr, expected) +} + +#[test] +fn a_lambda_takes_one_or_two_parameters() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x", "i", "j"], var("x"))?)?; + assert!(expr.bind(rows(lists()?, a())?.dtype()).is_err()); + Ok(()) +} + +#[test] +fn fixed_size_lists_keep_their_width() -> VortexResult<()> { + let fsl = FixedSizeListArray::try_new( + buffer![1_i32, 2, 3, 4].into_array(), + 2, + Validity::NonNullable, + 2, + )? + .into_array(); + let expr = list_transform(root(), lambda(["x"], eq(var("x"), lit(2_i32)))?)?; + let expected = FixedSizeListArray::try_new( + BoolArray::from_iter([false, true, false, false]).into_array(), + 2, + Validity::NonNullable, + 2, + )? + .into_array(); + assert_applies(fsl, &expr, expected) +} + +#[test] +fn list_views_are_transformed_through_overlapping_views() -> VortexResult<()> { + // `[[3, 4], [1], [2, 3]]`, sharing element 3. + 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, + )? + .into_array(); + let expr = list_transform(root(), lambda(["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) +} + +/// Rows `{ll, a}` with `ll = [[[1, 2], [3]], [[4]]]` and `a = [10, 20]`. +fn nested_rows() -> VortexResult { + let inner = list_of( + buffer![1_i32, 2, 3, 4].into_array(), + &[0, 2, 3, 4], + Validity::NonNullable, + )?; + let ll = list_of(inner, &[0, 2, 3], Validity::NonNullable)?; + Ok(StructArray::try_new( + FieldNames::from(["ll", "a"]), + vec![ll, buffer![10_i32, 20].into_array()], + 2, + Validity::NonNullable, + )? + .into_array()) +} + +/// [`nested_rows`] under `(xs) -> list_transform(xs, (y) -> y + a + len(xs))`. +fn nested_expected() -> VortexResult { + let expected_inner = list_of( + buffer![13_i32, 14, 14, 25].into_array(), + &[0, 2, 3, 4], + Validity::NonNullable, + )?; + list_of(expected_inner, &[0, 2, 3], Validity::NonNullable) +} + +fn i32_() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) +} + +#[test] +fn nested_transforms_chain_captures_outward() -> VortexResult<()> { + let input = nested_rows()?; + + // (xs) -> list_transform(xs, (y) -> y + a + len(xs)) + let length = cast(list_length(var("xs")), i32_()); + let expr = list_transform( + col("ll"), + lambda( + ["xs"], + list_transform( + var("xs"), + lambda(["y"], checked_add(checked_add(var("y"), col("a")), length))?, + )?, + )?, + )?; + // The inner lambda captures `$.a` and `xs`. In the outer body, `xs` is a parameter, so only + // `$.a` becomes an outer capture. + assert_eq!(expr.bind(input.dtype())?.children().len(), 2); + + assert_applies(input, &expr, nested_expected()?) +} + +#[test] +fn a_constant_body_fills_every_element() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x"], lit(7_i32))?)?; + let expected = transformed(buffer![7_i32, 7, 7, 7, 7].into_array())?; + assert_applies(rows(lists()?, a())?, &expr, expected) +} + +#[test] +fn lists_without_visible_elements_never_invoke_the_body() -> VortexResult<()> { + // `x + MAX` overflows on every hidden element, so the body must not run on any of them. + let body = lambda(["x"], checked_add(var("x"), lit(i32::MAX)))?; + + // `[[], null (hiding [1])]` + let list = list_of( + buffer![1_i32].into_array(), + &[0, 0, 1], + Validity::from_iter([true, false]), + )?; + let expected = list_of( + buffer![0_i32].into_array(), + &[0, 0, 1], + Validity::from_iter([true, false]), + )?; + assert_applies(list, &list_transform(root(), body.clone())?, expected)?; + + // `[null (hiding [1]), null (hiding [2])]` + let fsl = + FixedSizeListArray::try_new(buffer![1_i32, 2].into_array(), 1, Validity::AllInvalid, 2)? + .into_array(); + let expected = + FixedSizeListArray::try_new(buffer![0_i32, 0].into_array(), 1, Validity::AllInvalid, 2)? + .into_array(); + assert_applies(fsl, &list_transform(root(), body)?, expected) +} + +#[test] +fn validity_is_the_list_validity() -> VortexResult<()> { + let dtype = rows(lists()?, a())?.dtype().clone(); + let expr = list_transform(col("l"), lambda(["x"], checked_add(var("x"), col("a")))?)?; + assert_eq!( + expr.bind(&dtype)?.validity()?, + is_not_null(col("l")).bind(&dtype)? + ); + Ok(()) +} + +#[test] +fn both_forms_display() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x"], checked_add(var("x"), col("a")))?)?; + assert_eq!( + expr.to_string(), + "vortex.list.transform($.l, (x) -> ($x + $.a))" + ); + assert_eq!( + expr.bind(rows(lists()?, a())?.dtype())?.to_string(), + "vortex.list.transform($.l, $.a, opts=(x) -> ($x + $capture0))" + ); + Ok(()) +} + +#[test] +fn a_capture_accepts_a_narrowed_nullability() -> VortexResult<()> { + // In a nullable struct `$.a` is nullable. A layout reader that steps into the struct's fields + // rebuilds the node over the fields' own, non-nullable dtypes. + let input = rows(lists()?, a())?; + let expr = list_transform(col("l"), lambda(["x"], checked_add(var("x"), col("a")))?)?; + let bound = expr.bind(&input.dtype().as_nullable())?; + let stepped_in = + bound.with_children([col("l").bind(input.dtype())?, col("a").bind(input.dtype())?])?; + + let expected = transformed( + PrimitiveArray::from_option_iter([Some(1_i32), Some(2), None, Some(23), None]).into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(input.apply_bound(&stepped_in)?, expected, &mut ctx); + Ok(()) +} + +#[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(["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_lazy_transform_slices() -> VortexResult<()> { + let expr = list_transform(col("l"), lambda(["x"], checked_add(var("x"), col("a")))?)?; + let actual = rows(lists()?, a())?.apply(&expr)?.slice(1..3)?; + let expected = list_of( + PrimitiveArray::from_option_iter([Some(0_i32), Some(23), None]).into_array(), + &[0, 1, 3], + Validity::from_iter([false, true]), + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn list_sum_counts_non_null_elements() -> VortexResult<()> { + let expr = list_sum(list_transform( + col("l"), + lambda(["x"], is_not_null(var("x")))?, + )?); + let mut ctx = array_session().create_execution_ctx(); + let actual = rows(lists()?, a())? + .apply(&expr)? + .execute::(&mut ctx)? + .into_array(); + assert_eq!(actual.execute_scalar(0, &mut ctx)?, 2_u64.into()); + assert!(actual.execute_scalar(1, &mut ctx)?.is_null()); + assert_eq!(actual.execute_scalar(2, &mut ctx)?, 1_u64.into()); + 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 1d13ace5f988a638d73491e4f8f9ad3e1cf72e52 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 12 Oct 2026 16:12:48 -0400 Subject: [PATCH 2/2] list_transform lambda hooks on its ScalarFnVTable impl Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/list_transform/mod.rs | 63 +++++++++---------- 1 file changed, 28 insertions(+), 35 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 ea9979d5129..563d78d05c8 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -44,7 +44,6 @@ use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::HigherOrderVTable; use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnId; @@ -130,6 +129,34 @@ impl ScalarFnVTable for ListTransform { } } + fn lambda_params( + &self, + _options: &Self::Options, + lambda_index: usize, + param_count: usize, + arg_dtypes: &[DType], + ) -> VortexResult> { + vortex_ensure!(lambda_index == 0, "list_transform() takes one lambda"); + let Some(list) = arg_dtypes.first() else { + vortex_bail!("list_transform() requires a list argument"); + }; + let element = element_dtype(list)?.clone(); + match param_count { + 1 => Ok(vec![element]), + 2 => Ok(vec![element, index_dtype()]), + n => vortex_bail!("list_transform() lambda must take one or two parameters, got {n}"), + } + } + + fn with_lambdas( + &self, + _options: &Self::Options, + mut lambdas: Vec, + ) -> VortexResult { + vortex_ensure!(lambdas.len() == 1, "list_transform() takes one lambda"); + Ok(lambdas.pop().vortex_expect("checked length").into()) + } + fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { let lambda = options.bound()?; let Some((list, captures)) = arg_dtypes.split_first() else { @@ -188,40 +215,6 @@ impl ScalarFnVTable for ListTransform { fn is_infallible(&self, options: &Self::Options) -> bool { options.lambda().is_some_and(BoundLambda::is_infallible) } - - fn higher_order(&self) -> Option<&dyn HigherOrderVTable> { - Some(self) - } -} - -impl HigherOrderVTable for ListTransform { - fn lambda_params( - &self, - _options: &ListLambdaOptions, - lambda_index: usize, - param_count: usize, - arg_dtypes: &[DType], - ) -> VortexResult> { - vortex_ensure!(lambda_index == 0, "list_transform() takes one lambda"); - let Some(list) = arg_dtypes.first() else { - vortex_bail!("list_transform() requires a list argument"); - }; - let element = element_dtype(list)?.clone(); - match param_count { - 1 => Ok(vec![element]), - 2 => Ok(vec![element, index_dtype()]), - n => vortex_bail!("list_transform() lambda must take one or two parameters, got {n}"), - } - } - - fn with_lambdas( - &self, - _options: &ListLambdaOptions, - mut lambdas: Vec, - ) -> VortexResult { - vortex_ensure!(lambdas.len() == 1, "list_transform() takes one lambda"); - Ok(lambdas.pop().vortex_expect("checked length").into()) - } } fn execute_list(