diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index d77a3472d30..3203d89725e 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -18,6 +18,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::Lambda; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::scalar_fn::EmptyOptions; @@ -45,6 +46,8 @@ use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_contains::ListContainsOptions; use crate::scalar_fn::fns::list_length::ListLength; use crate::scalar_fn::fns::list_sum::ListSum; +use crate::scalar_fn::fns::list_transform::ListTransform; +use crate::scalar_fn::fns::list_transform::ListTransformOptions; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::DuplicateHandling; @@ -153,8 +156,7 @@ pub fn bound_get_item(field: impl Into, child: BoundExpression) -> Bo /// Creates an expression that reads the parameter `name` of the enclosing lambda. /// -/// It is only meaningful in the body of a [`Lambda`](crate::expr::Lambda). Anywhere else it fails -/// to bind. +/// It is only meaningful in the body of a [`Lambda`]. Anywhere else it fails to bind. /// /// ```rust /// # use vortex_array::dtype::{DType, Nullability, PType}; @@ -1382,6 +1384,22 @@ pub fn bound_list_sum_opts( .vortex_expect("list-sum expressions require a numeric list child") } +// ---- ListTransform ---- + +/// Creates an expression that applies a lambda to every element of each list. +/// +/// The lambda takes one parameter, the element. +/// +/// ```rust +/// # use vortex_array::expr::{Lambda, checked_add, col, list_transform, lit, var}; +/// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap(); +/// let expr = list_transform(col("l"), lambda); +/// assert_eq!(expr.to_string(), "vortex.list.transform($.l, (x) -> (@x + 1i32))"); +/// ``` +pub fn list_transform(list: Expression, lambda: Lambda) -> Expression { + ListTransform.new_expr(ListTransformOptions::Lambda(lambda), [list]) +} + /// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit. /// /// `ambiguous` is a UTF-8 expression containing `raise`, `earliest`, `latest`, or `null`. diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 230dae1f25d..d848c166344 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -92,6 +92,7 @@ pub use exprs::list_contains_opts; pub use exprs::list_length; pub use exprs::list_sum; pub use exprs::list_sum_opts; +pub use exprs::list_transform; pub use exprs::lit; pub use exprs::lt; pub use exprs::lt_eq; diff --git a/vortex-array/src/scalar_fn/fns/list_transform/mod.rs b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs new file mode 100644 index 00000000000..390c9e3c7de --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_transform/mod.rs @@ -0,0 +1,351 @@ +// 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 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; +use crate::ExecutionCtx; +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. +/// +/// The lambda takes one parameter, the element. The only child is the list, of `List` dtype in any +/// encoding, and the result is a `List` of the same shape. A lambda that can fail runs only on the +/// elements of valid lists. +#[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"); + }; + 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 { + 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 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 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, 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); + } + 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, + ) + })) +} 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..7b6783acb53 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_transform/tests.rs @@ -0,0 +1,625 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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; +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::arrays::listview::ListViewArraySlotsExt; +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) +} + +/// `[[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])]` + 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;