From e265a1fd75a2c124cbf2e3179ee35e2e4bd1c7fa Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 1 Sep 2026 17:38:24 -0400 Subject: [PATCH] first pass Signed-off-by: Matt Katz --- .../src/arrays/list_transform/array.rs | 163 +++++++ vortex-array/src/arrays/list_transform/mod.rs | 14 + .../src/arrays/list_transform/rules.rs | 93 ++++ .../src/arrays/list_transform/template.rs | 112 +++++ .../src/arrays/list_transform/tests.rs | 441 ++++++++++++++++++ .../src/arrays/list_transform/vtable.rs | 438 +++++++++++++++++ vortex-array/src/arrays/mod.rs | 10 + vortex-array/src/arrays/template/input.rs | 198 ++++++++ .../src/arrays/template/instantiate.rs | 134 ++++++ vortex-array/src/arrays/template/mod.rs | 17 + .../src/expr/analysis/immediate_access.rs | 8 +- vortex-array/src/expr/analysis/infallible.rs | 3 + vortex-array/src/expr/analysis/strict.rs | 1 + vortex-array/src/expr/bound_expression.rs | 243 ++++++++-- vortex-array/src/expr/display.rs | 14 + vortex-array/src/expr/expression.rs | 52 ++- vortex-array/src/expr/exprs.rs | 5 + vortex-array/src/expr/mod.rs | 14 + vortex-array/src/expr/optimize.rs | 32 +- vortex-array/src/expr/proto.rs | 23 + vortex-array/src/expr/traversal/mod.rs | 12 +- vortex-array/src/expression.rs | 15 + 22 files changed, 1993 insertions(+), 49 deletions(-) create mode 100644 vortex-array/src/arrays/list_transform/array.rs create mode 100644 vortex-array/src/arrays/list_transform/mod.rs create mode 100644 vortex-array/src/arrays/list_transform/rules.rs create mode 100644 vortex-array/src/arrays/list_transform/template.rs create mode 100644 vortex-array/src/arrays/list_transform/tests.rs create mode 100644 vortex-array/src/arrays/list_transform/vtable.rs create mode 100644 vortex-array/src/arrays/template/input.rs create mode 100644 vortex-array/src/arrays/template/instantiate.rs create mode 100644 vortex-array/src/arrays/template/mod.rs diff --git a/vortex-array/src/arrays/list_transform/array.rs b/vortex-array/src/arrays/list_transform/array.rs new file mode 100644 index 00000000000..897ba55441e --- /dev/null +++ b/vortex-array/src/arrays/list_transform/array.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ArraySlots; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::EmptyArrayData; +use crate::array::TypedArrayRef; +use crate::arrays::ListTransform; +use crate::arrays::list_transform::template::build_template; +use crate::dtype::DType; +use crate::expr::BoundLambda; + +/// A lazy list transformation has structural children only: +/// `list`, a zero-length lambda `body`, then outer-row capture arrays. +pub trait ListTransformArrayExt: TypedArrayRef { + fn list(&self) -> &ArrayRef { + self.as_ref().slots()[0] + .as_ref() + .vortex_expect("validated ListTransformArray list slot") + } + + fn body(&self) -> &ArrayRef { + self.as_ref().slots()[1] + .as_ref() + .vortex_expect("validated ListTransformArray body slot") + } + + fn captures(&self) -> impl Iterator + '_ { + self.as_ref().slots()[2..].iter().map(|capture| { + capture + .as_ref() + .vortex_expect("validated ListTransformArray capture slot") + }) + } + + fn capture_count(&self) -> usize { + self.as_ref().slots().len() - 2 + } +} +impl> ListTransformArrayExt for T {} + +impl Array { + /// Build a structural list transform from a bound lambda and outer-row captures. + pub fn try_new( + list: ArrayRef, + lambda: BoundLambda, + captures: impl IntoIterator, + ) -> VortexResult { + let captures = captures.into_iter().collect::>(); + validate_lambda(&list, &lambda, &captures)?; + let body = build_template(&lambda)?; + Self::try_new_from_parts(list, body, captures) + } + + /// Rebuild a transform after substituting its outer template inputs. + pub(crate) fn try_new_from_parts( + list: ArrayRef, + body: ArrayRef, + captures: impl IntoIterator, + ) -> VortexResult { + let captures = captures.into_iter().collect::>(); + let dtype = output_dtype(list.dtype(), body.dtype())?; + vortex_ensure!( + body.is_empty(), + "ListTransformArray body must be a zero-length template, got {}", + body.len() + ); + vortex_ensure!( + captures.iter().all(|capture| capture.len() == list.len()), + "ListTransformArray captures must have the outer list length {}", + list.len() + ); + let len = list.len(); + let slots = std::iter::once(list) + .chain(std::iter::once(body)) + .chain(captures) + .map(Some) + .collect::(); + Array::try_from_parts( + ArrayParts::new(ListTransform, dtype, len, EmptyArrayData).with_slots(slots), + ) + } +} + +pub(crate) fn output_dtype(list: &DType, body: &DType) -> VortexResult { + match list { + DType::List(_, nullability) => Ok(DType::List(body.clone().into(), *nullability)), + DType::FixedSizeList(_, size, nullability) => Ok(DType::FixedSizeList( + body.clone().into(), + *size, + *nullability, + )), + _ => vortex_bail!("list_transform() requires List, ListView, or FixedSizeList, got {list}"), + } +} + +fn validate_lambda( + list: &ArrayRef, + lambda: &BoundLambda, + captures: &[ArrayRef], +) -> VortexResult<()> { + let element_dtype = match list.dtype() { + DType::List(element, _) | DType::FixedSizeList(element, ..) => element.as_ref(), + _ => vortex_bail!( + "list_transform() requires List, ListView, or FixedSizeList, got {}", + list.dtype() + ), + }; + 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 {}", + lambda.param_dtypes()[0], + element_dtype + ); + if lambda.param_dtypes().len() == 2 { + let index = DType::Primitive( + crate::dtype::PType::U64, + crate::dtype::Nullability::NonNullable, + ); + vortex_ensure!( + lambda.param_dtypes()[1] == index, + "list_transform() index parameter expects dtype {index}, got {}", + lambda.param_dtypes()[1] + ); + } + vortex_ensure!( + lambda.captures().len() == captures.len(), + "list_transform() lambda requires {} captures, got {}", + lambda.captures().len(), + captures.len() + ); + for (index, (capture, array)) in lambda.captures().iter().zip(captures).enumerate() { + vortex_ensure!( + capture.dtype() == array.dtype(), + "list_transform() capture {index} expects dtype {}, got {}", + capture.dtype(), + array.dtype() + ); + vortex_ensure!( + array.len() == list.len(), + "list_transform() capture {index} has length {}, expected {}", + array.len(), + list.len() + ); + } + vortex_ensure!( + lambda.body().is_root_bound_to(element_dtype), + "list_transform() lambda root expects a different dtype than {element_dtype}" + ); + Ok(()) +} diff --git a/vortex-array/src/arrays/list_transform/mod.rs b/vortex-array/src/arrays/list_transform/mod.rs new file mode 100644 index 00000000000..4ed4491a20e --- /dev/null +++ b/vortex-array/src/arrays/list_transform/mod.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +pub(crate) mod array; +mod rules; +mod template; +mod vtable; + +pub use array::ListTransformArrayExt; +pub use vtable::ListTransform; +pub use vtable::ListTransformArray; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/list_transform/rules.rs b/vortex-array/src/arrays/list_transform/rules.rs new file mode 100644 index 00000000000..83f6053ec2c --- /dev/null +++ b/vortex-array/src/arrays/list_transform/rules.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Filter; +use crate::arrays::ListTransform; +use crate::arrays::ListTransformArray; +use crate::arrays::ListTransformArrayExt; +use crate::arrays::Slice; +use crate::arrays::dict::TakeReduce; +use crate::arrays::dict::TakeReduceAdaptor; +use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleSet; + +pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&ListTransformSlicePushDown), + ParentRuleSet::lift(&ListTransformFilterPushDown), + ParentRuleSet::lift(&TakeReduceAdaptor(ListTransform)), +]); + +#[derive(Debug)] +struct ListTransformSlicePushDown; + +impl ArrayParentReduceRule for ListTransformSlicePushDown { + type Parent = Slice; + + fn reduce_parent( + &self, + transform: ArrayView<'_, ListTransform>, + parent: ArrayView<'_, Slice>, + _child_idx: usize, + ) -> VortexResult> { + let range = parent.slice_range().clone(); + Ok(Some( + ListTransformArray::try_new_from_parts( + transform.list().slice(range.clone())?, + transform.body().clone(), + transform + .captures() + .map(|capture| capture.slice(range.clone())) + .collect::>>()?, + )? + .into_array(), + )) + } +} + +#[derive(Debug)] +struct ListTransformFilterPushDown; + +impl ArrayParentReduceRule for ListTransformFilterPushDown { + type Parent = Filter; + + fn reduce_parent( + &self, + transform: ArrayView<'_, ListTransform>, + parent: ArrayView<'_, Filter>, + _child_idx: usize, + ) -> VortexResult> { + let mask = parent.filter_mask().clone(); + Ok(Some( + ListTransformArray::try_new_from_parts( + transform.list().filter(mask.clone())?, + transform.body().clone(), + transform + .captures() + .map(|capture| capture.filter(mask.clone())) + .collect::>>()?, + )? + .into_array(), + )) + } +} + +impl TakeReduce for ListTransform { + fn take(transform: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + Ok(Some( + ListTransformArray::try_new_from_parts( + transform.list().take(indices.clone())?, + transform.body().clone(), + transform + .captures() + .map(|capture| capture.take(indices.clone())) + .collect::>>()?, + )? + .into_array(), + )) + } +} diff --git a/vortex-array/src/arrays/list_transform/template.rs b/vortex-array/src/arrays/list_transform/template.rs new file mode 100644 index 00000000000..5cd84969c4f --- /dev/null +++ b/vortex-array/src/arrays/list_transform/template.rs @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::arrays::TemplateInputArray; +use crate::arrays::template::TemplateInputArrayExt; +use crate::arrays::template::TemplateScope; +use crate::expr::BoundExpression; +use crate::expr::BoundLambda; +use crate::scalar_fn::fns::literal::Literal; + +/// Convert one bound lambda into its zero-length, lexically scoped body template. +pub(super) fn build_template(lambda: &BoundLambda) -> VortexResult { + let scope = TemplateScope::fresh(); + // Slots 0 and 1 are permanently reserved for the element and optional local index. Captures + // always start at 2, which keeps a one-parameter lambda's first capture distinct from the + // index slot after the bound lambda itself has been discarded. + let mut inputs = vec![None; 2 + lambda.captures().len()]; + inputs[0] = + Some(TemplateInputArray::new(scope, 0, lambda.param_dtypes()[0].clone()).into_array()); + if lambda.param_dtypes().len() == 2 { + inputs[1] = + Some(TemplateInputArray::new(scope, 1, lambda.param_dtypes()[1].clone()).into_array()); + } + for (index, capture) in lambda.captures().iter().enumerate() { + inputs[index + 2] = + Some(TemplateInputArray::new(scope, index + 2, capture.dtype().clone()).into_array()); + } + build_expression(lambda.body(), lambda, scope, &inputs) +} + +fn build_expression( + expression: &BoundExpression, + lambda: &BoundLambda, + scope: TemplateScope, + inputs: &[Option], +) -> VortexResult { + match expression { + BoundExpression::Root { .. } => template_input(inputs, 0), + BoundExpression::Variable(variable) => { + let slot = lambda + .param_refs() + .iter() + .position(|reference| *reference == variable.variable_ref()) + .or_else(|| { + lambda + .captures() + .iter() + .position(|capture| capture.variable_ref() == variable.variable_ref()) + .map(|index| 2 + index) + }) + .ok_or_else(|| { + vortex_error::vortex_err!( + "variable '{}' is unresolved while building a template", + variable + ) + })?; + let input = template_input(inputs, slot)?; + vortex_ensure!( + input.as_::().scope() == scope, + "template builder mixed scopes" + ); + Ok(input) + } + BoundExpression::Scalar { + scalar_fn, + children, + .. + } => { + if let Some(value) = scalar_fn.as_opt::() { + return Ok(ConstantArray::new(value.clone(), 0).into_array()); + } + let children = children + .iter() + .map(|child| build_expression(child, lambda, scope, inputs)) + .try_collect()?; + Ok(ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, 0)?.into_array()) + } + BoundExpression::Lambda(_) => { + vortex_bail!("a detached lambda cannot appear in a template body") + } + BoundExpression::ListTransform { + lambda: nested_lambda, + children, + .. + } => { + let list = build_expression(&children[0], lambda, scope, inputs)?; + let captures = children[1..] + .iter() + .map(|capture| build_expression(capture, lambda, scope, inputs)) + .collect::>>()?; + crate::arrays::ListTransformArray::try_new(list, nested_lambda.clone(), captures) + .map(IntoArray::into_array) + } + } +} + +fn template_input(inputs: &[Option], slot: usize) -> VortexResult { + inputs + .get(slot) + .and_then(Option::as_ref) + .cloned() + .ok_or_else(|| vortex_error::vortex_err!("template input slot {slot} is not bound")) +} diff --git a/vortex-array/src/arrays/list_transform/tests.rs b/vortex-array/src/arrays/list_transform/tests.rs new file mode 100644 index 00000000000..16ba17c9ba1 --- /dev/null +++ b/vortex-array/src/arrays/list_transform/tests.rs @@ -0,0 +1,441 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_proto::expr as pb; + +use super::ListTransform; +use super::ListTransformArray; +use super::ListTransformArrayExt; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::ScalarFn; +use crate::arrays::TemplateInput; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::scalar_fn::ScalarFnArrayExt; +use crate::arrays::template::TemplateInputArrayExt; +use crate::arrays::template::instantiate; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::BoundLambda; +use crate::expr::Expression; +use crate::expr::Lambda; +use crate::expr::Scope; +use crate::expr::Variable; +use crate::expr::binary; +use crate::expr::lambda; +use crate::expr::list_length; +use crate::expr::list_transform as list_transform_expr; +use crate::expr::lit; +use crate::expr::proto::ExprSerializeProtoExt; +use crate::expr::root; +use crate::expr::var; +use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +fn list(elements: ArrayRef, offsets: ArrayRef) -> VortexResult { + list_with_validity(elements, offsets, Validity::NonNullable) +} + +fn list_with_validity( + elements: ArrayRef, + offsets: ArrayRef, + validity: Validity, +) -> VortexResult { + ListArray::try_new(elements, offsets, validity).map(IntoArray::into_array) +} + +fn bind_lambda( + list: &ArrayRef, + params: &[&str], + body: Expression, + captures: &[(&str, DType)], +) -> VortexResult { + let element = match list.dtype() { + DType::List(element, _) | DType::FixedSizeList(element, ..) => element.as_ref().clone(), + _ => unreachable!("test only constructs lists"), + }; + let parameter_dtypes = std::iter::once(element.clone()).chain( + (params.len() == 2).then_some(DType::Primitive(PType::U64, Nullability::NonNullable)), + ); + let scope = Scope::new(element) + .with_bindings( + captures + .iter() + .map(|(name, dtype)| (Variable::new(*name), dtype.clone())), + )? + .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 transform<'a>( + list: ArrayRef, + params: &[&str], + body: Expression, + captures: impl IntoIterator, +) -> VortexResult { + let captures = captures.into_iter().collect::>(); + let capture_dtypes = captures + .iter() + .map(|(name, capture)| (*name, capture.dtype().clone())) + .collect::>(); + let lambda = bind_lambda(&list, params, body, &capture_dtypes)?; + ListTransformArray::try_new( + list, + lambda, + captures.into_iter().map(|(_, capture)| capture), + ) + .map(IntoArray::into_array) +} + +#[test] +fn transform_is_structural_and_reifies_a_capture() -> VortexResult<()> { + let input = list( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 2, 3].into_array(), + )?; + let capture = buffer![10_i32, 20].into_array(); + let lambda = bind_lambda( + &input, + &["x"], + binary(Operator::Add, var("x"), var("offset")), + &[("offset", capture.dtype().clone())], + )?; + + let transform = ListTransformArray::try_new(input, lambda, [capture])?; + assert!(transform.as_ref().is::()); + assert_eq!(transform.body().len(), 0); + let body = transform.body().as_::(); + assert!(body.child_at(0).is::()); + assert!(body.child_at(1).is::()); + + let expected = list( + buffer![11_i32, 12, 23].into_array(), + buffer![0_u32, 2, 3].into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(transform, expected, &mut ctx); + Ok(()) +} + +#[test] +fn index_parameter_resets_per_list() -> VortexResult<()> { + let input = list( + buffer![10_u64, 10, 10, 10].into_array(), + buffer![0_u32, 2, 2, 4].into_array(), + )?; + let lambda = bind_lambda( + &input, + &["x", "i"], + binary(Operator::Add, var("x"), var("i")), + &[], + )?; + let transform = ListTransformArray::try_new(input, lambda, [])?; + let expected = list( + buffer![10_u64, 11, 10, 11].into_array(), + buffer![0_u32, 2, 2, 4].into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(transform, expected, &mut ctx); + Ok(()) +} + +#[test] +fn constant_scalar_tree_uses_the_invocation_length() -> VortexResult<()> { + let input = list( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 2, 3].into_array(), + )?; + let lambda = bind_lambda( + &input, + &["x"], + binary(Operator::Add, lit(2_i32), lit(3_i32)), + &[], + )?; + let transform = ListTransformArray::try_new(input, lambda, [])?; + let expected = list( + buffer![5_i32, 5, 5].into_array(), + buffer![0_u32, 2, 3].into_array(), + )?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(transform, expected, &mut ctx); + Ok(()) +} + +#[test] +fn null_containers_do_not_evaluate_hidden_elements() -> VortexResult<()> { + let validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); + let input = list_with_validity( + buffer![1_i32, 0, 4].into_array(), + buffer![0_u32, 1, 2, 3].into_array(), + validity.clone(), + )?; + let transform = transform( + input, + &["x"], + binary(Operator::Div, lit(8_i32), var("x")), + [], + )?; + let expected = list_with_validity( + buffer![8_i32, 2].into_array(), + buffer![0_u32, 1, 1, 2].into_array(), + validity, + )?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(transform, expected, &mut ctx); + Ok(()) +} + +#[test] +fn nullable_elements_and_lazy_captures_are_preserved() -> VortexResult<()> { + let nullable = list( + PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]).into_array(), + buffer![0_u32, 3, 3].into_array(), + )?; + let nullable_transform = transform( + nullable, + &["x"], + binary(Operator::Add, var("x"), lit(1_i32)), + [], + )?; + let nullable_expected = list( + PrimitiveArray::from_option_iter([Some(2_i32), None, Some(4)]).into_array(), + buffer![0_u32, 3, 3].into_array(), + )?; + + let input = list( + buffer![0_u64, 1, 2, 3, 4].into_array(), + buffer![0_u32, 3, 3, 5].into_array(), + )?; + let lengths = input.clone().apply(&list_length(root()))?; + let capture_transform = transform( + input, + &["x"], + binary(Operator::Add, var("x"), var("lengths")), + [("lengths", lengths)], + )?; + let capture_expected = list( + buffer![3_u64, 4, 5, 5, 6].into_array(), + buffer![0_u32, 3, 3, 5].into_array(), + )?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(nullable_transform, nullable_expected, &mut ctx); + assert_arrays_eq!(capture_transform, capture_expected, &mut ctx); + Ok(()) +} + +#[test] +fn empty_and_all_null_fixed_size_domains_do_not_invoke_the_body() -> VortexResult<()> { + let empty = list( + PrimitiveArray::from_iter([0_i32; 0]).into_array(), + buffer![0_u32].into_array(), + )?; + let empty_transform = transform( + empty, + &["x"], + binary(Operator::Add, var("x"), lit(1_i32)), + [], + )?; + let empty_expected = list( + PrimitiveArray::from_iter([0_i32; 0]).into_array(), + buffer![0_u32].into_array(), + )?; + + let all_null = FixedSizeListArray::try_new( + buffer![0_i32, 0, 0, 0].into_array(), + 2, + Validity::AllInvalid, + 2, + )? + .into_array(); + let all_null_transform = transform( + all_null, + &["x"], + binary(Operator::Div, lit(8_i32), var("x")), + [], + )?; + let all_null_expected = FixedSizeListArray::try_new( + buffer![0_i32, 0, 0, 0].into_array(), + 2, + Validity::AllInvalid, + 2, + )? + .into_array(); + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(empty_transform, empty_expected, &mut ctx); + assert_arrays_eq!(all_null_transform, all_null_expected, &mut ctx); + Ok(()) +} + +#[test] +fn source_expression_round_trips_and_applies_to_a_structural_array() -> VortexResult<()> { + let expression = list_transform_expr( + root(), + lambda(["x"], binary(Operator::Add, var("x"), lit(1_i32)))?, + )?; + let encoded = expression.serialize_proto()?.encode_to_vec(); + let decoded = Expression::from_proto(&pb::Expr::decode(encoded.as_slice())?, &array_session())?; + assert_eq!(decoded, expression); + + let input = list( + buffer![1_i32, 2].into_array(), + buffer![0_u32, 2].into_array(), + )?; + let applied = input.apply(&expression)?; + assert!(applied.is::()); + let transform = applied.as_::(); + assert_eq!(transform.body().len(), 0); + assert!( + transform + .body() + .as_::() + .child_at(0) + .is::() + ); + Ok(()) +} + +#[test] +fn outer_row_rules_keep_the_template_body() -> VortexResult<()> { + let input = list( + buffer![1_i32, 2, 3].into_array(), + buffer![0_u32, 1, 2, 3].into_array(), + )?; + let capture = buffer![10_i32, 20, 30].into_array(); + let lambda = bind_lambda( + &input, + &["x"], + binary(Operator::Add, var("x"), var("offset")), + &[("offset", capture.dtype().clone())], + )?; + let transform = ListTransformArray::try_new(input, lambda, [capture])?.into_array(); + let original_body = transform.as_::().body().clone(); + + let sliced = transform.slice(1..3)?; + let filtered = transform.filter(Mask::from_iter([false, true, true]))?; + let taken = transform.take(buffer![2_u64, 0].into_array())?; + for rewritten in [sliced, filtered, taken] { + let rewritten = rewritten.as_::(); + assert!(ArrayRef::ptr_eq(rewritten.body(), &original_body)); + assert_eq!(rewritten.list().len(), 2); + assert_eq!(rewritten.captures().next().map(ArrayRef::len), Some(2)); + } + Ok(()) +} + +#[test] +fn fixed_size_and_overlapping_list_view_preserve_their_families() -> VortexResult<()> { + let fixed = FixedSizeListArray::try_new( + buffer![1_i32, 2, 3, 4].into_array(), + 2, + Validity::NonNullable, + 2, + )? + .into_array(); + let fixed_lambda = bind_lambda( + &fixed, + &["x"], + binary(Operator::Add, var("x"), lit(1_i32)), + &[], + )?; + let fixed_transform = ListTransformArray::try_new(fixed, fixed_lambda, [])?.into_array(); + let fixed_expected = FixedSizeListArray::try_new( + buffer![2_i32, 3, 4, 5].into_array(), + 2, + Validity::NonNullable, + 2, + )? + .into_array(); + + let view = 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 view_lambda = bind_lambda( + &view, + &["x"], + binary(Operator::Add, var("x"), lit(1_i32)), + &[], + )?; + let view_transform = ListTransformArray::try_new(view, view_lambda, [])?.into_array(); + let view_expected = list( + buffer![2_i32, 3, 3, 4].into_array(), + buffer![0_u32, 2, 4].into_array(), + )?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(fixed_transform, fixed_expected, &mut ctx); + assert_arrays_eq!(view_transform, view_expected, &mut ctx); + Ok(()) +} + +#[test] +fn nested_transform_reifies_only_the_outer_scope() -> VortexResult<()> { + let inner = list( + buffer![1_u64, 2, 3, 4].into_array(), + buffer![0_u32, 2, 3, 4].into_array(), + )?; + let input = list(inner, buffer![0_u32, 2, 3].into_array())?; + let inner_transform = list_transform_expr( + var("x"), + lambda( + ["y"], + binary(Operator::Add, var("y"), list_length(var("x"))), + )?, + )?; + let expression = list_transform_expr(root(), lambda(["x"], inner_transform)?)?; + let actual = input.clone().apply(&expression)?; + let outer = actual.as_::(); + let nested = outer.body().as_::(); + let outer_scope = nested.list().as_::().scope(); + let inner_scope = nested + .body() + .as_::() + .child_at(0) + .as_::() + .scope(); + assert_ne!(outer_scope, inner_scope); + + let reified = instantiate( + outer.body(), + outer_scope, + std::slice::from_ref(input.as_::().elements()), + )?; + let reified_nested = reified.as_::(); + assert!(!reified_nested.list().is::()); + assert!(ArrayRef::ptr_eq(reified_nested.body(), nested.body())); + + let expected_inner = list( + buffer![3_u64, 4, 4, 5].into_array(), + buffer![0_u32, 2, 3, 4].into_array(), + )?; + let expected = list(expected_inner, buffer![0_u32, 2, 3].into_array())?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/arrays/list_transform/vtable.rs b/vortex-array/src/arrays/list_transform/vtable.rs new file mode 100644 index 00000000000..49734c0728e --- /dev/null +++ b/vortex-array/src/arrays/list_transform/vtable.rs @@ -0,0 +1,438 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::array::output_dtype; +use super::rules::PARENT_RULES; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +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::ListTransformArrayExt; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::listview::ListViewRebuildMode; +use crate::arrays::scalar_fn::ScalarFnArrayExt; +use crate::arrays::template::TemplateInputArrayExt; +use crate::arrays::template::instantiate; +use crate::arrays::template::template_scope; +use crate::buffer::BufferHandle; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::matcher::Matcher; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; +use crate::serde::ArrayChildren; +use crate::validity::Validity; +use crate::vtable::NotSupported; + +/// A lazy structural list transformation. +pub type ListTransformArray = Array; + +#[derive(Clone, Debug)] +pub struct ListTransform; + +impl VTable for ListTransform { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = NotSupported; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.list-transform"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + slots.len() >= 2, + "ListTransformArray requires list and body slots, got {}", + slots.len() + ); + let list = slots[0] + .as_ref() + .ok_or_else(|| vortex_error::vortex_err!("ListTransformArray list slot is missing"))?; + let body = slots[1] + .as_ref() + .ok_or_else(|| vortex_error::vortex_err!("ListTransformArray body slot is missing"))?; + vortex_ensure!( + list.len() == len, + "ListTransformArray list length does not match outer length" + ); + vortex_ensure!( + body.is_empty(), + "ListTransformArray body must have length zero" + ); + vortex_ensure!( + output_dtype(list.dtype(), body.dtype())? == *dtype, + "ListTransformArray dtype does not match its list and body children" + ); + vortex_ensure!( + slots[2..] + .iter() + .all(|capture| capture.as_ref().is_some_and(|capture| capture.len() == len)), + "ListTransformArray captures must be present and match the outer length" + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("ListTransformArray has no buffers") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + match idx { + 0 => "list".to_string(), + 1 => "body".to_string(), + index if index < array.slots().len() => format!("capture[{}]", index - 2), + _ => vortex_panic!("ListTransformArray slot index {idx} out of bounds"), + } + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + // Template scopes are process-local. Expression serialization retains source lambdas, + // whereas persistent lazy-array serialization is deliberately deferred. + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("ListTransformArray is not serializable") + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let list = array.list().clone().execute_until::(ctx)?; + let captures = array.captures().cloned().collect::>(); + let body = array.body().clone(); + let result = if let Some(list) = list.as_opt::() { + execute_list(list.into_owned(), body, captures, ctx)? + } else if let Some(list) = list.as_opt::() { + execute_fixed_size_list(list.into_owned(), body, captures, ctx)? + } else if let Some(list) = list.as_opt::() { + execute_list_view(list.into_owned(), body, captures, ctx)? + } else { + unreachable!("AnyList only matches list encodings") + }; + Ok(ExecutionResult::done(result)) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl ValidityVTable for ListTransform { + fn validity(array: ArrayView<'_, ListTransform>) -> VortexResult { + array.list().validity() + } +} + +fn execute_list( + list: ListArray, + body: ArrayRef, + 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( + body, + list.elements().clone(), + sizes, + list.list_validity(), + captures, + ctx, + )?; + ListArray::try_new(transformed, list.offsets().clone(), list.list_validity()) + .map(IntoArray::into_array) +} + +fn execute_fixed_size_list( + list: FixedSizeListArray, + body: ArrayRef, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let size = list.list_size(); + let sizes = ConstantArray::new(u64::from(size), list.len()).into_array(); + let transformed = transform_elements( + body, + list.elements().clone(), + sizes, + list.fixed_size_list_validity(), + captures, + ctx, + )?; + FixedSizeListArray::try_new( + transformed, + size, + list.fixed_size_list_validity(), + list.len(), + ) + .map(IntoArray::into_array) +} + +fn execute_list_view( + list: ListViewArray, + body: ArrayRef, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // Rebuild to a logical element domain. This duplicates overlaps and omits elements hidden by + // null containers, giving every invocation exactly one outer-row parent. + let list = list.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?; + let transformed = transform_elements( + body, + list.elements().clone(), + list.sizes().clone(), + list.listview_validity(), + captures, + ctx, + )?; + // SAFETY: rebuild produced sequential non-overlapping views and the transformed elements have + // exactly the same invocation domain. + Ok(unsafe { + ListViewArray::new_unchecked( + transformed, + list.offsets().clone(), + list.sizes().clone(), + list.listview_validity(), + ) + .with_zero_copy_to_list(true) + } + .into_array()) +} + +fn transform_elements( + body: ArrayRef, + elements: ArrayRef, + sizes: ArrayRef, + validity: Validity, + captures: Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let invocation_count = elements.len(); + let parents = parent_indices(sizes.clone(), invocation_count)?; + let invocation_mask = if matches!(validity, Validity::NonNullable | Validity::AllValid) { + None + } else { + let validity = validity.take(&parents)?; + let mask = validity.execute_mask(invocation_count, ctx)?; + (mask.true_count() != invocation_count).then_some(mask) + }; + + let elements = match &invocation_mask { + Some(mask) => elements.filter(mask.clone())?, + None => elements, + }; + let parents = match &invocation_mask { + Some(mask) => parents.filter(mask.clone())?, + None => parents, + }; + let mut inputs = vec![elements]; + if template_uses_slot(&body, 1)? { + let local = local_indices(sizes, invocation_count)?; + inputs.push(match &invocation_mask { + Some(mask) => local.filter(mask.clone())?, + None => local, + }); + } else { + // Capture slots begin at 2 even for a one-parameter lambda. This unused placeholder + // preserves that structural numbering without creating a local-index sequence. + inputs.push(ConstantArray::new(0_u64, inputs[0].len()).into_array()); + } + inputs.extend( + captures + .into_iter() + .map(|capture| capture.take(parents.clone())) + .collect::>>()?, + ); + + let transformed = match template_scope(&body)? { + Some(scope) => instantiate(&body, scope, &inputs)?, + None => instantiate_constant_body(&body, inputs.first().map_or(0, ArrayRef::len))?, + }; + match invocation_mask { + Some(mask) => scatter_valid_invocations(transformed, &mask, body.dtype()), + None => Ok(transformed), + } +} + +fn template_uses_slot(body: &ArrayRef, slot: usize) -> VortexResult { + fn visit(body: &ArrayRef, slot: usize, found: &mut bool) -> VortexResult<()> { + if let Some(input) = body.as_opt::() { + *found |= input.slot() == slot; + return Ok(()); + } + if let Some(scalar) = body.as_opt::() { + for child in scalar.iter_children() { + visit(child, slot, found)?; + } + } else if let Some(transform) = body.as_opt::() { + visit(transform.list(), slot, found)?; + for capture in transform.captures() { + visit(capture, slot, found)?; + } + } + Ok(()) + } + let mut found = false; + visit(body, slot, &mut found)?; + Ok(found) +} + +fn instantiate_constant_body(body: &ArrayRef, len: usize) -> VortexResult { + if let Some(constant) = body.as_opt::() { + return Ok(ConstantArray::new(constant.scalar().clone(), len).into_array()); + } + // A constant-only scalar tree still needs the invocation length propagated through every + // scalar-function node. The dummy input supplies that length; it cannot be observed because + // this body has no template inputs. + let invocation = ConstantArray::new(0_u8, len).into_array(); + instantiate( + body, + crate::arrays::TemplateScope::fresh(), + std::slice::from_ref(&invocation), + ) +} + +fn parent_indices(sizes: ArrayRef, element_count: usize) -> VortexResult { + let parents = sizes.len(); + let dtype = DType::Primitive( + sizes.dtype().as_ptype().to_unsigned(), + Nullability::NonNullable, + ); + let sizes = sizes.cast(dtype)?; + let starts = PrimitiveArray::from_iter((0..parents).map(|index| index as u64)).into_array(); + let multipliers = ConstantArray::new(0_u64, parents).into_array(); + PiecewiseSequenceArray::try_new(starts, sizes, multipliers, element_count) + .map(IntoArray::into_array) +} + +fn local_indices(sizes: ArrayRef, element_count: usize) -> VortexResult { + let parents = sizes.len(); + let dtype = DType::Primitive( + sizes.dtype().as_ptype().to_unsigned(), + Nullability::NonNullable, + ); + let sizes = sizes.cast(dtype)?; + let starts = ConstantArray::new(0_u64, parents).into_array(); + let multipliers = ConstantArray::new(1_u64, parents).into_array(); + PiecewiseSequenceArray::try_new(starts, sizes, multipliers, element_count) + .map(IntoArray::into_array) +} + +fn scatter_valid_invocations( + transformed: ArrayRef, + mask: &Mask, + dtype: &DType, +) -> VortexResult { + vortex_ensure!( + transformed.len() == mask.true_count(), + "list_transform() produced {} visible invocations, expected {}", + transformed.len(), + mask.true_count() + ); + let mut arrays = BufferMut::::with_capacity(mask.len()); + let mut rows = BufferMut::::with_capacity(mask.len()); + let mut visible = 0_u64; + for valid in mask.iter() { + arrays.push(if valid { 0 } else { 1 }); + rows.push(if valid { + let row = visible; + visible += 1; + row + } else { + 0 + }); + } + let placeholder = if dtype.is_nullable() { + Scalar::null(dtype.clone()) + } else { + Scalar::zero_value(dtype) + }; + InterleaveArray::try_new( + vec![transformed, ConstantArray::new(placeholder, 1).into_array()], + arrays.into_array(), + rows.into_array(), + ) + .map(IntoArray::into_array) +} + +struct AnyList; +impl Matcher for AnyList { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option> { + (array.is::() || array.is::() || array.is::()).then_some(()) + } +} diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index f96eebead65..215860dfc4e 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -78,6 +78,11 @@ pub mod listview; pub use listview::ListView; pub use listview::ListViewArray; +pub mod list_transform; +pub use list_transform::ListTransform; +pub use list_transform::ListTransformArray; +pub use list_transform::ListTransformArrayExt; + pub mod map; pub use map::Map; pub use map::MapArray; @@ -106,6 +111,11 @@ pub mod scalar_fn; pub use scalar_fn::ScalarFn; pub use scalar_fn::ScalarFnArray; +pub mod template; +pub use template::TemplateInput; +pub use template::TemplateInputArray; +pub use template::TemplateScope; + pub mod shared; pub use shared::Shared; pub use shared::SharedArray; diff --git a/vortex-array/src/arrays/template/input.rs b/vortex-array/src/arrays/template/input.rs new file mode 100644 index 00000000000..a51111705d8 --- /dev/null +++ b/vortex-array/src/arrays/template/input.rs @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::EqMode; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::TypedArrayRef; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::serde::ArrayChildren; +use crate::validity::Validity; +use crate::vtable::NotSupported; + +/// An identity for one lexical template expansion. +/// +/// Template scopes are process-local implementation details. Lazy template arrays deliberately +/// have no persistent serialization representation. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct TemplateScope(u64); + +impl TemplateScope { + pub(crate) fn fresh() -> Self { + static NEXT_SCOPE: AtomicU64 = AtomicU64::new(0); + Self(NEXT_SCOPE.fetch_add(1, Ordering::Relaxed)) + } +} + +impl Display for TemplateScope { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +#[derive(Clone, Debug)] +pub struct TemplateInputData { + scope: TemplateScope, + slot: usize, +} + +impl Display for TemplateInputData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "scope: {}, slot: {}", self.scope, self.slot) + } +} + +impl ArrayHash for TemplateInputData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.scope.hash(state); + self.slot.hash(state); + } +} + +impl ArrayEq for TemplateInputData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.scope == other.scope && self.slot == other.slot + } +} + +/// A zero-length symbolic array input in a scoped template body. +pub type TemplateInputArray = Array; + +#[derive(Clone, Debug)] +pub struct TemplateInput; + +pub trait TemplateInputArrayExt: TypedArrayRef { + fn scope(&self) -> TemplateScope { + self.scope + } + + fn slot(&self) -> usize { + self.slot + } +} +impl> TemplateInputArrayExt for T {} + +impl Array { + /// Create one symbolic input. Template inputs intentionally have no physical values. + pub fn new(scope: TemplateScope, slot: usize, dtype: DType) -> Self { + unsafe { + Array::from_parts_unchecked(ArrayParts::new( + TemplateInput, + dtype, + 0, + TemplateInputData { scope, slot }, + )) + } + } +} + +impl VTable for TemplateInput { + type TypedArrayData = TemplateInputData; + type OperationsVTable = NotSupported; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.template-input"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + _dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + len == 0, + "TemplateInputArray must have length zero, got {len}" + ); + vortex_ensure!( + slots.is_empty(), + "TemplateInputArray must not have child slots" + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("TemplateInputArray has no buffers") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + vortex_panic!("TemplateInputArray slot index {idx} out of bounds") + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("TemplateInputArray is not serializable") + } + + fn execute(_array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + vortex_bail!( + "TemplateInputArray cannot be executed directly; instantiate its template first" + ) + } +} + +impl ValidityVTable for TemplateInput { + fn validity(_array: ArrayView<'_, TemplateInput>) -> VortexResult { + // Every template input is empty. AllValid works for nullable empty dtypes without + // pretending that it has material values. + Ok(Validity::AllValid) + } +} diff --git a/vortex-array/src/arrays/template/instantiate.rs b/vortex-array/src/arrays/template/instantiate.rs new file mode 100644 index 00000000000..c3f9a4a4cd7 --- /dev/null +++ b/vortex-array/src/arrays/template/instantiate.rs @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::ListTransformArrayExt; +use crate::arrays::ScalarFn; +use crate::arrays::ScalarFnArray; +use crate::arrays::scalar_fn::ScalarFnArrayExt; +use crate::arrays::template::TemplateInput; +use crate::arrays::template::TemplateInputArrayExt; +use crate::arrays::template::TemplateScope; + +/// Rebuild a template body in an invocation row domain. +/// +/// Only the encodings a lambda builder can emit are accepted. In particular, this intentionally +/// does not attempt generic arbitrary-array substitution. Nested list transforms are boundaries: +/// their list and capture children are rebuilt, while their own zero-length body is left sealed. +pub(crate) fn instantiate( + body: &ArrayRef, + scope: TemplateScope, + inputs: &[ArrayRef], +) -> VortexResult { + let invocation_len = inputs.first().map_or(0, ArrayRef::len); + vortex_ensure!( + inputs.iter().all(|input| input.len() == invocation_len), + "template invocation inputs must have a common length" + ); + instantiate_inner(body, scope, inputs, invocation_len) +} + +fn instantiate_inner( + body: &ArrayRef, + scope: TemplateScope, + inputs: &[ArrayRef], + invocation_len: usize, +) -> VortexResult { + if let Some(input) = body.as_opt::() { + vortex_ensure!( + input.scope() == scope, + "unresolved template input from a different template scope" + ); + let actual = inputs.get(input.slot()).ok_or_else(|| { + vortex_error::vortex_err!("template input slot {} is unresolved", input.slot()) + })?; + vortex_ensure!( + actual.dtype() == body.dtype(), + "template input slot {} expects dtype {}, got {}", + input.slot(), + body.dtype(), + actual.dtype() + ); + return Ok(actual.clone()); + } + + if let Some(constant) = body.as_opt::() { + return Ok(ConstantArray::new(constant.scalar().clone(), invocation_len).into_array()); + } + + if let Some(scalar_fn) = body.as_opt::() { + let children = scalar_fn + .iter_children() + .map(|child| instantiate_inner(child, scope, inputs, invocation_len)) + .collect::>>()?; + return Ok(ScalarFnArray::try_new_with_len( + scalar_fn.scalar_fn().clone(), + children, + invocation_len, + )? + .into_array()); + } + + if let Some(transform) = body.as_opt::() { + let list = instantiate_inner(transform.list(), scope, inputs, invocation_len)?; + let captures = transform + .captures() + .map(|capture| instantiate_inner(capture, scope, inputs, invocation_len)) + .collect::>>()?; + return crate::arrays::ListTransformArray::try_new_from_parts( + list, + transform.body().clone(), + captures, + ) + .map(IntoArray::into_array); + } + + vortex_bail!( + "unsupported symbolic encoding {} in a template body", + body.encoding_id() + ) +} + +/// Infer the outer scope represented by a template body. +/// +/// An all-constant body has no symbolic scope and needs no substitution. The walk deliberately +/// stops at a nested transform body so an inner lambda cannot be captured by its outer lambda. +pub(crate) fn template_scope(body: &ArrayRef) -> VortexResult> { + fn visit(body: &ArrayRef, found: &mut Option) -> VortexResult<()> { + if let Some(input) = body.as_opt::() { + if let Some(scope) = found { + vortex_ensure!( + *scope == input.scope(), + "template body contains inputs from more than one scope" + ); + } else { + *found = Some(input.scope()); + } + return Ok(()); + } + if let Some(scalar_fn) = body.as_opt::() { + for child in scalar_fn.iter_children() { + visit(child, found)?; + } + return Ok(()); + } + if let Some(transform) = body.as_opt::() { + visit(transform.list(), found)?; + for capture in transform.captures() { + visit(capture, found)?; + } + } + Ok(()) + } + + let mut scope = None; + visit(body, &mut scope)?; + Ok(scope) +} diff --git a/vortex-array/src/arrays/template/mod.rs b/vortex-array/src/arrays/template/mod.rs new file mode 100644 index 00000000000..e931cc7ef44 --- /dev/null +++ b/vortex-array/src/arrays/template/mod.rs @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scoped symbolic inputs used by lazy data-parallel templates. +//! +//! This is deliberately small: it supports substituting the scalar-function trees produced by a +//! lambda body. It is not a general protocol for rebuilding arbitrary array encodings. + +mod input; +mod instantiate; + +pub use input::TemplateInput; +pub use input::TemplateInputArray; +pub use input::TemplateInputArrayExt; +pub use input::TemplateScope; +pub(crate) use instantiate::instantiate; +pub(crate) use instantiate::template_scope; diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 64ec360b52a..7a1f9e8a8e5 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -42,7 +42,9 @@ pub fn make_free_field_annotator( ) -> impl AnnotationFn { move |expr: &Expression| match expr { Expression::Root => scope.names().iter().cloned().collect(), - Expression::Lambda(_) | Expression::Variable(_) => vec![], + Expression::Lambda(_) | Expression::Variable(_) | Expression::ListTransform { .. } => { + vec![] + } Expression::Scalar { scalar_fn, children, @@ -72,7 +74,9 @@ pub fn make_bound_free_field_annotator( ) -> impl AnnotationFn { move |expr: &BoundExpression| match expr { BoundExpression::Root { .. } => scope.names().iter().cloned().collect(), - BoundExpression::Lambda(_) | BoundExpression::Variable(_) => vec![], + BoundExpression::Lambda(_) + | BoundExpression::Variable(_) + | BoundExpression::ListTransform { .. } => vec![], BoundExpression::Scalar { scalar_fn, children, diff --git a/vortex-array/src/expr/analysis/infallible.rs b/vortex-array/src/expr/analysis/infallible.rs index 82a8c236a89..c565e5d5f30 100644 --- a/vortex-array/src/expr/analysis/infallible.rs +++ b/vortex-array/src/expr/analysis/infallible.rs @@ -19,6 +19,9 @@ pub fn label_infallible(expr: &Expression) -> BooleanLabels<'_> { Expression::Root => true, // Fallibility is determined by the enclosing HOF. Expression::Lambda(_) | Expression::Variable(_) => true, + // The lambda runs in a separate element domain. Do not claim an unbound transform is + // infallible merely because its outer list expression is. + Expression::ListTransform { .. } => false, }, |acc, &child| acc & child, ) diff --git a/vortex-array/src/expr/analysis/strict.rs b/vortex-array/src/expr/analysis/strict.rs index 008a6e1eaa7..f79dc3d7a10 100644 --- a/vortex-array/src/expr/analysis/strict.rs +++ b/vortex-array/src/expr/analysis/strict.rs @@ -17,6 +17,7 @@ pub fn label_strict(expr: &Expression) -> BooleanLabels<'_> { Expression::Root => true, // Strictness is determined by the enclosing HOF. Expression::Variable(_) | Expression::Lambda(_) => true, + Expression::ListTransform { .. } => false, }, |acc, &child| acc & child, ) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index e4cbce4f302..7aab380e9a6 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -16,6 +16,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; +use crate::arrays::list_transform::array::output_dtype; use crate::dtype::DType; use crate::expr::Expression; use crate::expr::Lambda; @@ -55,6 +56,16 @@ pub enum BoundExpression { /// A lambda is not independently executable; only an enclosing higher-order function may /// close it over captures and apply it to arguments. Lambda(BoundLambda), + /// A dedicated list transformation. Its ordinary children are the outer list followed by + /// capture expressions; its lambda body is a lexical boundary owned by `lambda`. + ListTransform { + /// The output list dtype. + dtype: DType, + /// The typed lambda, including its bound lexical body. + lambda: BoundLambda, + /// Slot 0 is the list input and remaining slots are captures. + children: Arc>, + }, /// The scope itself. Its dtype is the scope's root dtype. Root { /// The dtype this node evaluates to. @@ -101,6 +112,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 +147,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 +160,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 +181,11 @@ impl BoundLambda { &self.param_refs } + /// The outer lexical bindings read by this lambda body, in stable lexical order. + pub fn captures(&self) -> &[BoundVariable] { + &self.captures + } + /// The lexical frame containing the parameters. pub fn parameter_frame(&self) -> usize { self.parameter_frame @@ -182,33 +203,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 +218,42 @@ 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, .. } + | BoundExpression::ListTransform { 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 reference = capture.variable_ref(); + (reference.frame(), reference.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) @@ -254,10 +288,27 @@ impl PartialEq for ExactBoundExpr { && lhs_dtype == rhs_dtype } (BoundExpression::Lambda(lhs), BoundExpression::Lambda(rhs)) => lhs == rhs, + ( + BoundExpression::ListTransform { + dtype: lhs_dtype, + lambda: lhs_lambda, + children: lhs_children, + }, + BoundExpression::ListTransform { + dtype: rhs_dtype, + lambda: rhs_lambda, + children: rhs_children, + }, + ) => { + lhs_dtype == rhs_dtype + && lhs_lambda == rhs_lambda + && Arc::ptr_eq(lhs_children, rhs_children) + } (BoundExpression::Variable(lhs), BoundExpression::Variable(rhs)) => lhs == rhs, (BoundExpression::Root { .. }, _) | (BoundExpression::Scalar { .. }, _) | (BoundExpression::Lambda(_), _) + | (BoundExpression::ListTransform { .. }, _) | (BoundExpression::Variable(_), _) => false, } } @@ -280,6 +331,13 @@ impl Hash for ExactBoundExpr { state.write_u8(3); lambda.hash(state); } + BoundExpression::ListTransform { + lambda, children, .. + } => { + state.write_u8(4); + lambda.hash(state); + Arc::as_ptr(children).hash(state); + } BoundExpression::Scalar { scalar_fn, children, @@ -332,6 +390,36 @@ impl BoundExpression { }) } + /// Create a typed dedicated list-transform node from its already-bound outer children. + pub(crate) fn try_new_list_transform( + list: BoundExpression, + lambda: BoundLambda, + captures: impl IntoIterator, + ) -> VortexResult { + let captures = captures.into_iter().collect::>(); + vortex_ensure!( + captures.len() == lambda.captures().len(), + "list_transform() lambda requires {} captures, got {}", + lambda.captures().len(), + captures.len() + ); + for (index, (capture, expected)) in captures.iter().zip(lambda.captures()).enumerate() { + vortex_ensure!( + capture.dtype() == expected.dtype(), + "list_transform() capture {index} expects dtype {}, got {}", + expected.dtype(), + capture.dtype() + ); + } + let dtype = output_dtype(list.dtype(), lambda.body_dtype())?; + let children = std::iter::once(list).chain(captures).collect(); + Ok(Self::ListTransform { + dtype, + lambda, + children: Arc::new(children), + }) + } + /// Rebuild this node with new bound children, recomputing its dtype. pub fn with_children( self, @@ -342,6 +430,12 @@ impl BoundExpression { BoundExpression::Scalar { scalar_fn, .. } => { Self::try_new_vec(scalar_fn.clone(), children) } + BoundExpression::ListTransform { lambda, .. } => { + let Some((list, captures)) = children.split_first() else { + vortex_bail!("list_transform() requires a list child"); + }; + Self::try_new_list_transform(list.clone(), lambda.clone(), captures.to_vec()) + } BoundExpression::Lambda(_) | BoundExpression::Root { .. } | BoundExpression::Variable(_) => { @@ -358,7 +452,9 @@ impl BoundExpression { /// The dtype this expression evaluates to. pub fn dtype(&self) -> &DType { match self { - Self::Scalar { dtype, .. } | Self::Root { dtype } => dtype, + Self::Scalar { dtype, .. } + | Self::ListTransform { dtype, .. } + | Self::Root { dtype } => dtype, Self::Lambda(lambda) => lambda.body_dtype(), Self::Variable(variable) => variable.dtype(), } @@ -369,7 +465,9 @@ impl BoundExpression { /// A bound lambda body is available through [`BoundLambda::body`] instead. pub fn children(&self) -> &[BoundExpression] { match self { - Self::Scalar { children, .. } => children.as_slice(), + Self::Scalar { children, .. } | Self::ListTransform { children, .. } => { + children.as_slice() + } Self::Lambda(_) | Self::Root { .. } | Self::Variable(_) => &[], } } @@ -383,7 +481,10 @@ impl BoundExpression { pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Lambda(_) | Self::Root { .. } | Self::Variable(_) => None, + Self::Lambda(_) + | Self::ListTransform { .. } + | Self::Root { .. } + | Self::Variable(_) => None, } } @@ -391,7 +492,20 @@ impl BoundExpression { pub fn as_lambda(&self) -> Option<&BoundLambda> { match self { Self::Lambda(lambda) => Some(lambda), - Self::Scalar { .. } | Self::Root { .. } | Self::Variable(_) => None, + Self::Scalar { .. } + | Self::ListTransform { .. } + | Self::Root { .. } + | Self::Variable(_) => None, + } + } + + /// The lambda owned by a list transform node, if this is one. + pub fn as_list_transform(&self) -> Option<(&BoundLambda, &[BoundExpression])> { + match self { + Self::ListTransform { + lambda, children, .. + } => Some((lambda, children)), + Self::Lambda(_) | Self::Scalar { .. } | Self::Root { .. } | Self::Variable(_) => None, } } @@ -404,7 +518,10 @@ impl BoundExpression { pub fn as_variable(&self) -> Option<&BoundVariable> { match self { Self::Variable(variable) => Some(variable), - Self::Lambda(_) | Self::Scalar { .. } | Self::Root { .. } => None, + Self::Lambda(_) + | Self::ListTransform { .. } + | Self::Scalar { .. } + | Self::Root { .. } => None, } } @@ -483,6 +600,9 @@ impl Display for BoundExpression { match self { Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), Self::Lambda(lambda) => Display::fmt(lambda, f), + Self::ListTransform { + lambda, children, .. + } => write!(f, "list_transform({}, {lambda})", children[0]), Self::Root { .. } => f.write_str("$"), Self::Variable(variable) => write!(f, "${variable}"), } @@ -516,6 +636,43 @@ impl Expression { Expression::Lambda(_) => { vortex_bail!("a lambda can be bound only as an argument to a higher-order function") } + Expression::ListTransform { children } => { + let list = children[0].bind_scope(scope)?; + let lambda = children[1].as_lambda().ok_or_else(|| { + vortex_error::vortex_err!("list_transform() requires a lambda") + })?; + let element_dtype = match list.dtype() { + DType::List(element, _) | DType::FixedSizeList(element, ..) => { + element.as_ref().clone() + } + dtype => vortex_bail!( + "list_transform() requires List, ListView, or FixedSizeList, got {dtype}" + ), + }; + vortex_ensure!( + matches!(lambda.params().len(), 1 | 2), + "list_transform() lambda must take one or two parameters, got {}", + lambda.params().len() + ); + let parameter_dtypes = std::iter::once(element_dtype.clone()).chain( + (lambda.params().len() == 2).then_some(DType::Primitive( + crate::dtype::PType::U64, + crate::dtype::Nullability::NonNullable, + )), + ); + let lambda_scope = scope + .clone() + .with_root(element_dtype) + .with_bindings(lambda.params().iter().cloned().zip(parameter_dtypes))?; + let lambda = BoundLambda::bind(lambda, &lambda_scope)?; + let captures = lambda + .captures() + .iter() + .cloned() + .map(BoundExpression::Variable) + .collect::>(); + BoundExpression::try_new_list_transform(list, lambda, captures) + } Expression::Scalar { scalar_fn, children, @@ -540,6 +697,16 @@ impl Drop for BoundExpression { to_drop.append(children); } } + Self::ListTransform { + lambda, children, .. + } => { + if let Some(children) = Arc::get_mut(children) { + to_drop.append(children); + } + if let Some(body) = lambda.take_body() { + to_drop.push(body); + } + } Self::Lambda(lambda) => { if let Some(body) = lambda.take_body() { to_drop.push(body); @@ -555,6 +722,16 @@ impl Drop for BoundExpression { to_drop.append(grandchildren); } } + BoundExpression::ListTransform { + lambda, children, .. + } => { + if let Some(grandchildren) = Arc::get_mut(children) { + to_drop.append(grandchildren); + } + if let Some(body) = lambda.take_body() { + to_drop.push(body); + } + } BoundExpression::Lambda(lambda) => { if let Some(body) = lambda.take_body() { to_drop.push(body); diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 2a151716877..762994ca55d 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -4,6 +4,7 @@ use std::fmt; use std::fmt::Display; use std::fmt::Formatter; +use std::sync::Arc; use vortex_utils::tree::TreeDisplayAdapter; use vortex_utils::tree::write_branch_tree; @@ -64,6 +65,7 @@ impl DisplayTreeNode for Expression { fn tree_children(&self) -> &[Self] { match self { Expression::Lambda(lambda) => std::slice::from_ref(lambda.body()), + Expression::ListTransform { .. } => Expression::children(self), Expression::Scalar { .. } | Expression::Root | Expression::Variable(_) => { Expression::children(self) } @@ -74,6 +76,11 @@ impl DisplayTreeNode for Expression { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), Expression::Lambda(_) => ChildName::from("body"), + Expression::ListTransform { .. } => match index { + 0 => ChildName::from("list"), + 1 => ChildName::from("lambda"), + _ => unreachable!("list transform has two children"), + }, Expression::Root | Expression::Variable(_) => { unreachable!("a leaf expression has no children") } @@ -84,6 +91,7 @@ impl DisplayTreeNode for Expression { match self { Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), Expression::Lambda(lambda) => fmt_lambda(lambda.params(), f), + Expression::ListTransform { .. } => write!(f, "vortex.list_transform()"), Expression::Root => write!(f, "{ROOT_DISPLAY}"), Expression::Variable(variable) => write!(f, "vortex.var({variable})"), } @@ -95,6 +103,7 @@ impl DisplayTreeNode for BoundExpression { match self { BoundExpression::Lambda(lambda) => std::slice::from_ref(lambda.body()), BoundExpression::Scalar { .. } + | BoundExpression::ListTransform { .. } | BoundExpression::Root { .. } | BoundExpression::Variable(_) => BoundExpression::children(self), } @@ -104,6 +113,10 @@ impl DisplayTreeNode for BoundExpression { match self { BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), BoundExpression::Lambda(_) => ChildName::from("body"), + BoundExpression::ListTransform { .. } => match index { + 0 => ChildName::from("list"), + index => ChildName::from(Arc::::from(format!("capture[{}]", index - 1))), + }, BoundExpression::Root { .. } | BoundExpression::Variable(_) => { unreachable!("a leaf bound expression has no children") } @@ -114,6 +127,7 @@ impl DisplayTreeNode for BoundExpression { match self { BoundExpression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), BoundExpression::Lambda(lambda) => fmt_lambda(lambda.params(), f), + BoundExpression::ListTransform { .. } => write!(f, "vortex.list_transform()"), BoundExpression::Root { .. } => write!(f, "{ROOT_DISPLAY}"), BoundExpression::Variable(variable) => write!(f, "vortex.var({variable})"), } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index dfe1ec326f2..522563d15c5 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -46,6 +46,11 @@ pub enum Expression { }, /// Lambda syntax owned and invoked by an enclosing higher-order function. Lambda(Lambda), + /// A dedicated list transformation with two syntax children: list input and lambda. + /// + /// This is intentionally not a scalar function: the lambda body executes in an element + /// domain, not the outer row domain. + ListTransform { children: Arc> }, /// The full scope of the expression evaluation. Root, /// A named value resolved from the surrounding scope when the expression is bound. @@ -77,6 +82,17 @@ impl Expression { }) } + /// Create the dedicated `list_transform(list, lambda)` syntax node. + pub fn try_new_list_transform(list: Expression, lambda: Expression) -> VortexResult { + vortex_ensure!( + lambda.is_lambda(), + "list_transform() requires a lambda as its second argument" + ); + Ok(Self::ListTransform { + children: Arc::new(vec![list, lambda]), + }) + } + /// Whether this expression is the scope root. pub fn is_root(&self) -> bool { matches!(self, Self::Root) @@ -86,7 +102,7 @@ impl Expression { pub fn as_variable(&self) -> Option<&Variable> { match self { Self::Variable(variable) => Some(variable), - Self::Lambda(_) | Self::Root | Self::Scalar { .. } => None, + Self::Lambda(_) | Self::ListTransform { .. } | Self::Root | Self::Scalar { .. } => None, } } @@ -94,7 +110,9 @@ impl Expression { pub fn as_lambda(&self) -> Option<&Lambda> { match self { Self::Lambda(lambda) => Some(lambda), - Self::Root | Self::Scalar { .. } | Self::Variable(_) => None, + Self::Root | Self::ListTransform { .. } | Self::Scalar { .. } | Self::Variable(_) => { + None + } } } @@ -103,11 +121,19 @@ impl Expression { self.as_lambda().is_some() } + /// Return the list input and lambda syntax of this node, if it is a list transform. + pub fn as_list_transform(&self) -> Option<(&Expression, &Lambda)> { + let Self::ListTransform { children } = self else { + return None; + }; + Some((&children[0], children[1].as_lambda()?)) + } + /// Returns the scalar fn for this expression, or `None` if it is not a scalar node. pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Lambda(_) | Self::Root | Self::Variable(_) => None, + Self::Lambda(_) | Self::ListTransform { .. } | Self::Root | Self::Variable(_) => None, } } @@ -136,7 +162,7 @@ impl Expression { /// A lambda body is binder-owned syntax and is available through [`Lambda::body`] instead. pub fn children(&self) -> &[Expression] { match self { - Self::Scalar { children, .. } => children.as_slice(), + Self::Scalar { children, .. } | Self::ListTransform { children } => children.as_slice(), Self::Lambda(_) | Self::Root | Self::Variable(_) => NO_CHILDREN, } } @@ -173,6 +199,15 @@ impl Expression { children: children.into(), }) } + Self::ListTransform { .. } => { + vortex_ensure!( + children.len() == 2 && children[1].is_lambda(), + "list_transform() requires exactly a list and lambda child" + ); + Ok(Self::ListTransform { + children: children.into(), + }) + } } } @@ -189,6 +224,7 @@ impl Expression { Self::Lambda(_) => vortex_bail!( "a lambda has no standalone dtype; it must be bound by a higher-order function" ), + Self::ListTransform { .. } => self.bind(root_dtype).map(|bound| bound.dtype().clone()), Self::Scalar { scalar_fn, children, @@ -214,6 +250,7 @@ impl Expression { Self::Lambda(_) => vortex_bail!( "a lambda has no standalone validity expression; it must be applied by a higher-order function" ), + Self::ListTransform { children } => children[0].validity(), Self::Scalar { scalar_fn, .. } => scalar_fn.validity(self), } } @@ -227,6 +264,9 @@ impl Expression { Self::Root => write!(f, "$"), Self::Variable(variable) => write!(f, "${variable}"), Self::Lambda(lambda) => Display::fmt(lambda, f), + Self::ListTransform { children } => { + write!(f, "list_transform({}, {})", children[0], children[1]) + } Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), } } @@ -351,7 +391,7 @@ impl Drop for Expression { fn drop(&mut self) { let mut children_to_drop = Vec::new(); match self { - Self::Scalar { children, .. } => { + Self::Scalar { children, .. } | Self::ListTransform { children } => { if let Some(children) = Arc::get_mut(children) { children_to_drop.append(children); } @@ -369,7 +409,7 @@ impl Drop for Expression { None => { while let Some(mut child) = children_to_drop.pop() { match &mut child { - Self::Scalar { children, .. } => { + Self::Scalar { children, .. } | Self::ListTransform { children } => { if let Some(expr_children) = Arc::get_mut(children) { children_to_drop.append(expr_children); } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 256872ad956..26208ac97f5 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -87,6 +87,11 @@ pub fn lambda( Ok(Lambda::try_new(params, body)?.into()) } +/// Creates a dedicated list transformation expression from a list input and lambda syntax. +pub fn list_transform(list: Expression, lambda: Expression) -> VortexResult { + Expression::try_new_list_transform(list, lambda) +} + /// Return whether the expression is a root expression. pub fn is_root(expr: &Expression) -> bool { expr.is_root() diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index d8d16aa695e..20c8ae9cf52 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -103,6 +103,7 @@ pub use exprs::list_contains; 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; @@ -171,6 +172,14 @@ impl PartialEq for ExactExpr { (Expression::Root, Expression::Root) => true, (Expression::Variable(lhs), Expression::Variable(rhs)) => lhs == rhs, (Expression::Lambda(lhs), Expression::Lambda(rhs)) => lhs == rhs, + ( + Expression::ListTransform { + children: lhs_children, + }, + Expression::ListTransform { + children: rhs_children, + }, + ) => Arc::ptr_eq(lhs_children, rhs_children), ( Expression::Scalar { scalar_fn: lhs_fn, @@ -184,6 +193,7 @@ impl PartialEq for ExactExpr { (Expression::Root, _) | (Expression::Scalar { .. }, _) | (Expression::Lambda(_), _) + | (Expression::ListTransform { .. }, _) | (Expression::Variable(_), _) => false, } } @@ -202,6 +212,10 @@ impl Hash for ExactExpr { state.write_u8(3); lambda.hash(state); } + Expression::ListTransform { children } => { + state.write_u8(4); + Arc::as_ptr(children).hash(state); + } Expression::Scalar { scalar_fn, children, diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index aecd2c74114..4fd389cf394 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -32,7 +32,10 @@ impl Expression { fn simplify_untyped_node(&self) -> VortexResult> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify_untyped(self), - Expression::Lambda(_) | Expression::Root | Expression::Variable(_) => Ok(None), + Expression::Lambda(_) + | Expression::ListTransform { .. } + | Expression::Root + | Expression::Variable(_) => Ok(None), } } @@ -40,7 +43,10 @@ impl Expression { fn simplify_node(&self, ctx: &dyn SimplifyCtx) -> VortexResult> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self, ctx), - Expression::Lambda(_) | Expression::Root | Expression::Variable(_) => Ok(None), + Expression::Lambda(_) + | Expression::ListTransform { .. } + | Expression::Root + | Expression::Variable(_) => Ok(None), } } @@ -51,7 +57,10 @@ impl Expression { ) -> VortexResult>> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node), - Expression::Lambda(_) | Expression::Root | Expression::Variable(_) => Ok(None), + Expression::Lambda(_) + | Expression::ListTransform { .. } + | Expression::Root + | Expression::Variable(_) => Ok(None), } } @@ -78,7 +87,8 @@ impl Expression { Expression::Lambda(_) => { vortex_bail!("cannot optimize a lambda outside a higher-order function") } - Expression::Root | Expression::Scalar { .. } => {} + Expression::Root | Expression::Scalar { .. } | Expression::ListTransform { .. } => { + } } let mut changed = false; @@ -141,6 +151,17 @@ impl Expression { &self, cache: &SimplifyCache<'_>, ) -> VortexResult> { + // The lambda body is a different lexical/domain boundary. Optimizing ordinary children + // may optimize the outer list input, but must not schedule a detached lambda body. + if let Expression::ListTransform { children } = self { + if let Some(list) = children[0].try_optimize_recursive_inner(cache)? { + return Ok(Some(Expression::try_new_list_transform( + list, + children[1].clone(), + )?)); + } + return Ok(None); + } // First optimize the root let mut current = self.try_optimize(cache)?; @@ -202,6 +223,9 @@ impl SimplifyCtx for SimplifyCache<'_> { Expression::Lambda(_) => vortex_bail!( "cannot determine the standalone dtype of a lambda; it must be bound by a higher-order function" ), + Expression::ListTransform { .. } => { + expr.bind(self.scope).map(|bound| bound.dtype().clone())? + } Expression::Scalar { scalar_fn, children, diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 8e76a70011f..537c390a223 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -32,6 +32,9 @@ pub(crate) const VARIABLE_ID: &str = "vortex.var"; /// The wire id for [`Expression::Lambda`]. pub(crate) const LAMBDA_ID: &str = "vortex.lambda"; +/// The wire id for the dedicated list-transform syntax node. +pub(crate) const LIST_TRANSFORM_ID: &str = "vortex.list_transform"; + impl Lambda { /// Serialize this lambda to its protobuf representation. fn serialize_proto(&self) -> VortexResult { @@ -127,6 +130,14 @@ impl ExprSerializeProtoExt for Expression { }), Expression::Variable(variable) => Ok(variable.serialize_proto()), Expression::Lambda(lambda) => lambda.serialize_proto(), + Expression::ListTransform { children } => Ok(pb::Expr { + id: LIST_TRANSFORM_ID.to_string(), + children: vec![ + children[0].serialize_proto()?, + children[1].serialize_proto()?, + ], + metadata: Some(vec![]), + }), Expression::Scalar { scalar_fn, children, @@ -155,6 +166,18 @@ impl Expression { return Ok(Lambda::from_proto(expr, session)?.into()); } + if expr.id == LIST_TRANSFORM_ID { + vortex_ensure!( + expr.children.len() == 2, + "list_transform() must have a list and lambda child, got {}", + expr.children.len() + ); + return Expression::try_new_list_transform( + Expression::from_proto(&expr.children[0], session)?, + Expression::from_proto(&expr.children[1], session)?, + ); + } + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let expr_id = ScalarFnId::new(expr.id.as_str()); let children = expr diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index 839f8d7622f..3e04bff56b1 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -534,7 +534,8 @@ impl Node for BoundExpression { mut f: F, ) -> VortexResult { let children = match self { - BoundExpression::Scalar { children, .. } => children, + BoundExpression::Scalar { children, .. } + | BoundExpression::ListTransform { children, .. } => children, BoundExpression::Lambda(_) | BoundExpression::Root { .. } | BoundExpression::Variable(_) => { @@ -557,7 +558,8 @@ impl Node for BoundExpression { mut f: F, ) -> VortexResult> { let children = match &self { - BoundExpression::Scalar { children, .. } => children, + BoundExpression::Scalar { children, .. } + | BoundExpression::ListTransform { children, .. } => children, BoundExpression::Lambda(_) | BoundExpression::Root { .. } | BoundExpression::Variable(_) => { @@ -602,7 +604,8 @@ impl Node for BoundExpression { fn iter_children(&self, f: impl FnOnce(&mut dyn Iterator) -> T) -> T { match self { - BoundExpression::Scalar { children, .. } => f(&mut children.iter()), + BoundExpression::Scalar { children, .. } + | BoundExpression::ListTransform { children, .. } => f(&mut children.iter()), BoundExpression::Lambda(_) | BoundExpression::Root { .. } | BoundExpression::Variable(_) => f(&mut std::iter::empty()), @@ -611,7 +614,8 @@ impl Node for BoundExpression { fn children_count(&self) -> usize { match self { - BoundExpression::Scalar { children, .. } => children.len(), + BoundExpression::Scalar { children, .. } + | BoundExpression::ListTransform { children, .. } => children.len(), BoundExpression::Lambda(_) | BoundExpression::Root { .. } | BoundExpression::Variable(_) => 0, diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d5cd410b63c..d5e26886d12 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -8,6 +8,7 @@ use vortex_error::vortex_bail; use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ConstantArray; +use crate::arrays::ListTransformArray; use crate::arrays::ScalarFnArray; use crate::expr::BoundExpression; use crate::expr::Expression; @@ -23,6 +24,16 @@ impl ArrayRef { BoundExpression::Lambda(_) => { vortex_bail!("cannot apply a lambda outside a higher-order function") } + BoundExpression::ListTransform { + lambda, children, .. + } => { + let list = self.clone().apply_bound(&children[0])?; + let captures = children[1..] + .iter() + .map(|capture| self.clone().apply_bound(capture)) + .collect::>>()?; + Ok(ListTransformArray::try_new(list, lambda.clone(), captures)?.into_array()) + } BoundExpression::Variable(variable) => { vortex_bail!("cannot apply variable '{variable}' without a provided value") } @@ -41,6 +52,10 @@ impl ArrayRef { Expression::Lambda(_) => { vortex_bail!("cannot apply a lambda outside a higher-order function") } + Expression::ListTransform { .. } => { + let bound = expr.bind(self.dtype())?; + self.apply_bound(&bound) + } Expression::Variable(variable) => { vortex_bail!("cannot apply unbound variable '{variable}'") }