From 9276eaa1a520f08dbc8e6dc02c1fe9509de0abb0 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 14 Oct 2026 16:57:32 -0400 Subject: [PATCH 1/7] 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 8634ae901b3..4904c3791c3 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}; @@ -1391,6 +1393,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 bbdbbfdf01a..6345f8e992d 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -88,6 +88,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 baa3214a95b3672e9728c2d65921da132797b182 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 22 Oct 2026 15:59:54 -0400 Subject: [PATCH 2/7] 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 d1c0b755280f7315b176660fdc6c9d278c1c3bb2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 23 Oct 2026 11:35:56 -0400 Subject: [PATCH 3/7] 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])]` From 03678156d1c7a007e5ff621d9d1460df9cebff0c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 10:10:56 +0000 Subject: [PATCH 4/7] update list_transform tests for develop Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/list_transform/tests.rs | 37 +++++++++---------- vortex-expr/src/transform/replace.rs | 19 ++++++++++ 2 files changed, 37 insertions(+), 19 deletions(-) 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 7b6783acb53..01b0fdf62cf 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -2,10 +2,12 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::sync::LazyLock; use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_session::VortexSession; use super::ListTransform; use super::ListTransformOptions; @@ -14,7 +16,6 @@ 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; @@ -42,12 +43,13 @@ 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; +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + /// `[[1, 2], null (hiding [i32::MAX]), [3, null]]` fn lists() -> VortexResult { list_of( @@ -68,7 +70,7 @@ fn rows(l: ArrayRef) -> VortexResult { } fn list_of(values: ArrayRef, offsets: &[u32], validity: Validity) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); Ok(ListArray::try_new( values, PrimitiveArray::from_iter(offsets.iter().copied()).into_array(), @@ -94,7 +96,7 @@ fn plus_one() -> VortexResult { fn assert_applies(input: ArrayRef, expr: &Expression, expected: ArrayRef) -> VortexResult<()> { let actual = input.apply(expr)?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } @@ -160,7 +162,7 @@ fn sliced_list_with_leftovers() -> VortexResult { /// `[[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 mut ctx = SESSION.create_execution_ctx(); let view = ListViewArray::try_new( buffer![i32::MAX, 1_i32, 2].into_array(), buffer![1_u32, 2].into_array(), @@ -174,7 +176,7 @@ fn list_view_with_a_leftover_in_front() -> VortexResult { /// `[[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(); + let mut ctx = SESSION.create_execution_ctx(); Ok(ListViewArray::try_new( buffer![1_i32, i32::MAX, 2].into_array(), buffer![0_u32, 2].into_array(), @@ -204,7 +206,7 @@ fn a_lambda_that_can_fail_sees_only_the_lists_elements( #[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 mut ctx = 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); @@ -217,7 +219,7 @@ fn a_lambda_that_can_fail_runs_on_a_slice_without_its_leftovers() -> VortexResul 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 mut ctx = SESSION.create_execution_ctx(); let result = lists()? .apply(&list_transform( root(), @@ -251,7 +253,7 @@ fn a_failure_in_a_valid_list_is_an_error() -> VortexResult<()> { &[0, 1, 2], Validity::from_iter([false, true]), )?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let error = lists .apply(&list_transform(root(), plus_one()?))? .execute::(&mut ctx) @@ -403,7 +405,7 @@ fn three_nested_lambdas() -> VortexResult<()> { 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 mut ctx = SESSION.create_execution_ctx(); let view = list_of( buffer![1_i32, 2, 3, 4, 5].into_array(), &[0, 2, 3, 5], @@ -423,7 +425,7 @@ fn sliced_list_views() -> VortexResult<()> { #[test] fn overlapping_list_views() -> VortexResult<()> { // `[[3, 4], [1], [2, 3]]`, sharing element 3. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let view = ListViewArray::try_new( buffer![1_i32, 2, 3, 4].into_array(), buffer![2_u32, 0, 1].into_array(), @@ -518,7 +520,7 @@ fn applying_the_bound_node() -> VortexResult<()> { 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(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(input.apply_bound(&bound)?, expected, &mut ctx); Ok(()) } @@ -552,9 +554,8 @@ fn a_lambda_bound_to_another_element_dtype_is_an_error() -> VortexResult<()> { } #[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. +fn a_list_in_a_nested_field() -> VortexResult<()> { + // Rows `{s: {l}}`. let input = StructArray::try_new( FieldNames::from(["s"]), vec![rows(lists()?)?], @@ -562,9 +563,7 @@ fn replacing_the_root_rereads_the_list() -> VortexResult<()> { 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")); + let expr = list_transform(get_item("l", col("s")), plus_one()?); assert_eq!( expr.to_string(), "vortex.list.transform($.s.l, (x) -> (@x + 1i32))" @@ -592,7 +591,7 @@ fn an_infallible_transform_pushes_into_dictionary_values() -> VortexResult<()> { &[0, 2, 3, 5, 6], Validity::NonNullable, )?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } diff --git a/vortex-expr/src/transform/replace.rs b/vortex-expr/src/transform/replace.rs index a990b601d0f..12e408981d2 100644 --- a/vortex-expr/src/transform/replace.rs +++ b/vortex-expr/src/transform/replace.rs @@ -28,9 +28,16 @@ pub fn replace(expr: Expression, needle: &Expression, replacement: Expression) - #[cfg(test)] mod test { use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::expr::Lambda; + use vortex_array::expr::checked_add; + use vortex_array::expr::col; use vortex_array::expr::get_item; + use vortex_array::expr::list_transform; use vortex_array::expr::lit; use vortex_array::expr::pack; + use vortex_array::expr::root; + use vortex_array::expr::var; + use vortex_error::VortexResult; use super::replace; @@ -51,4 +58,16 @@ mod test { let replaced_expr = replace(e, &needle, replacement); assert_eq!(replaced_expr.to_string(), "pack(a: 1i32, b: 42i32)"); } + + #[test] + fn test_replace_leaves_a_lambda_alone() -> VortexResult<()> { + // The lambda is in `list_transform`'s options, not among its children. + let expr = list_transform(col("l"), Lambda::new(["x"], checked_add(var("x"), lit(1)))?); + assert_eq!(replace(expr.clone(), &var("x"), lit(0)), expr); + assert_eq!( + replace(expr, &root(), col("s")).to_string(), + "vortex.list.transform($.s.l, (x) -> (@x + 1i32))" + ); + Ok(()) + } } From c26c935c0d7d6832879f3b9f55b11002f0136b5c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 11:28:16 +0000 Subject: [PATCH 5/7] list_transform: use rebuild_with_referenced_valid_elements, drop the Lists enum Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/list_transform/mod.rs | 226 +++++------------- .../src/scalar_fn/fns/list_transform/tests.rs | 27 ++- 2 files changed, 73 insertions(+), 180 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 390c9e3c7de..7f0f3a1094b 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -10,14 +10,10 @@ use std::fmt::Formatter; 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; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use vortex_mask::AllOr; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -26,25 +22,20 @@ use crate::IntoArray; use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::ListViewArray; -use crate::arrays::PrimitiveArray; use crate::arrays::list::ListArrayExt; use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewArraySlotsExt; -use crate::arrays::listview::ListViewRebuildMode; -use crate::arrays::listview::list_from_list_view; use crate::dtype::DType; use crate::dtype::FieldName; use crate::expr::BoundLambda; use crate::expr::Lambda; use crate::expr::display::ExprDisplay; -use crate::match_each_integer_ptype; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; 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. /// @@ -140,12 +131,10 @@ impl ScalarFnVTable for ListTransform { let ListTransformOptions::Bound(lambda) = options else { vortex_bail!("list_transform()'s lambda is not bound"); }; - 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())) + match args.get(0)?.try_downcast::() { + Ok(list) => transform_list(list, lambda, ctx), + Err(list) => transform_list_view(list.execute::(ctx)?, lambda, ctx), + } } fn is_strict(&self, _options: &ListTransformOptions) -> bool { @@ -185,167 +174,62 @@ fn element_dtype(list: &DType, params: &[FieldName]) -> VortexResult<[DType; 1]> Ok([element.as_ref().clone()]) } -/// 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), -} - -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)?; - match_each_integer_ptype!(offsets.ptype(), |O| { - let offsets = offsets.as_slice::(); - // The offsets are sorted, so a run of null lists is empty when it starts and ends at the - // same offset. - 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 lambda over a list's elements, inside the list's offsets and validity. +/// +/// A lambda that can fail runs only on the elements valid lists reference. One that can't runs on +/// every element from the first list to the last: results under null lists are never read. +fn transform_list( + list: ListArray, + lambda: &BoundLambda, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let list = if lambda.is_infallible() { + trim_ends(list, ctx)? + } else { + list.rebuild_with_referenced_valid_elements(ctx)? + }; + let values = lambda.evaluate(vec![list.elements().clone()], list.elements().len())?; + 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(), + ) } -/// The runs of null rows, as `(start, end)` ranges. -fn null_runs( - validity: &Validity, - len: usize, +/// The lambda over a list view's elements, inside the view's offsets, sizes and validity. +/// +/// A lambda that can fail runs only on the elements valid lists reference. +fn transform_list_view( + view: ListViewArray, + lambda: &BoundLambda, ctx: &mut ExecutionCtx, -) -> VortexResult> { - if matches!(validity, Validity::NonNullable | Validity::AllValid) { - return Ok(Vec::new()); +) -> VortexResult { + let view = if lambda.is_infallible() { + view + } else { + view.rebuild_with_referenced_valid_elements(ctx)? + }; + let values = lambda.evaluate(vec![view.elements().clone()], view.elements().len())?; + 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()) } - Ok(match (!&validity.execute_mask(len, ctx)?).slices() { - AllOr::All => vec![(0, len)], - AllOr::None => Vec::new(), - AllOr::Some(runs) => runs.to_vec(), - }) + .into_array()) } -/// 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(None); +/// The list without what comes before its first list and after its last. +fn trim_ends(list: ListArray, ctx: &mut ExecutionCtx) -> VortexResult { + if list.offset_at(0, ctx)? == 0 && list.offset_at(list.len(), ctx)? == list.elements().len() { + return Ok(list); } - let sizes = list.sizes().clone().execute::(ctx)?; - let sizes = match_each_integer_ptype!(sizes.ptype(), |T| { - let sizes = sizes.as_slice::(); - // Writers usually leave null lists empty already. - if nulls - .iter() - .all(|&(start, end)| sizes[start..end].iter().all(|size| *size == T::default())) - { - return Ok(None); - } - 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, - // which the rebuild must drop. - Ok(Some(unsafe { - ListViewArray::new_unchecked( - list.elements().clone(), - list.offsets().clone(), - sizes.into_array(), - validity, - ) - })) + list.reset_offsets(false, ctx) } 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 01b0fdf62cf..a93bc64edeb 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -7,11 +7,12 @@ use std::sync::LazyLock; use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_session::VortexSession; use super::ListTransform; use super::ListTransformOptions; -use super::Lists; +use super::transform_list; use crate::ArrayRef; use crate::IntoArray; use crate::RecursiveCanonical; @@ -20,10 +21,12 @@ use crate::arrays::BoolArray; use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -203,15 +206,21 @@ fn a_lambda_that_can_fail_sees_only_the_lists_elements( 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. +#[rstest] +#[case::cannot_fail(eq(var("x"), lit(3_i32)))] +#[case::can_fail(checked_add(var("x"), lit(1_i32)))] +fn the_lambda_runs_on_a_slice_without_its_leftovers(#[case] body: Expression) -> VortexResult<()> { + // Canonicalizing the result drops the leftovers too, so look at the transformed list itself. let mut ctx = SESSION.create_execution_ctx(); - let 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()) - })?; + let lambda = Lambda::new(["x"], body)? + .bind(&[DType::Primitive(PType::I32, Nullability::NonNullable)])?; + let list = sliced_list_with_leftovers()? + .try_downcast::() + .map_err(|_| vortex_err!("not a list"))?; + let result = transform_list(list, &lambda, &mut ctx)? + .try_downcast::() + .map_err(|_| vortex_err!("not a list"))?; + assert_eq!(result.elements().len(), 2); Ok(()) } From 6067b187445de2ed930b2e6629cf30e504492e5d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 16:37:41 +0000 Subject: [PATCH 6/7] use MakeExactWithEmptyNulls for list views Signed-off-by: Matt Katz --- vortex-array/src/scalar_fn/fns/list_transform/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 7f0f3a1094b..33be6d2aa68 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -26,6 +26,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::dtype::DType; use crate::dtype::FieldName; use crate::expr::BoundLambda; @@ -208,7 +209,7 @@ fn transform_list_view( let view = if lambda.is_infallible() { view } else { - view.rebuild_with_referenced_valid_elements(ctx)? + view.rebuild(ListViewRebuildMode::MakeExactWithEmptyNulls, ctx)? }; let values = lambda.evaluate(vec![view.elements().clone()], view.elements().len())?; vortex_ensure_eq!(values.len(), view.elements().len()); From c546c665fa9a167d506a089931c32c26a3a27bcd Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 16:59:05 +0000 Subject: [PATCH 7/7] use ListArray::compact Signed-off-by: Matt Katz --- vortex-array/src/scalar_fn/fns/list_transform/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 33be6d2aa68..c7dc2e77ae3 100644 --- a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -187,7 +187,7 @@ fn transform_list( let list = if lambda.is_infallible() { trim_ends(list, ctx)? } else { - list.rebuild_with_referenced_valid_elements(ctx)? + list.compact(ctx)? }; let values = lambda.evaluate(vec![list.elements().clone()], list.elements().len())?; vortex_ensure_eq!(values.len(), list.elements().len());