From c6f98c3bccd7624618c71f8bb1c1576c3544c7ae Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 31 Aug 2026 16:17:26 -0400 Subject: [PATCH] Add scalar list transform Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 78 +- vortex-array/src/expression.rs | 209 +++- .../src/scalar_fn/fns/list_transform.rs | 1083 +++++++++++++++++ vortex-array/src/scalar_fn/fns/mod.rs | 1 + vortex-array/src/scalar_fn/session.rs | 2 + 5 files changed, 1328 insertions(+), 45 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/list_transform.rs diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index e4cbce4f302..f2da83238d8 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -101,6 +101,7 @@ pub struct BoundLambda { params: Box<[Variable]>, param_dtypes: Box<[DType]>, param_refs: Box<[VariableRef]>, + captures: Box<[BoundVariable]>, parameter_frame: usize, body: Arc, } @@ -135,6 +136,9 @@ impl BoundLambda { "lambda parameters must be bound in the innermost lexical frame" ); + let body = lambda.body().bind_scope(scope)?; + let captures = collect_captures(&body, parameter_frame); + Ok(Self { params: lambda.params().into(), param_dtypes: parameter_bindings @@ -145,8 +149,9 @@ impl BoundLambda { .into_iter() .map(|(_, variable_ref)| variable_ref) .collect(), + captures, parameter_frame, - body: Arc::new(lambda.body().bind_scope(scope)?), + body: Arc::new(body), }) } @@ -165,6 +170,11 @@ impl BoundLambda { &self.param_refs } + /// The outer lexical bindings read by this lambda body. + pub fn captures(&self) -> &[BoundVariable] { + &self.captures + } + /// The lexical frame containing the parameters. pub fn parameter_frame(&self) -> usize { self.parameter_frame @@ -182,33 +192,10 @@ impl BoundLambda { /// The outer lexical bindings read by this lambda body. pub fn free_variables(&self) -> Vec { - fn collect( - expression: &BoundExpression, - parameter_frame: usize, - variables: &mut Vec, - ) { - match expression { - BoundExpression::Variable(variable) - if variable.variable_ref().frame() < parameter_frame - && !variables.contains(&variable.variable_ref()) => - { - variables.push(variable.variable_ref()); - } - BoundExpression::Scalar { children, .. } => { - for child in children.iter() { - collect(child, parameter_frame, variables); - } - } - BoundExpression::Lambda(_) - | BoundExpression::Root { .. } - | BoundExpression::Variable(_) => {} - } - } - - let mut variables = Vec::new(); - collect(&self.body, self.parameter_frame, &mut variables); - variables.sort_by_key(|variable_ref| (variable_ref.frame(), variable_ref.slot())); - variables + self.captures + .iter() + .map(BoundVariable::variable_ref) + .collect() } fn take_body(&mut self) -> Option { @@ -220,6 +207,41 @@ impl BoundLambda { } } +fn collect_captures(expression: &BoundExpression, parameter_frame: usize) -> Box<[BoundVariable]> { + fn collect( + expression: &BoundExpression, + parameter_frame: usize, + captures: &mut Vec, + ) { + match expression { + BoundExpression::Variable(variable) + if variable.variable_ref().frame() < parameter_frame + && !captures + .iter() + .any(|capture| capture.variable_ref() == variable.variable_ref()) => + { + captures.push(variable.clone()); + } + BoundExpression::Scalar { children, .. } => { + for child in children.iter() { + collect(child, parameter_frame, captures); + } + } + BoundExpression::Lambda(_) + | BoundExpression::Root { .. } + | BoundExpression::Variable(_) => {} + } + } + + let mut captures = Vec::new(); + collect(expression, parameter_frame, &mut captures); + captures.sort_by_key(|capture| { + let variable_ref = capture.variable_ref(); + (variable_ref.frame(), variable_ref.slot()) + }); + captures.into_boxed_slice() +} + impl Display for BoundLambda { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { write!(f, "({}) -> {}", self.params.iter().join(", "), self.body) diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d5cd410b63c..d44aef69b66 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -4,34 +4,34 @@ use itertools::Itertools; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; +use crate::dtype::FieldName; +use crate::dtype::Nullability; use crate::expr::BoundExpression; +use crate::expr::BoundLambda; use crate::expr::Expression; +use crate::expr::VariableRef; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::fns::pack::Pack; +use crate::scalar_fn::fns::pack::PackOptions; impl ArrayRef { /// Apply a bound expression to this array, producing a new array in constant time. pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult { - match expr { - BoundExpression::Root { .. } => Ok(self), - BoundExpression::Lambda(_) => { - vortex_bail!("cannot apply a lambda outside a higher-order function") - } - BoundExpression::Variable(variable) => { - vortex_bail!("cannot apply variable '{variable}' without a provided value") - } - BoundExpression::Scalar { - scalar_fn, - children, - .. - } => apply_bound_scalar_fn(self, scalar_fn, children), + BoundApplyCtx { + root: &self, + bindings: None, } + .apply(expr) } /// Apply the expression to this array, producing a new array in constant time. @@ -52,21 +52,158 @@ impl ArrayRef { } } +impl BoundLambda { + /// Apply this lambda to arrays in a common invocation row domain. + /// + /// Parameters and captures are packed into a lazy non-nullable struct. Variable projections + /// reduce through that pack to the original arrays, leaving only the lambda body's lazy scalar + /// function array tree. + pub fn apply( + &self, + root: ArrayRef, + parameters: &[ArrayRef], + captures: &[ArrayRef], + ) -> VortexResult { + vortex_ensure!( + parameters.len() == self.param_dtypes().len(), + "lambda takes {} parameters but was applied with {} arguments", + self.param_dtypes().len(), + parameters.len() + ); + vortex_ensure!( + captures.len() == self.captures().len(), + "lambda requires {} captures but was applied with {}", + self.captures().len(), + captures.len() + ); + vortex_ensure!( + self.body().is_root_bound_to(root.dtype()), + "lambda root expects a different dtype than {}", + root.dtype() + ); + + for (index, (expected_dtype, parameter)) in + self.param_dtypes().iter().zip(parameters).enumerate() + { + vortex_ensure!( + parameter.dtype() == expected_dtype, + "lambda parameter {index} expects dtype {expected_dtype}, got {}", + parameter.dtype() + ); + vortex_ensure!( + parameter.len() == root.len(), + "lambda parameter {index} has length {}, expected {}", + parameter.len(), + root.len() + ); + } + for (index, (capture, array)) in self.captures().iter().zip(captures).enumerate() { + vortex_ensure!( + array.dtype() == capture.dtype(), + "lambda capture {index} expects dtype {}, got {}", + capture.dtype(), + array.dtype() + ); + vortex_ensure!( + array.len() == root.len(), + "lambda capture {index} has length {}, expected {}", + array.len(), + root.len() + ); + } + + let names = self + .param_refs() + .iter() + .copied() + .chain(self.captures().iter().map(|capture| capture.variable_ref())) + .map(binding_name) + .collect::>() + .into(); + let fields = parameters.iter().chain(captures).cloned().collect(); + let bindings = ScalarFnArray::try_new_with_len( + Pack.bind(PackOptions { + names, + nullability: Nullability::NonNullable, + }), + fields, + root.len(), + )? + .into_array(); + + let result = BoundApplyCtx { + root: &root, + bindings: Some(&bindings), + } + .apply(self.body())?; + vortex_ensure!( + result.dtype() == self.body_dtype(), + "lambda produced dtype {}, expected {}", + result.dtype(), + self.body_dtype() + ); + vortex_ensure!( + result.len() == root.len(), + "lambda produced {} rows, expected {}", + result.len(), + root.len() + ); + Ok(result) + } +} + +struct BoundApplyCtx<'a> { + root: &'a ArrayRef, + bindings: Option<&'a ArrayRef>, +} + +impl BoundApplyCtx<'_> { + fn apply(&self, expr: &BoundExpression) -> VortexResult { + match expr { + BoundExpression::Root { .. } => Ok(self.root.clone()), + BoundExpression::Lambda(_) => { + vortex_bail!("cannot apply a lambda outside a higher-order function") + } + BoundExpression::Variable(variable) => { + let Some(bindings) = self.bindings else { + vortex_bail!("cannot apply variable '{variable}' without a provided value"); + }; + GetItem::try_new(bindings.clone(), binding_name(variable.variable_ref()))? + .into_array() + .optimize() + } + BoundExpression::Scalar { + scalar_fn, + children, + .. + } => apply_bound_scalar_fn(self, scalar_fn, children), + } + } +} + +fn binding_name(variable_ref: VariableRef) -> FieldName { + FieldName::from(format!( + "frame[{}].slot[{}]", + variable_ref.frame(), + variable_ref.slot() + )) +} + fn apply_bound_scalar_fn( - root: ArrayRef, + ctx: &BoundApplyCtx<'_>, scalar_fn: &ScalarFnRef, children: &[BoundExpression], ) -> VortexResult { if let Some(scalar) = scalar_fn.as_opt::() { - return Ok(ConstantArray::new(scalar.clone(), root.len()).into_array()); + return Ok(ConstantArray::new(scalar.clone(), ctx.root.len()).into_array()); } let children: Vec<_> = children .iter() - .map(|child| root.clone().apply_bound(child)) + .map(|child| ctx.apply(child)) .try_collect()?; let array = - ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, root.len())?.into_array(); + ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, ctx.root.len())?.into_array(); array.optimize() } @@ -92,12 +229,20 @@ fn apply_scalar_fn( mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use crate::ArrayRef; use crate::IntoArray; + use crate::arrays::ScalarFn; + use crate::arrays::scalar_fn::ScalarFnArrayExt; + use crate::expr::Lambda; use crate::expr::Scope; use crate::expr::Variable; + use crate::expr::binary; use crate::expr::lambda; use crate::expr::var; + use crate::scalar_fn::fns::binary::Binary; + use crate::scalar_fn::fns::operators::Operator; #[test] fn variable_application_requires_a_runtime_binding() -> VortexResult<()> { @@ -120,4 +265,34 @@ mod tests { assert!(root.apply(&expression).is_err()); Ok(()) } + + #[test] + fn bound_lambda_pack_is_eliminated() -> VortexResult<()> { + let parameter = buffer![1_i32, 2, 3].into_array(); + let capture = buffer![10_i32, 20, 30].into_array(); + let scope = Scope::new(parameter.dtype().clone()) + .with_bindings([(Variable::new("capture"), capture.dtype().clone())])? + .with_bindings([(Variable::new("x"), parameter.dtype().clone())])?; + + let identity = crate::expr::BoundLambda::bind(&Lambda::try_new(["x"], var("x"))?, &scope)? + .apply(parameter.clone(), std::slice::from_ref(¶meter), &[])?; + assert!(ArrayRef::ptr_eq(&identity, ¶meter)); + + let lambda = crate::expr::BoundLambda::bind( + &Lambda::try_new(["x"], binary(Operator::Add, var("x"), var("capture")))?, + &scope, + )?; + let result = lambda.apply( + parameter.clone(), + std::slice::from_ref(¶meter), + std::slice::from_ref(&capture), + )?; + let Some(result) = result.as_opt::() else { + vortex_bail!("bound lambda did not produce a ScalarFnArray"); + }; + assert!(result.scalar_fn().is::()); + assert!(ArrayRef::ptr_eq(result.child_at(0), ¶meter)); + assert!(ArrayRef::ptr_eq(result.child_at(1), &capture)); + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/list_transform.rs b/vortex-array/src/scalar_fn/fns/list_transform.rs new file mode 100644 index 00000000000..add59e6364e --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_transform.rs @@ -0,0 +1,1083 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::FixedSizeListArray; +use crate::arrays::InterleaveArray; +use crate::arrays::List; +use crate::arrays::ListArray; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::ScalarFnArray; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::listview::ListViewRebuildMode; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::BoundLambda; +use crate::matcher::Matcher; +use crate::scalar::Scalar; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +/// Lazily transform every element of a list with a bound lambda. +/// +/// The typed lambda is stored in scalar-function options. The first child is the list and all +/// remaining children are its captures in the outer row domain. The first lambda parameter is the +/// element; an optional second parameter is its zero-based index within the containing list. +#[derive(Clone)] +pub struct ListTransform; + +impl ListTransform { + /// Create a lazy list transformation. + pub fn try_new( + list: ArrayRef, + lambda: BoundLambda, + captures: impl IntoIterator, + ) -> VortexResult { + let children = std::iter::once(list).chain(captures).collect(); + ScalarFnArray::try_new(ListTransform.bind(lambda), children) + } +} + +impl ScalarFnVTable for ListTransform { + type Options = BoundLambda; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.transform"); + *ID + } + + fn arity(&self, lambda: &Self::Options) -> Arity { + Arity::Exact(1 + lambda.captures().len()) + } + + fn child_name(&self, _lambda: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("list"), + index => ChildName::from(Arc::from(format!("capture[{}]", index - 1))), + } + } + + fn return_dtype(&self, lambda: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let Some((list_dtype, capture_dtypes)) = arg_dtypes.split_first() else { + vortex_bail!("list_transform() requires a list argument"); + }; + match list_dtype { + DType::List(element_dtype, list_nullability) => { + validate_lambda(lambda, element_dtype, capture_dtypes)?; + Ok(DType::List( + Arc::new(lambda.body_dtype().clone()), + *list_nullability, + )) + } + DType::FixedSizeList(element_dtype, list_size, list_nullability) => { + validate_lambda(lambda, element_dtype, capture_dtypes)?; + Ok(DType::FixedSizeList( + Arc::new(lambda.body_dtype().clone()), + *list_size, + *list_nullability, + )) + } + _ => vortex_bail!("list_transform() requires List or FixedSizeList, got {list_dtype}"), + } + } + + fn execute( + &self, + lambda: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let captures = (1..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>>()?; + let input = input.execute_until::(ctx)?; + if let Some(list) = input.as_opt::() { + execute_list(lambda, list.into_owned(), captures, ctx) + } else if let Some(list) = input.as_opt::() { + execute_fixed_size_list(lambda, list.into_owned(), captures, ctx) + } else if let Some(list) = input.as_opt::() { + execute_list_view(lambda, list.into_owned(), captures, ctx) + } else { + unreachable!("AnyList matcher returned a non-list array") + } + } +} + +fn execute_list( + lambda: &BoundLambda, + list: ListArray, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let list = list.reset_offsets(false, ctx)?; + let offsets = list.offsets(); + let sizes = offsets + .slice(1..offsets.len())? + .binary(offsets.slice(0..list.len())?, Operator::Sub)?; + let transformed = transform_elements( + lambda, + list.elements().clone(), + sizes, + list.list_validity(), + captures, + lambda.body_dtype(), + ctx, + )?; + + Ok(ListArray::try_new(transformed, list.offsets().clone(), list.list_validity())?.into_array()) +} + +fn execute_fixed_size_list( + lambda: &BoundLambda, + list: FixedSizeListArray, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let list_size = list.list_size(); + let sizes = ConstantArray::new(u64::from(list_size), list.len()).into_array(); + let transformed = transform_elements( + lambda, + list.elements().clone(), + sizes, + list.fixed_size_list_validity(), + captures, + lambda.body_dtype(), + ctx, + )?; + + Ok(FixedSizeListArray::try_new( + transformed, + list_size, + list.fixed_size_list_validity(), + list.len(), + )? + .into_array()) +} + +fn execute_list_view( + lambda: &BoundLambda, + list: ListViewArray, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // Logical flattening removes overlaps and gives null lists empty element ranges, so each + // invocation has one unambiguous parent row. + let list = list.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?; + let transformed = transform_elements( + lambda, + list.elements().clone(), + list.sizes().clone(), + list.listview_validity(), + captures, + lambda.body_dtype(), + ctx, + )?; + + // SAFETY: the rebuilt view is zero-copyable to List and `transform_elements` preserves its + // element domain, so its offsets, sizes, and validity remain valid and zero-copyable. + Ok(unsafe { + ListViewArray::new_unchecked( + transformed, + list.offsets().clone(), + list.sizes().clone(), + list.listview_validity(), + ) + .with_zero_copy_to_list(true) + } + .into_array()) +} + +fn validate_lambda( + lambda: &BoundLambda, + element_dtype: &DType, + capture_dtypes: &[DType], +) -> VortexResult<()> { + vortex_ensure!( + matches!(lambda.param_dtypes().len(), 1 | 2), + "list_transform() lambda must take one or two parameters, got {}", + lambda.param_dtypes().len() + ); + vortex_ensure!( + &lambda.param_dtypes()[0] == element_dtype, + "list_transform() element parameter expects dtype {}, got {element_dtype}", + lambda.param_dtypes()[0] + ); + if lambda.param_dtypes().len() == 2 { + let index_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); + vortex_ensure!( + lambda.param_dtypes()[1] == index_dtype, + "list_transform() index parameter expects dtype {index_dtype}, got {}", + lambda.param_dtypes()[1] + ); + } + vortex_ensure!( + lambda.captures().len() == capture_dtypes.len(), + "list_transform() lambda requires {} captures, got {}", + lambda.captures().len(), + capture_dtypes.len() + ); + for (index, (capture, dtype)) in lambda.captures().iter().zip(capture_dtypes).enumerate() { + vortex_ensure!( + capture.dtype() == dtype, + "list_transform() capture {index} expects dtype {}, got {dtype}", + capture.dtype() + ); + } + vortex_ensure!( + lambda.body().is_root_bound_to(element_dtype), + "list_transform() lambda root expects a different dtype than {element_dtype}" + ); + Ok(()) +} + +fn transform_elements( + lambda: &BoundLambda, + elements: ArrayRef, + sizes: ArrayRef, + validity: Validity, + captures: Vec, + body_dtype: &DType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let invocation_count = elements.len(); + let parent_indices = parent_indices(sizes.clone(), invocation_count)?; + let invocation_mask = if matches!(validity, Validity::NonNullable | Validity::AllValid) { + None + } else { + let invocation_validity = validity.take(&parent_indices)?; + let invocation_mask = invocation_validity.execute_mask(invocation_count, ctx)?; + if invocation_mask.true_count() == invocation_count { + None + } else { + Some(invocation_mask) + } + }; + + let elements = if let Some(mask) = &invocation_mask { + elements.filter(mask.clone())? + } else { + elements + }; + let capture_indices = match &invocation_mask { + Some(mask) => parent_indices.filter(mask.clone())?, + None => parent_indices, + }; + let captures = captures + .into_iter() + .map(|capture| capture.take(capture_indices.clone())) + .collect::>>()?; + + let mut parameters = vec![elements]; + if lambda.param_dtypes().len() > 1 { + let local_indices = local_indices(sizes, invocation_count)?; + let local_indices = match &invocation_mask { + Some(mask) => local_indices.filter(mask.clone())?, + None => local_indices, + }; + parameters.push(local_indices); + } + let transformed = lambda.apply(parameters[0].clone(), ¶meters, &captures)?; + + match invocation_mask { + Some(mask) => scatter_valid_invocations(transformed, &mask, body_dtype), + None => Ok(transformed), + } +} + +fn parent_indices(sizes: ArrayRef, element_count: usize) -> VortexResult { + let parent_count = sizes.len(); + let sizes_dtype = DType::Primitive( + sizes.dtype().as_ptype().to_unsigned(), + Nullability::NonNullable, + ); + let sizes = sizes.cast(sizes_dtype)?; + let starts = + PrimitiveArray::from_iter((0..parent_count).map(|index| index as u64)).into_array(); + let multipliers = ConstantArray::new(0_u64, parent_count).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, sizes, multipliers, element_count)?.into_array()) +} + +fn local_indices(sizes: ArrayRef, element_count: usize) -> VortexResult { + let parent_count = sizes.len(); + let sizes_dtype = DType::Primitive( + sizes.dtype().as_ptype().to_unsigned(), + Nullability::NonNullable, + ); + let sizes = sizes.cast(sizes_dtype)?; + let starts = ConstantArray::new(0_u64, parent_count).into_array(); + let multipliers = ConstantArray::new(1_u64, parent_count).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, sizes, multipliers, element_count)?.into_array()) +} + +fn scatter_valid_invocations( + transformed: ArrayRef, + mask: &Mask, + body_dtype: &DType, +) -> VortexResult { + vortex_ensure!( + transformed.len() == mask.true_count(), + "list_transform() produced {} valid invocations, expected {}", + transformed.len(), + mask.true_count() + ); + + let mut array_indices = BufferMut::::with_capacity(mask.len()); + let mut row_indices = BufferMut::::with_capacity(mask.len()); + let mut valid_index = 0_u64; + for valid in mask.iter() { + if valid { + array_indices.push(0); + row_indices.push(valid_index); + valid_index += 1; + } else { + array_indices.push(1); + row_indices.push(0); + } + } + + let placeholder = if body_dtype.is_nullable() { + Scalar::null(body_dtype.clone()) + } else { + Scalar::zero_value(body_dtype) + }; + let default = ConstantArray::new(placeholder, 1).into_array(); + Ok(InterleaveArray::try_new( + vec![transformed, default], + array_indices.into_array(), + row_indices.into_array(), + )? + .into_array()) +} + +/// Matches a `List`, `ListView`, or `FixedSizeList` physical array. +struct AnyList; + +impl Matcher for AnyList { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option> { + (array.as_opt::().is_some() + || array.as_opt::().is_some() + || array.as_opt::().is_some()) + .then_some(()) + } +} + +#[cfg(test)] +mod tests { + //! Execution examples for `list_transform`. + //! + //! Every test documents the logical input, lambda, and expected output. The physical inputs cover + //! ordinary lists, overlapping list views, nullable lists and elements, captured arrays, lazy + //! capture expressions, and nested list trees. + + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::*; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::ListArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::ScalarFn; + use crate::arrays::scalar_fn::ScalarFnArrayExt; + use crate::assert_arrays_eq; + use crate::expr::Expression; + use crate::expr::Lambda; + use crate::expr::Scope; + use crate::expr::Variable; + use crate::expr::lit; + use crate::expr::var; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::fns::list_length::ListLength; + use crate::scalar_fn::fns::operators::Operator; + use crate::validity::Validity; + + fn binary(operator: Operator, lhs: Expression, rhs: Expression) -> Expression { + crate::expr::binary(operator, lhs, rhs) + } + + fn lit_i32(value: i32) -> Expression { + lit(value) + } + + fn list_element_dtype(dtype: &DType) -> VortexResult<&DType> { + match dtype { + DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => { + Ok(element_dtype) + } + dtype => vortex_bail!("test list lambda requires a list dtype, got {dtype}"), + } + } + + fn bind_lambda_for_list_dtype( + list_dtype: &DType, + params: &[&str], + body: Expression, + captures: &[(&str, DType)], + ) -> VortexResult { + vortex_ensure!( + matches!(params.len(), 1 | 2), + "test list lambda must have one or two parameters" + ); + let element_dtype = list_element_dtype(list_dtype)?; + let parameter_dtypes = std::iter::once(element_dtype.clone()) + .chain( + (params.len() == 2) + .then_some(DType::Primitive(PType::U64, Nullability::NonNullable)), + ) + .collect::>(); + let scope = Scope::new(element_dtype.clone()); + let scope = if captures.is_empty() { + scope + } else { + scope.with_bindings( + captures + .iter() + .map(|(name, dtype)| (Variable::new(name), dtype.clone())), + )? + }; + let scope = scope.with_bindings( + params + .iter() + .zip(parameter_dtypes) + .map(|(name, dtype)| (Variable::new(name), dtype)), + )?; + BoundLambda::bind(&Lambda::try_new(params.iter().copied(), body)?, &scope) + } + + fn bind_list_lambda( + list: &ArrayRef, + params: &[&str], + body: Expression, + captures: &[(&str, ArrayRef)], + ) -> VortexResult { + let capture_dtypes = captures + .iter() + .map(|(name, array)| (*name, array.dtype().clone())) + .collect::>(); + bind_lambda_for_list_dtype(list.dtype(), params, body, &capture_dtypes) + } + + fn list_transform<'a>( + list: ArrayRef, + params: &[&str], + body: Expression, + captures: impl IntoIterator, + ) -> VortexResult { + let captures = captures.into_iter().collect::>(); + let lambda = bind_list_lambda(&list, params, body, &captures)?; + ListTransform::try_new(list, lambda, captures.into_iter().map(|(_, array)| array)) + } + + fn list(elements: ArrayRef, offsets: ArrayRef, validity: Validity) -> VortexResult { + ListArray::try_new(elements, offsets, validity).map(IntoArray::into_array) + } + + fn fixed_size_list( + elements: ArrayRef, + list_size: u32, + validity: Validity, + len: usize, + ) -> VortexResult { + FixedSizeListArray::try_new(elements, list_size, validity, len).map(IntoArray::into_array) + } + + fn nested_list( + elements: ArrayRef, + inner_offsets: ArrayRef, + outer_offsets: ArrayRef, + ) -> VortexResult { + let inner = list(elements, inner_offsets, Validity::NonNullable)?; + list(inner, outer_offsets, Validity::NonNullable) + } + + fn assert_transform(transform: ScalarFnArray, expected: ArrayRef) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(transform, expected, &mut ctx); + Ok(()) + } + + fn assert_list_transform(transform: ScalarFnArray, expected: ArrayRef) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let actual = transform.into_array().execute_until::(&mut ctx)?; + assert!(actual.as_opt::().is_some()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + fn assert_fixed_size_list_transform( + transform: ScalarFnArray, + expected: ArrayRef, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let actual = transform.into_array().execute_until::(&mut ctx)?; + assert!(actual.as_opt::().is_some()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + fn assert_list_view_transform( + transform: ScalarFnArray, + expected: ArrayRef, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let actual = transform.into_array().execute_until::(&mut ctx)?; + assert!(actual.as_opt::().is_some()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// Input: `[[1, 2], [], [3]]` + /// + /// Lambda: `x -> x + 1` + /// + /// Output: `[[2, 3], [], [4]]` + #[test] + fn scalar_call() -> VortexResult<()> { + let body = binary(Operator::Add, var("x"), lit_i32(1)); + let input = list( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 2, 2, 3].into_array(), + Validity::NonNullable, + )?; + let lambda = bind_list_lambda(&input, &["x"], body, &[])?; + let transform = ListTransform::try_new(input, lambda.clone(), [])?; + assert_eq!(transform.child_count(), 1); + assert!(transform.as_ref().is::()); + assert_eq!(transform.scalar_fn().as_::(), &lambda); + + let expected = list( + buffer![2_i32, 3, 4].into_array(), + buffer![0_u32, 2, 2, 3].into_array(), + Validity::NonNullable, + )?; + assert_list_transform(transform, expected) + } + + /// Input: `[[10, 10], [], [10, 10, 10]]` + /// + /// Lambda: `(x, i) -> x + i`, where `i` is zero-based within each list. + /// + /// Output: `[[10, 11], [], [10, 11, 12]]` + #[test] + fn local_index_parameter_on_list() -> VortexResult<()> { + let body = binary(Operator::Add, var("x"), var("i")); + let input = list( + buffer![10_u64, 10, 10, 10, 10].into_array(), + buffer![0_u32, 2, 2, 5].into_array(), + Validity::NonNullable, + )?; + let expected = list( + buffer![10_u64, 11, 10, 11, 12].into_array(), + buffer![0_u32, 2, 2, 5].into_array(), + Validity::NonNullable, + )?; + + assert_list_transform(list_transform(input, &["x", "i"], body, [])?, expected) + } + + /// Input: `[[-2, 0, 3], [5]]` + /// + /// Lambda: `x -> (x * 2) + 3` + /// + /// Output: `[[-1, 3, 9], [13]]` + #[test] + fn nested_scalar_calls() -> VortexResult<()> { + let doubled = binary(Operator::Mul, var("x"), lit_i32(2)); + let body = binary(Operator::Add, doubled, lit_i32(3)); + let input = list( + buffer![-2_i32, 0, 3, 5].into_array(), + buffer![0_u32, 3, 4].into_array(), + Validity::NonNullable, + )?; + let expected = list( + buffer![-1_i32, 3, 9, 13].into_array(), + buffer![0_u32, 3, 4].into_array(), + Validity::NonNullable, + )?; + + assert_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: `[[1, 2], [], [3]]` + /// + /// Lambda: `_ -> 7` + /// + /// Output: `[[7, 7], [], [7]]` + #[test] + fn unused_parameter_and_literal_body() -> VortexResult<()> { + let input = list( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 2, 2, 3].into_array(), + Validity::NonNullable, + )?; + let expected = list( + buffer![7_i32, 7, 7].into_array(), + buffer![0_u32, 2, 2, 3].into_array(), + Validity::NonNullable, + )?; + + assert_transform(list_transform(input, &["_"], lit_i32(7), [])?, expected) + } + + /// Physical elements: `[1, 2, 3]` + /// + /// Input views: `[[1, 2], [2, 3]]` + /// + /// Lambda: `x -> x * 10` + /// + /// Output: `[[10, 20], [20, 30]]` + #[test] + fn overlapping_list_view_without_captures() -> VortexResult<()> { + let input = ListViewArray::new( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 1].into_array(), + buffer![2_u32, 2].into_array(), + Validity::NonNullable, + ) + .into_array(); + let body = binary(Operator::Mul, var("x"), lit_i32(10)); + let expected = list( + buffer![10_i32, 20, 20, 30].into_array(), + buffer![0_u32, 2, 4].into_array(), + Validity::NonNullable, + )?; + + assert_list_view_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Physical elements: `[10, 20, 30]` + /// + /// Input views: `[[10, 20], [20, 30]]` + /// + /// Lambda: `(x, i) -> x + i`, where each view receives local indices `[0, 1]`. + /// + /// Output: `[[10, 21], [20, 31]]` + #[test] + fn local_index_parameter_on_overlapping_list_view() -> VortexResult<()> { + let input = ListViewArray::new( + buffer![10_u64, 20, 30].into_array(), + buffer![0_u32, 1].into_array(), + buffer![2_u32, 2].into_array(), + Validity::NonNullable, + ) + .into_array(); + let body = binary(Operator::Add, var("x"), var("i")); + let expected = list( + buffer![10_u64, 21, 20, 31].into_array(), + buffer![0_u32, 2, 4].into_array(), + Validity::NonNullable, + )?; + + assert_list_view_transform(list_transform(input, &["x", "i"], body, [])?, expected) + } + + /// Physical elements: `[1, 2, 3]` + /// + /// Input views: `[[1, 2], [2, 3]]`; capture: `[10, 20]` + /// + /// Lambda: `x -> x + capture` + /// + /// Output: `[[11, 12], [22, 23]]` + #[test] + fn captures_are_spread_by_parent_occurrence() -> VortexResult<()> { + let body = binary(Operator::Add, var("x"), var("capture")); + let input = ListViewArray::new( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 1].into_array(), + buffer![2_u32, 2].into_array(), + Validity::NonNullable, + ) + .into_array(); + let capture = buffer![10_i32, 20].into_array(); + let transform = list_transform(input, &["x"], body, [("capture", capture)])?; + assert_eq!(transform.child_count(), 2); + + let expected = list( + buffer![11_i32, 12, 22, 23].into_array(), + buffer![0_u32, 2, 4].into_array(), + Validity::NonNullable, + )?; + assert_list_view_transform(transform, expected) + } + + /// Input: `[[0, 1, 2], [], [3, 4]]` + /// + /// Captured expression: `list_length(input) = [3, 0, 2]` + /// + /// Lambda: `x -> x + list_length(input)` + /// + /// Output: `[[3, 4, 5], [], [5, 6]]` + #[test] + fn lazy_scalar_function_capture() -> VortexResult<()> { + let input = list( + buffer![0_u64, 1, 2, 3, 4].into_array(), + buffer![0_u32, 3, 3, 5].into_array(), + Validity::NonNullable, + )?; + let lengths = ScalarFnArray::try_new( + ListLength.bind(EmptyOptions), + std::slice::from_ref(&input).to_vec(), + )? + .into_array(); + let body = binary(Operator::Add, var("x"), var("lengths")); + let expected = list( + buffer![3_u64, 4, 5, 5, 6].into_array(), + buffer![0_u32, 3, 3, 5].into_array(), + Validity::NonNullable, + )?; + + assert_list_transform( + list_transform(input, &["x"], body, [("lengths", lengths)])?, + expected, + ) + } + + /// Input: `[[1], null-containing-[0], [4]]` + /// + /// Lambda: `x -> 8 / x` + /// + /// Output: `[[8], null, [2]]` + /// + /// The physical zero belongs only to the null list and must not be evaluated. + #[test] + fn null_lists_do_not_evaluate_hidden_elements() -> VortexResult<()> { + let body = binary(Operator::Div, lit_i32(8), var("x")); + let validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); + let input = list( + buffer![1_i32, 0, 4].into_array(), + buffer![0_u32, 1, 2, 3].into_array(), + validity.clone(), + )?; + let expected = list( + buffer![8_i32, 2].into_array(), + buffer![0_u32, 1, 1, 2].into_array(), + validity, + )?; + + assert_list_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: fixed-size lists `[[1, 2], [3, 4]]`; capture: `[10, 20]` + /// + /// Lambda: `x -> x + capture` + /// + /// Output: fixed-size lists `[[11, 12], [23, 24]]` + #[test] + fn fixed_size_list_with_capture_preserves_encoding() -> VortexResult<()> { + let input = fixed_size_list( + buffer![1_i32, 2, 3, 4].into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let body = binary(Operator::Add, var("x"), var("capture")); + let expected = fixed_size_list( + buffer![11_i32, 12, 23, 24].into_array(), + 2, + Validity::NonNullable, + 2, + )?; + + assert_fixed_size_list_transform( + list_transform( + input, + &["x"], + body, + [("capture", buffer![10_i32, 20].into_array())], + )?, + expected, + ) + } + + /// Input: fixed-size lists `[[10, 10], [10, 10]]` + /// + /// Lambda: `(x, i) -> x + i`, where each row receives local indices `[0, 1]`. + /// + /// Output: fixed-size lists `[[10, 11], [10, 11]]` + #[test] + fn local_index_parameter_on_fixed_size_list() -> VortexResult<()> { + let input = fixed_size_list( + buffer![10_u64, 10, 10, 10].into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let body = binary(Operator::Add, var("x"), var("i")); + let expected = fixed_size_list( + buffer![10_u64, 11, 10, 11].into_array(), + 2, + Validity::NonNullable, + 2, + )?; + + assert_fixed_size_list_transform(list_transform(input, &["x", "i"], body, [])?, expected) + } + + /// Input: fixed-size lists `[[1], null-containing-[0], [4]]` + /// + /// Lambda: `x -> 8 / x` + /// + /// Output: fixed-size lists `[[8], null, [2]]` + /// + /// The lambda is evaluated only for valid outer rows, while the fixed-size physical shape is + /// retained. + #[test] + fn nullable_fixed_size_list_skips_hidden_elements() -> VortexResult<()> { + let body = binary(Operator::Div, lit_i32(8), var("x")); + let validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); + let input = fixed_size_list(buffer![1_i32, 0, 4].into_array(), 1, validity.clone(), 3)?; + let expected = fixed_size_list(buffer![8_i32, 0, 2].into_array(), 1, validity, 3)?; + + assert_fixed_size_list_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: two null fixed-size lists with physical zeros. + /// + /// Lambda: `x -> 8 / x` + /// + /// Output: `[null, null]` + /// + /// No lambda invocation is executed when every outer row is null. + #[test] + fn all_null_fixed_size_list_has_no_invocations() -> VortexResult<()> { + let body = binary(Operator::Div, lit_i32(8), var("x")); + let input = fixed_size_list( + buffer![0_i32, 0, 0, 0].into_array(), + 2, + Validity::AllInvalid, + 2, + )?; + let expected = fixed_size_list( + buffer![0_i32, 0, 0, 0].into_array(), + 2, + Validity::AllInvalid, + 2, + )?; + + assert_fixed_size_list_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: three fixed-size lists of width zero: `[[], [], []]` + /// + /// Lambda: `x -> x + 1` + /// + /// Output: `[[], [], []]` + #[test] + fn degenerate_fixed_size_list_preserves_encoding() -> VortexResult<()> { + let input = fixed_size_list( + PrimitiveArray::from_iter([0_i32; 0]).into_array(), + 0, + Validity::NonNullable, + 3, + )?; + let body = binary(Operator::Add, var("x"), lit_i32(1)); + let expected = fixed_size_list( + PrimitiveArray::from_iter([0_i32; 0]).into_array(), + 0, + Validity::NonNullable, + 3, + )?; + + assert_fixed_size_list_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: `[[1, null, 3], []]` + /// + /// Lambda: `x -> x + 1` + /// + /// Output: `[[2, null, 4], []]` + #[test] + fn nullable_elements_propagate_through_scalar_functions() -> VortexResult<()> { + let input = list( + PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]).into_array(), + buffer![0_u32, 3, 3].into_array(), + Validity::NonNullable, + )?; + let body = binary(Operator::Add, var("x"), lit_i32(1)); + let expected = list( + PrimitiveArray::from_option_iter([Some(2_i32), None, Some(4)]).into_array(), + buffer![0_u32, 3, 3].into_array(), + Validity::NonNullable, + )?; + + assert_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: an empty `List` with zero outer rows. + /// + /// Lambda: `x -> x + 1` + /// + /// Output: an empty `List` with zero outer rows. + #[test] + fn empty_outer_array() -> VortexResult<()> { + let input = list( + PrimitiveArray::from_iter([0_i32; 0]).into_array(), + buffer![0_u32].into_array(), + Validity::NonNullable, + )?; + let body = binary(Operator::Add, var("x"), lit_i32(1)); + let expected = list( + PrimitiveArray::from_iter([0_i32; 0]).into_array(), + buffer![0_u32].into_array(), + Validity::NonNullable, + )?; + + assert_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: `[[[1, 2], [3]], [[]]]` + /// + /// Lambda: `x -> list_length(x)` + /// + /// Output: `[[2, 1], [0]]` + #[test] + fn scalar_function_over_nested_list_elements() -> VortexResult<()> { + let input = nested_list( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 2, 3, 3].into_array(), + buffer![0_u32, 2, 3].into_array(), + )?; + let body = crate::expr::list_length(var("x")); + let expected = list( + buffer![2_u64, 1, 0].into_array(), + buffer![0_u32, 2, 3].into_array(), + Validity::NonNullable, + )?; + + assert_transform(list_transform(input, &["x"], body, [])?, expected) + } + + /// Input: `[[[1, 2], [3]], [[], [4, 5, 6]]]` + /// + /// Lambda: `x -> list_transform(x, y -> y + 1)` + /// + /// Output: `[[[2, 3], [4]], [[], [5, 6, 7]]]` + #[test] + fn nested_list_transform() -> VortexResult<()> { + let input = nested_list( + buffer![1_i32, 2, 3, 4, 5, 6].into_array(), + buffer![0_u32, 2, 3, 3, 6].into_array(), + buffer![0_u32, 2, 4].into_array(), + )?; + let inner_list_dtype = list_element_dtype(input.dtype())?; + let inner_body = binary(Operator::Add, var("y"), lit_i32(1)); + let inner_lambda = bind_lambda_for_list_dtype(inner_list_dtype, &["y"], inner_body, &[])?; + let outer_body = Expression::try_new(ListTransform.bind(inner_lambda), [var("x")])?; + let expected = nested_list( + buffer![2_i32, 3, 4, 5, 6, 7].into_array(), + buffer![0_u32, 2, 3, 3, 6].into_array(), + buffer![0_u32, 2, 4].into_array(), + )?; + + assert_transform(list_transform(input, &["x"], outer_body, [])?, expected) + } + + /// Input: `[[[1, 2], [3]], [[], [4, 5, 6]]]` + /// + /// Lambda: `x -> list_transform(x, y -> y + list_length(x))` + /// + /// Output: `[[[3, 4], [4]], [[], [7, 8, 9]]]` + #[test] + fn nested_lambda_captures_outer_parameter_expression() -> VortexResult<()> { + let input = nested_list( + buffer![1_u64, 2, 3, 4, 5, 6].into_array(), + buffer![0_u32, 2, 3, 3, 6].into_array(), + buffer![0_u32, 2, 4].into_array(), + )?; + let inner_list_dtype = list_element_dtype(input.dtype())?; + let length_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); + let inner_body = binary(Operator::Add, var("y"), var("length")); + let inner_lambda = bind_lambda_for_list_dtype( + inner_list_dtype, + &["y"], + inner_body, + &[("length", length_dtype)], + )?; + let outer_body = Expression::try_new( + ListTransform.bind(inner_lambda), + [var("x"), crate::expr::list_length(var("x"))], + )?; + let expected = nested_list( + buffer![3_u64, 4, 4, 7, 8, 9].into_array(), + buffer![0_u32, 2, 3, 3, 6].into_array(), + buffer![0_u32, 2, 4].into_array(), + )?; + + assert_transform(list_transform(input, &["x"], outer_body, [])?, expected) + } + + /// Input: `[[[1, 2], [3]], [[], [4, 5, 6]]]`; outer capture: `[10, 20]` + /// + /// Lambda: `x -> list_transform(x, y -> y + outer_capture)` + /// + /// Output: `[[[11, 12], [13]], [[], [24, 25, 26]]]` + #[test] + fn capture_is_spread_across_two_nested_list_domains() -> VortexResult<()> { + let input = nested_list( + buffer![1_i32, 2, 3, 4, 5, 6].into_array(), + buffer![0_u32, 2, 3, 3, 6].into_array(), + buffer![0_u32, 2, 4].into_array(), + )?; + let inner_list_dtype = list_element_dtype(input.dtype())?; + let capture_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let inner_body = binary(Operator::Add, var("y"), var("outer_capture")); + let inner_lambda = bind_lambda_for_list_dtype( + inner_list_dtype, + &["y"], + inner_body, + &[("outer_capture", capture_dtype)], + )?; + let outer_body = Expression::try_new( + ListTransform.bind(inner_lambda), + [var("x"), var("outer_capture")], + )?; + let expected = nested_list( + buffer![11_i32, 12, 13, 24, 25, 26].into_array(), + buffer![0_u32, 2, 3, 3, 6].into_array(), + buffer![0_u32, 2, 4].into_array(), + )?; + + assert_transform( + list_transform( + input, + &["x"], + outer_body, + [("outer_capture", buffer![10_i32, 20].into_array())], + )?, + expected, + ) + } +} diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index 087f540ce7b..c7703bbf70c 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -16,6 +16,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; diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 211227858ca..5d7e6ac3f1d 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -25,6 +25,7 @@ use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::list_contains::ListContains; 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::literal::Literal; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::Merge; @@ -76,6 +77,7 @@ impl Default for ScalarFnSession { this.register(ListContains); this.register(ListLength); this.register(ListSum); + this.register(ListTransform); this.register(Literal); this.register(Mask); this.register(Merge);