diff --git a/vortex-array/src/arrays/higher_order/mod.rs b/vortex-array/src/arrays/higher_order/mod.rs new file mode 100644 index 00000000000..5c4e884d021 --- /dev/null +++ b/vortex-array/src/arrays/higher_order/mod.rs @@ -0,0 +1,299 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! An array that applies a higher-order function to its children, with its lambdas bound. + +mod rules; + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; + +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayRef; +use crate::ArraySlots; +use crate::EqMode; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::ProbeState; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::array::TypedArrayRef; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::arrays::ConstantArray; +use crate::buffer::BufferHandle; +use crate::columnar::Columnar; +use crate::dtype::DType; +use crate::executor::ExecutionResult; +use crate::expr::BoundLambda; +use crate::higher_order_fn::HigherOrderFnRef; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::VecExecutionArgs; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// An array that applies a higher-order function to its children. +pub type HigherOrderArray = Array; + +/// The array vtable of [`HigherOrderArray`]. Its id is the function's. +#[derive(Clone, Debug)] +pub struct HigherOrder { + id: ScalarFnId, +} + +/// The function and its bound lambdas. +#[derive(Clone, Debug)] +pub struct HigherOrderData { + function: HigherOrderFnRef, + lambdas: Arc>, +} + +impl HigherOrderData { + /// The function. + pub fn function(&self) -> &HigherOrderFnRef { + &self.function + } + + /// The bound lambdas. + pub fn lambdas(&self) -> &[BoundLambda] { + &self.lambdas + } +} + +impl Display for HigherOrderData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "higher_order_fn: {} {}", + self.function, + self.lambdas.iter().format(", ") + ) + } +} + +impl ArrayHash for HigherOrderData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.function.hash(state); + self.lambdas.hash(state); + } +} + +impl ArrayEq for HigherOrderData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.function == other.function && self.lambdas == other.lambdas + } +} + +impl Array { + /// Apply a function, with its lambdas bound, to these children of `len` rows. + pub fn try_new( + function: HigherOrderFnRef, + lambdas: Arc>, + children: impl IntoIterator, + len: usize, + ) -> VortexResult { + let slots: ArraySlots = children.into_iter().map(Some).collect(); + vortex_ensure!( + function.arity().matches(slots.len()), + "{} requires {} children, got {}", + function.id(), + function.arity(), + slots.len() + ); + vortex_ensure!( + slots.iter().flatten().all(|child| child.len() == len), + "HigherOrderArray requires children of length {len}" + ); + let arg_dtypes = slots + .iter() + .flatten() + .map(|child| child.dtype().clone()) + .collect_vec(); + let dtype = function.return_dtype(&arg_dtypes, &lambdas)?; + let vtable = HigherOrder { id: function.id() }; + let data = HigherOrderData { function, lambdas }; + Ok( + unsafe { + Array::from_parts_unchecked(ArrayParts::new(vtable, dtype, len, data, slots)) + }, + ) + } +} + +/// Accessors for a [`HigherOrderArray`]. +pub trait HigherOrderArrayExt: TypedArrayRef { + /// The value children. + fn children(&self) -> Vec { + self.slots() + .iter() + .map(|slot| { + slot.as_ref() + .vortex_expect("HigherOrderArray child slot") + .clone() + }) + .collect() + } +} + +impl> HigherOrderArrayExt for T {} + +impl VTable for HigherOrder { + type TypedArrayData = HigherOrderData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + self.id + } + + fn validate( + &self, + data: &HigherOrderData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure_eq!(data.function.id(), self.id); + vortex_ensure!(slots.iter().all(Option::is_some)); + vortex_ensure!(slots.iter().flatten().all(|child| child.len() == len)); + let arg_dtypes = slots + .iter() + .flatten() + .map(|child| child.dtype().clone()) + .collect_vec(); + vortex_ensure_eq!( + data.function.return_dtype(&arg_dtypes, &data.lambdas)?, + *dtype + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("HigherOrderArray buffer index {idx} out of bounds") + } + + 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 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!("HigherOrderArray is not serializable") + } + + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + array.function.child_name(idx).as_ref().to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let args = VecExecutionArgs::new(array.children(), array.len()); + array + .function + .execute(&args, &array.lambdas, ctx) + .map(ExecutionResult::done) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + rules::PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for HigherOrder { + type ProbeState = (); + + fn probe_scalar( + state: &mut ProbeState<'_, HigherOrder>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::scalar_at(state.array(), index, ctx) + } + + fn scalar_at( + array: ArrayView<'_, HigherOrder>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let inputs: Vec<_> = array + .children() + .iter() + .map(|child| Ok(ConstantArray::new(child.execute_scalar(index, ctx)?, 1).into_array())) + .collect::>()?; + let result = + array + .function + .execute(&VecExecutionArgs::new(inputs, 1), &array.lambdas, ctx)?; + Ok(match result.execute::(ctx)? { + Columnar::Canonical(canonical) => canonical.into_array().execute_scalar(0, ctx)?, + Columnar::Constant(constant) => constant.scalar().clone(), + }) + } +} + +impl ValidityVTable for HigherOrder { + fn validity(view: ArrayView<'_, HigherOrder>) -> VortexResult { + if view.function.is_strict() { + // Valid where every child is. + return view + .children() + .iter() + .try_fold(Validity::from(view.dtype().nullability()), |acc, child| { + acc.and(child.validity()?) + }); + } + // Prototype: a function that isn't strict would need a lazy validity, like the one + // scalar function arrays build. + vortex_bail!( + "validity of a non-strict {} is not supported yet", + view.function.id() + ) + } +} diff --git a/vortex-array/src/arrays/higher_order/rules.rs b/vortex-array/src/arrays/higher_order/rules.rs new file mode 100644 index 00000000000..b93fbe93382 --- /dev/null +++ b/vortex-array/src/arrays/higher_order/rules.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Slice and filter pushdown for [`HigherOrderArray`], as scalar function +//! arrays have: selection comes before evaluation, so a lambda never runs on unselected rows. + +use std::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexResult; + +use super::HigherOrder; +use super::HigherOrderArray; +use super::HigherOrderArrayExt; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Filter; +use crate::arrays::Slice; +use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleSet; + +pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&HigherOrderFilterPushdownRule), + ParentRuleSet::lift(&HigherOrderSliceReduceRule), +]); + +#[derive(Debug)] +struct HigherOrderSliceReduceRule; + +impl ArrayParentReduceRule for HigherOrderSliceReduceRule { + type Parent = Slice; + + fn reduce_parent( + &self, + array: ArrayView<'_, HigherOrder>, + parent: ArrayView<'_, Slice>, + _child_idx: usize, + ) -> VortexResult> { + let range = parent.slice_range(); + let children: Vec<_> = array + .children() + .iter() + .map(|child| child.slice(range.clone())) + .try_collect()?; + Ok(Some( + HigherOrderArray::try_new( + array.function.clone(), + Arc::clone(&array.lambdas), + children, + range.len(), + )? + .into_array(), + )) + } +} + +#[derive(Debug)] +struct HigherOrderFilterPushdownRule; + +impl ArrayParentReduceRule for HigherOrderFilterPushdownRule { + type Parent = Filter; + + fn reduce_parent( + &self, + array: ArrayView<'_, HigherOrder>, + parent: ArrayView<'_, Filter>, + _child_idx: usize, + ) -> VortexResult> { + let children: Vec<_> = array + .children() + .iter() + .map(|child| child.filter(parent.filter_mask().clone())) + .try_collect()?; + Ok(Some( + HigherOrderArray::try_new( + array.function.clone(), + Arc::clone(&array.lambdas), + children, + parent.len(), + )? + .into_array(), + )) + } +} diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index ed6e272b2b1..33cb4404481 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -105,6 +105,9 @@ pub use primitive::PrimitiveArray; pub mod scalar_fn; pub use scalar_fn::ScalarFn; pub use scalar_fn::ScalarFnArray; +pub mod higher_order; +pub use higher_order::HigherOrder; +pub use higher_order::HigherOrderArray; pub mod slice; pub use slice::Slice; diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 6258c35649f..3484b1b2982 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -40,6 +40,11 @@ pub fn make_bound_free_field_annotator( scope: &StructFields, ) -> impl AnnotationFn { move |expr: &BoundExpression| { + // A higher-order node reads through its value children; its lambdas read only their + // parameters. + if expr.as_higher_order().is_some() { + return vec![]; + } let Some(scalar_fn) = expr.as_scalar() else { return scope.names().iter().cloned().collect(); }; diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 97666966b62..1c81a82215c 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -15,11 +15,16 @@ use vortex_error::vortex_ensure; use vortex_session::VortexSession; use crate::dtype::DType; +use crate::expr::BoundLambda; use crate::expr::Expression; +use crate::expr::Lambda; use crate::expr::display::DisplayTreeExpr; +use crate::expr::expression::unbound_variable; use crate::expr::scope::Scope; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; +use crate::higher_order_fn::HigherOrderFnRef; +use crate::higher_order_fn::rebind_if_needed; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExpressionReduceNode; use crate::scalar_fn::ReduceNode; @@ -28,6 +33,7 @@ use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::is_not_null::IsNotNull; +use crate::scalar_fn::union_child_validities; use crate::stats::rewrite::falsify; use crate::stats::rewrite::satisfy; @@ -57,6 +63,20 @@ pub enum BoundExpression { /// The dtype this node evaluates to. dtype: DType, }, + /// A higher-order function applied to bound value children, with its lambdas bound. + HigherOrder { + /// The dtype this node evaluates to. + dtype: DType, + /// The function and its options. + function: HigherOrderFnRef, + /// The bound value children, in argument order. + children: Arc>, + /// The lambdas as written, kept to display them and to rebind them when the children's + /// dtypes change. + lambdas: Arc>, + /// The lambdas bound to the children's dtypes. + bound: Arc>, + }, } /// A bound-expression wrapper that compares shared tree identity instead of structure. @@ -86,6 +106,27 @@ impl PartialEq for ExactBoundExpr { && Arc::ptr_eq(lhs_children, rhs_children) && lhs_dtype == rhs_dtype } + ( + BoundExpression::HigherOrder { + dtype: lhs_dtype, + function: lhs_fn, + children: lhs_children, + bound: lhs_bound, + .. + }, + BoundExpression::HigherOrder { + dtype: rhs_dtype, + function: rhs_fn, + children: rhs_children, + bound: rhs_bound, + .. + }, + ) => { + lhs_fn == rhs_fn + && Arc::ptr_eq(lhs_children, rhs_children) + && lhs_bound == rhs_bound + && lhs_dtype == rhs_dtype + } _ => false, } } @@ -108,6 +149,13 @@ impl Hash for ExactBoundExpr { scalar_fn.hash(state); Arc::as_ptr(children).hash(state); } + BoundExpression::HigherOrder { + function, children, .. + } => { + state.write_u8(2); + function.hash(state); + Arc::as_ptr(children).hash(state); + } } } } @@ -147,12 +195,77 @@ impl BoundExpression { }) } + /// Create a bound higher-order node: bind its lambdas against its children's dtypes. + pub fn try_new_higher_order( + function: HigherOrderFnRef, + children: impl IntoIterator, + lambdas: Arc>, + ) -> VortexResult { + let children = Vec::from_iter(children); + vortex_ensure!( + function.arity().matches(children.len()), + "Expression arity mismatch: {} expects {} children but got {}", + function.id(), + function.arity(), + children.len() + ); + let arg_dtypes = children + .iter() + .map(|child| child.dtype().clone()) + .collect_vec(); + let (bound, dtype) = function.bind_node(&arg_dtypes, &lambdas)?; + Ok(Self::HigherOrder { + dtype, + function, + children: children.into(), + lambdas, + bound: bound.into(), + }) + } + + /// The higher-order function and its bound lambdas, if this is a higher-order node. + pub fn as_higher_order(&self) -> Option<(&HigherOrderFnRef, &[BoundLambda])> { + match self { + Self::HigherOrder { + function, bound, .. + } => Some((function, bound.as_slice())), + _ => None, + } + } + /// Rebuild this node with new bound children, recomputing its dtype. pub fn with_children( self, children: impl IntoIterator, ) -> VortexResult { let children = Vec::from_iter(children); + if let BoundExpression::HigherOrder { + function, + lambdas, + bound, + .. + } = &self + { + vortex_ensure!( + function.arity().matches(children.len()), + "Expression arity mismatch: {} expects {} children but got {}", + function.id(), + function.arity(), + children.len() + ); + let arg_dtypes = children + .iter() + .map(|child| child.dtype().clone()) + .collect_vec(); + let (bound, dtype) = rebind_if_needed(function, &arg_dtypes, lambdas, bound)?; + return Ok(Self::HigherOrder { + dtype, + function: function.clone(), + children: children.into(), + lambdas: Arc::clone(lambdas), + bound, + }); + } let BoundExpression::Scalar { scalar_fn, .. } = &self else { vortex_ensure!( children.is_empty(), @@ -168,14 +281,18 @@ 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::Root { dtype } | Self::HigherOrder { dtype, .. } => { + dtype + } } } /// The bound children of this node, in argument order. Empty for [`BoundExpression::Root`]. pub fn children(&self) -> &[BoundExpression] { match self { - Self::Scalar { children, .. } => children.as_slice(), + Self::Scalar { children, .. } | Self::HigherOrder { children, .. } => { + children.as_slice() + } Self::Root { .. } => &[], } } @@ -189,7 +306,7 @@ impl BoundExpression { pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root { .. } => None, + Self::Root { .. } | Self::HigherOrder { .. } => None, } } @@ -256,6 +373,11 @@ impl BoundExpression { match self { // The scope is exactly as valid as itself. Self::Root { .. } => Ok(self.clone()), + // A strict function is valid where all its value children are. + Self::HigherOrder { function, .. } if function.is_strict() => { + Ok(union_child_validities(&ExpressionReduceNode::new(self))?.into_expression()) + } + Self::HigherOrder { .. } => IsNotNull.try_new_bound_expr(EmptyOptions, [self.clone()]), Self::Scalar { .. } => { let node = ExpressionReduceNode::new(self); Ok(match node.validity()? { @@ -290,6 +412,9 @@ impl Display for BoundExpression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + Self::HigherOrder { + function, lambdas, .. + } => function.fmt_sql(self, lambdas, f), Self::Root { .. } => f.write_str("$"), } } @@ -310,15 +435,28 @@ impl Expression { if self.is_root() { return Ok(BoundExpression::new_root(scope.root().clone())); } + if let Expression::Variable(name) = self { + return Err(unbound_variable(name)); + } let children: Vec<_> = self .children() .iter() .map(|child| child.bind_scope(scope)) .try_collect()?; + if let Expression::HigherOrder { + function, lambdas, .. + } = self + { + return BoundExpression::try_new_higher_order( + function.clone(), + children, + Arc::clone(lambdas), + ); + } let scalar_fn = self .as_scalar() - .vortex_expect("root was handled above, so this is a scalar node"); + .vortex_expect("leaves were handled above, so this is a scalar node"); BoundExpression::try_new(scalar_fn.clone(), children) } } @@ -326,7 +464,7 @@ impl Expression { /// Iterative drop to avoid stack overflows on deep trees. impl Drop for BoundExpression { fn drop(&mut self) { - let Self::Scalar { children, .. } = self else { + let (Self::Scalar { children, .. } | Self::HigherOrder { children, .. }) = self else { return; }; let Some(children) = Arc::get_mut(children) else { @@ -335,7 +473,8 @@ impl Drop for BoundExpression { let mut to_drop = std::mem::take(children); while let Some(mut child) = to_drop.pop() { - if let BoundExpression::Scalar { children, .. } = &mut child + if let BoundExpression::Scalar { children, .. } + | BoundExpression::HigherOrder { children, .. } = &mut child && let Some(grandchildren) = Arc::get_mut(children) { to_drop.append(grandchildren); diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index b22ae246522..7af7b070f1d 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -5,6 +5,7 @@ use std::fmt; use std::fmt::Display; use std::fmt::Formatter; +use itertools::Itertools; use vortex_utils::tree::TreeDisplayAdapter; use vortex_utils::tree::write_branch_tree; @@ -68,14 +69,20 @@ impl DisplayTreeNode for Expression { fn tree_child_name(&self, index: usize) -> ChildName { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), + Expression::HigherOrder { function, .. } => function.child_name(index), Expression::Root => unreachable!("the scope root has no children"), + Expression::Variable(_) => unreachable!("a variable has no children"), } } fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + Expression::HigherOrder { + function, lambdas, .. + } => write!(f, "{function} {}", lambdas.iter().format(", ")), Expression::Root => write!(f, "{ROOT_DISPLAY}"), + Expression::Variable(name) => write!(f, "@{name}"), } } } @@ -88,6 +95,7 @@ impl DisplayTreeNode for BoundExpression { fn tree_child_name(&self, index: usize) -> ChildName { match self { BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), + BoundExpression::HigherOrder { function, .. } => function.child_name(index), BoundExpression::Root { .. } => unreachable!("the scope root has no children"), } } @@ -95,6 +103,9 @@ impl DisplayTreeNode for BoundExpression { fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { BoundExpression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + BoundExpression::HigherOrder { + function, lambdas, .. + } => write!(f, "{function} {}", lambdas.iter().format(", ")), BoundExpression::Root { .. } => write!(f, "{ROOT_DISPLAY}"), } } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index d228f646363..82558d27fd3 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -10,14 +10,19 @@ use std::hash::Hash; use std::sync::Arc; use itertools::Itertools; +use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::expr::Lambda; use crate::expr::display::DisplayTreeExpr; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; +use crate::higher_order_fn::HigherOrderFnRef; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; @@ -26,10 +31,13 @@ const NO_CHILDREN: &[Expression] = &[]; /// A node in a Vortex expression tree. /// -/// Most nodes are a scalar function applied to child expressions. [`Expression::Root`] is the scope -/// itself: a language primitive rather than a registered function, because its dtype comes from the -/// scope rather than from children and it is not executable. A [`ScalarFnVTable`] can answer neither -/// of those, so `Root` is a variant instead. +/// Most nodes are a scalar function applied to child expressions. A higher-order node applies a +/// function that also takes lambdas, which it holds next to its children. The other two are leaves +/// that name a value rather than compute one: [`Expression::Root`], the scope itself, and +/// [`Expression::Variable`], a parameter of an enclosing lambda. They are language primitives +/// rather than registered functions, because their dtype comes from where they appear rather than +/// from children and they are not executable. A [`ScalarFnVTable`] can answer neither of those, so +/// they are variants instead. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Expression { /// A scalar function applied to child expressions. @@ -39,8 +47,22 @@ pub enum Expression { /// Any children of this expression. children: Arc>, }, + /// A higher-order function applied to value children and lambdas. + HigherOrder { + /// The function and its options. + function: HigherOrderFnRef, + /// The value children. + children: Arc>, + /// The lambdas, as written. + lambdas: Arc>, + }, /// The full scope of the expression evaluation. Root, + /// A parameter of the enclosing lambda, read by name. + /// + /// Binding a [`Lambda`] turns its variables into reads of the struct of + /// its inputs. Anywhere else a variable fails to bind, and it is not serializable yet. + Variable(FieldName), } impl Expression { @@ -64,6 +86,45 @@ impl Expression { }) } + /// Create a higher-order node from a function, its value children and its lambdas. + pub fn try_new_higher_order( + function: HigherOrderFnRef, + children: impl IntoIterator, + lambdas: impl IntoIterator, + ) -> VortexResult { + let children = Vec::from_iter(children); + let lambdas = Vec::from_iter(lambdas); + vortex_ensure!( + function.arity().matches(children.len()), + "Expression arity mismatch: {} expects {} children but got {}", + function.id(), + function.arity(), + children.len() + ); + vortex_ensure!( + lambdas.len() == function.lambda_count(), + "{} takes {} lambdas, got {}", + function.id(), + function.lambda_count(), + lambdas.len() + ); + Ok(Self::HigherOrder { + function, + children: children.into(), + lambdas: lambdas.into(), + }) + } + + /// The higher-order function and its lambdas, if this is a higher-order node. + pub fn as_higher_order(&self) -> Option<(&HigherOrderFnRef, &[Lambda])> { + match self { + Self::HigherOrder { + function, lambdas, .. + } => Some((function, lambdas.as_slice())), + _ => None, + } + } + /// Whether this expression is the scope root. pub fn is_root(&self) -> bool { matches!(self, Self::Root) @@ -73,7 +134,7 @@ impl Expression { pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root => None, + Self::HigherOrder { .. } | Self::Root | Self::Variable(_) => None, } } @@ -100,8 +161,10 @@ impl Expression { /// Returns the children of this expression. pub fn children(&self) -> &[Expression] { match self { - Self::Scalar { children, .. } => children.as_slice(), - Self::Root => NO_CHILDREN, + Self::Scalar { children, .. } | Self::HigherOrder { children, .. } => { + children.as_slice() + } + Self::Root | Self::Variable(_) => NO_CHILDREN, } } @@ -125,6 +188,14 @@ impl Expression { ); Ok(Self::Root) } + Self::Variable(name) => { + vortex_ensure!( + children.is_empty(), + "Expression arity mismatch: variable expects 0 children but got {}", + children.len() + ); + Ok(Self::Variable(name.clone())) + } Self::Scalar { scalar_fn, .. } => { vortex_ensure!( scalar_fn.signature().arity().matches(children.len()), @@ -137,6 +208,9 @@ impl Expression { children: children.into(), }) } + Self::HigherOrder { + function, lambdas, .. + } => Self::try_new_higher_order(function.clone(), children, lambdas.iter().cloned()), } } @@ -144,6 +218,7 @@ impl Expression { pub fn return_dtype(&self, scope: &DType) -> VortexResult { match self { Self::Root => Ok(scope.clone()), + Self::Variable(name) => Err(unbound_variable(name)), Self::Scalar { scalar_fn, children, @@ -154,6 +229,17 @@ impl Expression { .try_collect()?; scalar_fn.return_dtype(&dtypes) } + Self::HigherOrder { + function, + children, + lambdas, + } => { + let dtypes: Vec<_> = children + .iter() + .map(|c| c.return_dtype(scope)) + .try_collect()?; + Ok(function.bind_node(&dtypes, lambdas)?.1) + } } } @@ -164,7 +250,11 @@ impl Expression { pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Self::Root => write!(f, "$"), + Self::Variable(name) => write!(f, "@{name}"), Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + Self::HigherOrder { + function, lambdas, .. + } => function.fmt_sql(self, lambdas, f), } } @@ -249,6 +339,11 @@ impl Expression { } } +/// The error for a variable that no enclosing lambda declares. +pub(crate) fn unbound_variable(name: &FieldName) -> VortexError { + vortex_err!("variable @{name} is not a parameter of an enclosing lambda") +} + /// The default display implementation for expressions uses the 'SQL'-style format. impl Display for Expression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { @@ -286,7 +381,7 @@ impl Drop for DropDepthGuard { impl Drop for Expression { fn drop(&mut self) { - let Self::Scalar { children, .. } = self else { + let (Self::Scalar { children, .. } | Self::HigherOrder { children, .. }) = self else { return; }; let Some(children) = Arc::get_mut(children) else { @@ -302,7 +397,8 @@ impl Drop for Expression { Some(_guard) => drop(children_to_drop), None => { while let Some(mut child) = children_to_drop.pop() { - if let Self::Scalar { children, .. } = &mut child + if let Self::Scalar { children, .. } | Self::HigherOrder { children, .. } = + &mut child && 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 b0bed8565f6..98213ab2cff 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -18,6 +18,9 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::Lambda; +use crate::higher_order_fn::HigherOrderFnVTableExt; +use crate::higher_order_fn::fns::list_transform::ListTransform; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::scalar_fn::EmptyOptions; @@ -149,6 +152,24 @@ pub fn bound_get_item(field: impl Into, child: BoundExpression) -> Bo .vortex_expect("get-item expressions must reference a field in the child dtype") } +// ---- Variable ---- + +/// Creates an expression that reads the parameter `name` of the enclosing lambda. +/// +/// It is only meaningful in the body of a [`Lambda`]. Anywhere else it fails to bind. +/// +/// ```rust +/// # use vortex_array::dtype::{DType, Nullability, PType}; +/// # use vortex_array::expr::{checked_add, lit, var}; +/// let expr = checked_add(var("x"), lit(1_i32)); +/// assert_eq!(expr.to_string(), "(@x + 1i32)"); +/// // Outside a lambda body, nothing gives `x` a value. +/// assert!(expr.bind(&DType::Primitive(PType::I32, Nullability::NonNullable)).is_err()); +/// ``` +pub fn var(name: impl Into) -> Expression { + Expression::Variable(name.into()) +} + // ---- VariantGet ---- /// Creates an expression that extracts a path from a Variant expression. @@ -1363,6 +1384,24 @@ pub fn bound_list_sum_opts( .vortex_expect("list-sum expressions require a numeric list child") } +// ---- ListTransform ---- + +/// Creates an expression that applies a lambda to every element of each list. +/// +/// The lambda takes one parameter, the element. +/// +/// ```rust +/// # use vortex_array::expr::{Lambda, checked_add, col, list_transform, lit, var}; +/// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap(); +/// let expr = list_transform(col("l"), lambda); +/// assert_eq!(expr.to_string(), "vortex.list.transform($.l, (x) -> (@x + 1i32))"); +/// ``` +pub fn list_transform(list: Expression, lambda: Lambda) -> Expression { + ListTransform + .new_expr(EmptyOptions, [list], [lambda]) + .vortex_expect("list_transform() takes one list and one lambda") +} + /// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit. /// /// `ambiguous` is a UTF-8 expression containing `raise`, `earliest`, `latest`, or `null`. diff --git a/vortex-array/src/expr/lambda/mod.rs b/vortex-array/src/expr/lambda/mod.rs new file mode 100644 index 00000000000..7b27ed56c4a --- /dev/null +++ b/vortex-array/src/expr/lambda/mod.rs @@ -0,0 +1,293 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lambdas, the function arguments of higher-order functions such as `list_transform`. +//! +//! A [`Lambda`] is a list of parameter names and a body that reads them with +//! [`var`](crate::expr::var). A higher-order node holds its lambdas next to its value children. +//! When the node is bound, its function binds each lambda to its parameters' dtypes, once, and the +//! bound node keeps the [`BoundLambda`]s. +//! +//! Binding turns the lambda into an ordinary expression over a struct of its inputs, one field per +//! parameter: `@x` becomes `$.x`. To run it, the function puts the parameters' arrays in that +//! struct, such as one row per list element, and applies the bound body to it. The body reads +//! nothing but its parameters, so the function's children are everything it reads. + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::StructArray; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::expr::BoundExpression; +use crate::expr::Expression; +use crate::expr::get_item; +use crate::expr::root; +use crate::expr::traversal::NodeExt; +use crate::expr::traversal::Transformed; +use crate::expr::traversal::TraversalOrder; +use crate::expr::traversal::pre_order_visit_down; +use crate::scalar_fn::internal::row_count::RowCount; +use crate::validity::Validity; + +/// Parameter names and a body: the function argument of a higher-order scalar function, as +/// written. +/// +/// The body is an ordinary [`Expression`] that reads its parameters with +/// [`var`](crate::expr::var), and nothing else from outside the lambda. +/// +/// ```rust +/// # use vortex_array::expr::{Lambda, checked_add, lit, var}; +/// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap(); +/// assert_eq!(lambda.to_string(), "(x) -> (@x + 1i32)"); +/// ``` +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Lambda { + params: FieldNames, + body: Expression, +} + +impl Lambda { + /// Create a lambda with these parameters and body. + /// + /// Returns an error when a parameter repeats, or when the body reads from outside the lambda: + /// `$` or a column, another lambda's parameter, or the scan's `row_idx()` or `row_count()`. + pub fn new( + params: impl IntoIterator>, + body: Expression, + ) -> VortexResult { + let params: FieldNames = params.into_iter().map(Into::into).collect(); + if let Some(param) = params.iter().duplicates().next() { + vortex_bail!("lambda parameter {param} is declared more than once"); + } + pre_order_visit_down(&body, |node: &Expression| { + vortex_ensure!( + !reads_outside_lambda(node, params.as_ref()), + "lambda ({}) can only read its parameters, not {node}; reading anything else is \ + not supported yet", + params.iter().format(", ") + ); + Ok(TraversalOrder::Continue) + })?; + + Ok(Self { params, body }) + } + + /// The parameter names, in order. + pub fn params(&self) -> &[FieldName] { + self.params.as_ref() + } + + /// The body, as written. + pub fn body(&self) -> &Expression { + &self.body + } + + /// The dtype of the body, given the parameters' dtypes, without binding it. + pub fn return_dtype(&self, param_dtypes: &[DType]) -> VortexResult { + let inputs = inputs_dtype(self.param_fields(param_dtypes)?); + self.body_over_inputs().return_dtype(&inputs) + } + + /// Bind the body to these parameter dtypes: `@x` becomes `$.x`, a field of the struct of the + /// lambda's inputs. + /// + /// A higher-order function calls this from its + /// [`bind_lambdas`](crate::higher_order_fn::HigherOrderFnVTable::bind_lambdas), so the body is + /// bound once, with the node that holds it, and never while the function runs. + pub fn bind(&self, param_dtypes: &[DType]) -> VortexResult { + let params = self.param_fields(param_dtypes)?; + let body = self + .body_over_inputs() + .bind(&inputs_dtype(params.clone()))?; + Ok(BoundLambda { params, body }) + } + + /// Whether nothing in the body can fail, so that the body may run over values whose results + /// are never read, such as dictionary values that no row uses. + pub fn is_infallible(&self) -> bool { + let mut infallible = true; + pre_order_visit_down(&self.body, |node: &Expression| { + infallible &= node + .as_scalar() + .is_none_or(|scalar_fn| scalar_fn.signature().is_infallible()); + Ok(if infallible { + TraversalOrder::Continue + } else { + TraversalOrder::Stop + }) + }) + .vortex_expect("expression traversal cannot fail"); + infallible + } + + fn param_fields(&self, param_dtypes: &[DType]) -> VortexResult { + vortex_ensure!( + param_dtypes.len() == self.params.len(), + "lambda {self} needs one dtype per parameter, got {}", + param_dtypes.len() + ); + Ok(StructFields::new( + self.params.clone(), + param_dtypes.to_vec(), + )) + } + + /// The body over the struct of the lambda's inputs, where `@x` reads the field `x`. + /// + /// A lambda in the body is a function's options, which this leaves alone: its variables are its + /// own parameters. + fn body_over_inputs(&self) -> Expression { + self.body + .clone() + .transform_down(|node| { + Ok(match &node { + Expression::Variable(name) => Transformed { + value: get_item(name.clone(), root()), + changed: true, + order: TraversalOrder::Skip, + }, + _ => Transformed::no(node), + }) + }) + .vortex_expect("expression traversal cannot fail") + .into_inner() + } +} + +/// The struct of a lambda's inputs, one field per parameter. +fn inputs_dtype(params: StructFields) -> DType { + DType::Struct(params, Nullability::NonNullable) +} + +impl Display for Lambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "({}) -> {}", self.params.iter().format(", "), self.body) + } +} + +/// A [`Lambda`] bound to its parameters' dtypes, ready to run. +/// +/// The body is a [`BoundExpression`] over the struct of the lambda's inputs, one field per +/// parameter: `@x` was bound as `$.x`. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct BoundLambda { + params: StructFields, + body: BoundExpression, +} + +impl BoundLambda { + /// The parameters' names and dtypes, in order. + pub fn params(&self) -> &StructFields { + &self.params + } + + /// The bound body, over the struct of the lambda's inputs. + pub fn body(&self) -> &BoundExpression { + &self.body + } + + /// The dtype of the body. + pub fn dtype(&self) -> &DType { + self.body.dtype() + } + + /// Whether the lambda is bound to these parameter dtypes. + pub fn is_bound_to(&self, param_dtypes: &[DType]) -> bool { + self.params.fields().eq(param_dtypes.iter().cloned()) + } + + /// Whether nothing in the body can fail. See [`Lambda::is_infallible`]. + pub fn is_infallible(&self) -> bool { + let mut infallible = true; + pre_order_visit_down(&self.body, |node: &BoundExpression| { + infallible &= node + .as_scalar() + .is_none_or(|scalar_fn| scalar_fn.signature().is_infallible()); + Ok(if infallible { + TraversalOrder::Continue + } else { + TraversalOrder::Stop + }) + }) + .vortex_expect("bound expression traversal cannot fail"); + infallible + } + + /// Evaluate the body over `len` rows of `args`, one array per parameter, with the dtypes the + /// lambda is bound to. They become the fields of the struct the body is applied to. + /// + /// The result is computed later, like any applied expression, so errors from running the body + /// come from executing it. + pub fn evaluate(&self, args: Vec, len: usize) -> VortexResult { + vortex_ensure!( + args.len() == self.params.nfields(), + "lambda {self} needs one array per parameter, got {}", + args.len() + ); + for ((name, dtype), arg) in self + .params + .names() + .iter() + .zip(self.params.fields()) + .zip(&args) + { + vortex_ensure!( + arg.dtype() == &dtype, + "lambda {self}: parameter {name} is bound to {dtype}, not {}", + arg.dtype() + ); + } + let inputs = StructArray::try_new( + self.params.names().clone(), + args, + len, + Validity::NonNullable, + )? + .into_array(); + inputs.apply_bound(&self.body) + } +} + +impl Display for BoundLambda { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "({}) -> {}", + self.params.names().iter().format(", "), + self.body + ) + } +} + +/// Whether `node`, in the body of a lambda with these parameters, reads a value from outside the +/// lambda, which lambdas can't do yet: `$`, another lambda's parameter, or the scan's `row_idx()` +/// or `row_count()`. The scan replaces those two before anything runs, but never looks inside a +/// lambda's body. +fn reads_outside_lambda(node: &Expression, params: &[FieldName]) -> bool { + match node { + Expression::Root => true, + // Its children are part of this body, and its own lambdas read only their parameters. + Expression::HigherOrder { .. } => false, + Expression::Variable(name) => !params.contains(name), + // `row_idx` is defined in vortex-layout, so it is recognized by its id. + Expression::Scalar { scalar_fn, .. } => { + scalar_fn.is::() || scalar_fn.id().as_ref() == "vortex.row_idx" + } + } +} diff --git a/vortex-array/src/expr/lambda/tests.rs b/vortex-array/src/expr/lambda/tests.rs new file mode 100644 index 00000000000..40f2d025855 --- /dev/null +++ b/vortex-array/src/expr/lambda/tests.rs @@ -0,0 +1,316 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::StructFields; +use crate::expr::BoundLambda; +use crate::expr::Expression; +use crate::expr::Lambda; +use crate::expr::checked_add; +use crate::expr::col; +use crate::expr::eq; +use crate::expr::get_item; +use crate::expr::lit; +use crate::expr::proto::ExprSerializeProtoExt; +use crate::expr::root; +use crate::expr::var; +use crate::higher_order_fn::HigherOrderFnVTable; +use crate::higher_order_fn::HigherOrderFnVTableExt; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::internal::row_count::RowCount; + +/// `apply(value, lambda)` gives the lambda's one parameter the value of `value`, row by row: the +/// smallest function that takes a lambda. +#[derive(Clone)] +struct Apply; + +impl HigherOrderFnVTable for Apply { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.apply"); + *ID + } + + fn arity(&self, _options: &EmptyOptions) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &EmptyOptions, _child_idx: usize) -> ChildName { + ChildName::from("value") + } + + fn lambda_count(&self, _options: &EmptyOptions) -> usize { + 1 + } + + fn bind_lambdas( + &self, + _options: &EmptyOptions, + args: &[DType], + lambdas: &[Lambda], + ) -> VortexResult> { + lambdas.iter().map(|lambda| lambda.bind(args)).collect() + } + + fn return_dtype( + &self, + _options: &EmptyOptions, + _args: &[DType], + lambdas: &[BoundLambda], + ) -> VortexResult { + Ok(lambdas[0].dtype().clone()) + } + + fn execute( + &self, + _options: &EmptyOptions, + args: &dyn ExecutionArgs, + lambdas: &[BoundLambda], + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + lambdas[0].evaluate(vec![args.get(0)?], args.row_count()) + } + + fn is_infallible(&self, _options: &EmptyOptions) -> bool { + true + } +} + +fn apply(value: Expression, param: &str, body: Expression) -> VortexResult { + Apply.new_expr(EmptyOptions, [value], [Lambda::new([param], body)?]) +} + +fn i32_() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) +} + +/// Rows `{a: i32, s: {i: i32}}`. +fn rows() -> VortexResult { + let s = StructArray::from_fields(&[("i", buffer![10_i32, 20, 30].into_array())])?; + Ok(StructArray::from_fields(&[ + ("a", buffer![1_i32, 2, 3].into_array()), + ("s", s.into_array()), + ])? + .into_array()) +} + +fn assert_applies(expr: &Expression, expected: ArrayRef) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(rows()?.apply(expr)?, expected, &mut ctx); + Ok(()) +} + +#[test] +fn new_keeps_params_and_body() -> VortexResult<()> { + let lambda = Lambda::new( + ["x", "p"], + eq(checked_add(var("x"), lit(1_i32)), get_item("k", var("p"))), + )?; + assert_eq!(lambda.params(), ["x", "p"].map(FieldName::from)); + assert_eq!( + lambda.body(), + &eq(checked_add(var("x"), lit(1_i32)), get_item("k", var("p"))) + ); + assert_eq!(lambda.to_string(), "(x, p) -> ((@x + 1i32) = @p.k)"); + Ok(()) +} + +#[rstest] +#[case::root(checked_add(var("x"), root()), "not $;")] +#[case::column(checked_add(var("x"), col("a")), "not $;")] +#[case::other_variable(checked_add(var("x"), var("p")), "not @p;")] +#[case::field_of_another_variable(get_item("k", var("p")), "not @p;")] +#[case::function_of_the_scan(RowCount.new_expr(EmptyOptions, []), "not row_count();")] +fn a_body_reads_only_its_parameters(#[case] body: Expression, #[case] read: &str) { + let error = Lambda::new(["x"], body).expect_err("reads from outside"); + let message = error.to_string(); + assert!( + message.contains("lambda (x) can only read its parameters") && message.contains(read), + "{message}" + ); +} + +#[test] +fn duplicate_parameters_are_rejected() { + let error = Lambda::new(["x", "x"], var("x")).expect_err("x is declared twice"); + assert!( + error.to_string().contains("declared more than once"), + "{error}" + ); +} + +#[test] +fn apply_runs_the_body_per_row() -> VortexResult<()> { + let expr = apply(col("a"), "x", checked_add(var("x"), lit(1_i32)))?; + assert_eq!( + expr.to_string(), + "vortex.test.apply($.a, (x) -> (@x + 1i32))" + ); + assert_applies(&expr, buffer![2_i32, 3, 4].into_array()) +} + +#[test] +fn a_parameter_named_like_a_column() -> VortexResult<()> { + // `var("a")` is the parameter, whatever the table's column `a` holds. + let expr = apply( + get_item("i", col("s")), + "a", + checked_add(var("a"), lit(1_i32)), + )?; + assert_applies(&expr, buffer![11_i32, 21, 31].into_array()) +} + +#[test] +fn nested_lambdas_read_their_own_parameters() -> VortexResult<()> { + // p -> apply(p.i, y -> y + 1): the inner value reads the outer parameter. + let inner = apply( + get_item("i", var("p")), + "y", + checked_add(var("y"), lit(1_i32)), + )?; + let expr = apply(col("s"), "p", inner)?; + assert_eq!( + expr.to_string(), + "vortex.test.apply($.s, (p) -> vortex.test.apply(@p.i, (y) -> (@y + 1i32)))" + ); + assert_applies(&expr, buffer![11_i32, 21, 31].into_array()) +} + +#[test] +fn an_inner_parameter_shadows_an_outer_one() -> VortexResult<()> { + // x -> apply(x + 1, x -> x + x) + let inner = apply( + checked_add(var("x"), lit(1_i32)), + "x", + checked_add(var("x"), var("x")), + )?; + let expr = apply(col("a"), "x", inner)?; + assert_applies(&expr, buffer![4_i32, 6, 8].into_array()) +} + +#[rstest] +#[case::bind(|expr: &Expression, rows: ArrayRef| expr.bind(rows.dtype()).map(drop))] +#[case::return_dtype(|expr: &Expression, rows: ArrayRef| expr.return_dtype(rows.dtype()).map(drop))] +#[case::apply(|expr: &Expression, rows: ArrayRef| rows.apply(expr).map(drop))] +#[case::serialize(|expr: &Expression, _rows: ArrayRef| expr.serialize_proto().map(drop))] +fn a_variable_outside_a_lambda_is_an_error( + #[case] use_expr: fn(&Expression, ArrayRef) -> VortexResult<()>, +) -> VortexResult<()> { + let error = + use_expr(&checked_add(var("y"), lit(1_i32)), rows()?).expect_err("y is not a parameter"); + assert!( + error + .to_string() + .contains("variable @y is not a parameter of an enclosing lambda"), + "{error}" + ); + Ok(()) +} + +#[test] +fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { + let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?; + assert_eq!(lambda.return_dtype(&[i32_()])?, i32_()); + let bound = lambda.bind(&[i32_()])?; + assert_eq!(bound.dtype(), &i32_()); + assert!(bound.is_bound_to(&[i32_()])); + // The body reads the struct of the lambda's inputs: `@x` became `$.x`. + assert_eq!(bound.to_string(), "(x) -> ($.x + 1i32)"); + let inputs = DType::Struct( + StructFields::new(["x"].into(), vec![i32_()]), + Nullability::NonNullable, + ); + assert_eq!( + bound.body(), + &checked_add(col("x"), lit(1_i32)).bind(&inputs)? + ); + assert_eq!( + lambda.bind(&[i32_().as_nullable()])?.dtype(), + &i32_().as_nullable() + ); + for error in [ + lambda.return_dtype(&[]).expect_err("x has no dtype"), + lambda.bind(&[]).expect_err("x has no dtype"), + ] { + assert!( + error + .to_string() + .contains("lambda (x) -> (@x + 1i32) needs one dtype per parameter, got 0"), + "{error}" + ); + } + Ok(()) +} + +#[test] +fn binding_an_expression_binds_its_lambdas() -> VortexResult<()> { + let expr = apply(col("a"), "x", checked_add(var("x"), lit(1_i32)))?; + let (_, lambdas) = expr.as_higher_order().expect("a higher-order node"); + assert_eq!(lambdas[0].to_string(), "(x) -> (@x + 1i32)"); + let bound = expr.bind(rows()?.dtype())?; + let (_, bound_lambdas) = bound.as_higher_order().expect("a higher-order node"); + assert!(bound_lambdas[0].is_bound_to(&[i32_()])); + // The bound node shows its lambdas as written. + assert_eq!(bound.to_string(), expr.to_string()); + Ok(()) +} + +#[test] +fn evaluate_runs_the_bound_body() -> VortexResult<()> { + let lambda = + Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?.bind(&[i32_().as_nullable()])?; + let x = PrimitiveArray::from_option_iter([Some(1_i32), None]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + lambda.evaluate(vec![x], 2)?, + PrimitiveArray::from_option_iter([Some(2_i32), None]).into_array(), + &mut ctx + ); + + // The parameters must have the dtypes the lambda is bound to. + let error = lambda + .evaluate(vec![buffer![1_i32, 2].into_array()], 2) + .expect_err("bound to i32?, given i32"); + assert!( + error + .to_string() + .contains("parameter x is bound to i32?, not i32"), + "{error}" + ); + Ok(()) +} + +#[rstest] +#[case::comparison(eq(var("x"), lit(3_i32)), true)] +#[case::checked_arithmetic(checked_add(var("x"), lit(1_i32)), false)] +fn infallibility_comes_from_the_body( + #[case] body: Expression, + #[case] infallible: bool, +) -> VortexResult<()> { + let lambda = Lambda::new(["x"], body)?; + assert_eq!(lambda.is_infallible(), infallible); + assert_eq!(lambda.bind(&[i32_()])?.is_infallible(), infallible); + Ok(()) +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 8dbb3c92bcb..d848c166344 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -49,6 +49,7 @@ pub(crate) mod expression; mod exprs; pub(crate) mod field; pub mod forms; +pub mod lambda; mod optimize; pub mod proto; pub mod scope; @@ -91,6 +92,7 @@ pub use exprs::list_contains_opts; pub use exprs::list_length; pub use exprs::list_sum; pub use exprs::list_sum_opts; +pub use exprs::list_transform; pub use exprs::lit; pub use exprs::lt; pub use exprs::lt_eq; @@ -109,8 +111,10 @@ pub use exprs::replace_time_zone; pub use exprs::root; pub use exprs::select; pub use exprs::select_exclude; +pub use exprs::var; pub use exprs::variant_get; pub use exprs::zip_expr; +pub use lambda::*; pub use scope::*; #[cfg(feature = "_test-harness")] diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index ecf7e3d21e7..e3907936ea7 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -22,7 +22,8 @@ impl BoundExpression { fn simplify_node(&self) -> VortexResult> { match self { BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self), - BoundExpression::Root { .. } => Ok(None), + // Higher-order functions have no simplification rules. + BoundExpression::Root { .. } | BoundExpression::HigherOrder { .. } => Ok(None), } } @@ -33,7 +34,8 @@ impl BoundExpression { ) -> VortexResult>> { match self { BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node), - BoundExpression::Root { .. } => Ok(None), + // Higher-order functions have no reduce rules. + BoundExpression::Root { .. } | BoundExpression::HigherOrder { .. } => Ok(None), } } diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index c09da2c2a2f..800cfe935a2 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -3,11 +3,13 @@ use itertools::Itertools; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; use crate::expr::Expression; +use crate::expr::expression::unbound_variable; use crate::proto::expr as pb; use crate::scalar_fn::ForeignScalarFnVTable; use crate::scalar_fn::ScalarFnId; @@ -24,6 +26,15 @@ pub(crate) const ROOT_ID: &str = "vortex.root"; impl ExprSerializeProtoExt for Expression { fn serialize_proto(&self) -> VortexResult { + if let Expression::Variable(name) = self { + return Err(unbound_variable(name)); + } + if let Some((function, _)) = self.as_higher_order() { + vortex_bail!( + "higher-order function {} is not serializable yet", + function.id() + ); + } let Some(scalar_fn) = self.as_scalar() else { return Ok(pb::Expr { id: ROOT_ID.to_string(), diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index 5b80816c30a..ea3e423f64b 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -475,7 +475,9 @@ impl Node for BoundExpression { &'a self, mut f: F, ) -> VortexResult { - let BoundExpression::Scalar { children, .. } = self else { + let (BoundExpression::Scalar { children, .. } + | BoundExpression::HigherOrder { children, .. }) = self + else { return Ok(TraversalOrder::Continue); }; @@ -493,7 +495,9 @@ impl Node for BoundExpression { self, mut f: F, ) -> VortexResult> { - let BoundExpression::Scalar { children, .. } = &self else { + let (BoundExpression::Scalar { children, .. } + | BoundExpression::HigherOrder { children, .. }) = &self + else { return Ok(Transformed::no(self)); }; @@ -534,14 +538,16 @@ 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::HigherOrder { children, .. } => f(&mut children.iter()), BoundExpression::Root { .. } => f(&mut std::iter::empty()), } } fn children_count(&self) -> usize { match self { - BoundExpression::Scalar { children, .. } => children.len(), + BoundExpression::Scalar { children, .. } + | BoundExpression::HigherOrder { children, .. } => children.len(), BoundExpression::Root { .. } => 0, } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index 3d2684f2dc7..8d9d15d2515 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use smallvec::SmallVec; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -8,15 +10,37 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ConstantArray; +use crate::arrays::HigherOrderArray; use crate::arrays::ScalarFnArray; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::expression::unbound_variable; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::literal::Literal; impl ArrayRef { /// Apply a bound expression to this array, producing a new array in constant time. pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult { + if let BoundExpression::HigherOrder { + function, + children, + bound, + .. + } = expr + { + let len = self.len(); + let children = children + .iter() + .map(|child| self.clone().apply_bound(child)) + .collect::>>()?; + return Ok(HigherOrderArray::try_new( + function.clone(), + Arc::clone(bound), + children, + len, + )? + .into_array()); + } let BoundExpression::Scalar { scalar_fn, children, @@ -47,6 +71,9 @@ impl ArrayRef { if expr.is_root() { return Ok(self); } + if let Expression::Variable(name) = expr { + return Err(unbound_variable(name)); + } // Manually convert literals to ConstantArray. if let Some(scalar) = expr.as_opt::() { @@ -61,10 +88,28 @@ impl ArrayRef { .map(|e| self.clone().apply(e)) .collect::>>()?; + if let Expression::HigherOrder { + function, lambdas, .. + } = expr + { + let arg_dtypes = children + .iter() + .map(|child| child.dtype().clone()) + .collect::>(); + let bound = function.bind_lambdas(&arg_dtypes, lambdas)?; + return Ok(HigherOrderArray::try_new( + function.clone(), + Arc::new(bound), + children, + len, + )? + .into_array()); + } + // And wrap the scalar function up in an array. let scalar_fn = expr .as_scalar() - .vortex_expect("root and literal were handled above, so this is a scalar node"); + .vortex_expect("leaves and literals were handled above, so this is a scalar node"); let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array(); // Optimize the resulting array's root. diff --git a/vortex-array/src/higher_order_fn/fns/list_transform/mod.rs b/vortex-array/src/higher_order_fn/fns/list_transform/mod.rs new file mode 100644 index 00000000000..2cacd02879c --- /dev/null +++ b/vortex-array/src/higher_order_fn/fns/list_transform/mod.rs @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(test)] +mod tests; + +use std::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::listview::list_from_list_view; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::expr::BoundLambda; +use crate::expr::Lambda; +use crate::higher_order_fn::HigherOrderFnVTable; +use crate::match_each_integer_ptype; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + +/// Apply a lambda to every element of each list, keeping each list's length and validity. +/// +/// The lambda takes one parameter, the element. The only child is the list, of `List` dtype in any +/// encoding, and the result is a `List`. The lambda runs only on the elements of valid lists. +#[derive(Clone)] +pub struct ListTransform; + +impl HigherOrderFnVTable for ListTransform { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.transform"); + *ID + } + + fn arity(&self, _options: &EmptyOptions) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &EmptyOptions, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("list"), + _ => unreachable!("Invalid child index {child_idx} for list_transform()"), + } + } + + fn lambda_count(&self, _options: &EmptyOptions) -> usize { + 1 + } + + fn bind_lambdas( + &self, + _options: &EmptyOptions, + args: &[DType], + lambdas: &[Lambda], + ) -> VortexResult> { + let [lambda] = lambdas else { + vortex_bail!("list_transform() takes one lambda, got {}", lambdas.len()); + }; + Ok(vec![ + lambda.bind(&element_dtype(&args[0], lambda.params())?)?, + ]) + } + + fn return_dtype( + &self, + _options: &EmptyOptions, + args: &[DType], + lambdas: &[BoundLambda], + ) -> VortexResult { + let lambda = one_lambda(lambdas)?; + check_bound_to(lambda, &args[0])?; + Ok(DType::List( + Arc::new(lambda.dtype().clone()), + args[0].nullability(), + )) + } + + fn execute( + &self, + _options: &EmptyOptions, + args: &dyn ExecutionArgs, + lambdas: &[BoundLambda], + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let lambda = one_lambda(lambdas)?; + // Every element belongs to exactly one valid list, in order: nothing comes before the + // first list or after the last, and a null list has no elements. + let list = args.get(0)?.execute::(ctx)?; + let list = list_from_list_view(empty_null_lists(list, ctx)?, ctx)?; + let elements = list.elements(); + let values = lambda.evaluate(vec![elements.clone()], elements.len())?; + vortex_ensure_eq!(values.len(), elements.len()); + // SAFETY: the offsets are the list's, over one value per element. + Ok(unsafe { + ListArray::new_unchecked(values, list.offsets().clone(), list.list_validity()) + } + .into_array()) + } + + fn is_strict(&self, _options: &EmptyOptions) -> bool { + true + } + + fn is_infallible(&self, _options: &EmptyOptions) -> bool { + true + } +} + +fn one_lambda(lambdas: &[BoundLambda]) -> VortexResult<&BoundLambda> { + let [lambda] = lambdas else { + vortex_bail!("list_transform() takes one lambda, got {}", lambdas.len()); + }; + Ok(lambda) +} + +/// Check that a bound lambda's parameter has the element dtype of a list of this dtype. +fn check_bound_to(lambda: &BoundLambda, list: &DType) -> VortexResult<()> { + let element = element_dtype(list, lambda.params().names().as_ref())?; + vortex_ensure!( + lambda.is_bound_to(&element), + "list_transform()'s lambda {lambda} is bound to element dtype {}, not {}", + lambda.params().fields().format(", "), + element[0] + ); + Ok(()) +} + +/// The dtype of the lambda's one parameter, the element, for a list of this dtype. +fn element_dtype(list: &DType, params: &[FieldName]) -> VortexResult<[DType; 1]> { + let DType::List(element, _) = list else { + vortex_bail!("list_transform() takes a List, got {list}"); + }; + vortex_ensure!( + params.len() == 1, + "list_transform() takes a lambda of one parameter, the element, not ({})", + params.iter().format(", ") + ); + Ok([element.as_ref().clone()]) +} + +/// The lists with every null list emptied, so that nothing under a null list reaches the lambda. +fn empty_null_lists(list: ListViewArray, ctx: &mut ExecutionCtx) -> VortexResult { + let validity = list.listview_validity(); + if matches!(validity, Validity::NonNullable | Validity::AllValid) { + return Ok(list); + } + let valid = validity.execute_mask(list.len(), ctx)?; + let sizes = list.sizes().clone().execute::(ctx)?; + let sizes = match_each_integer_ptype!(sizes.ptype(), |T| { + let sizes = sizes.as_slice::(); + let null = |index: usize| !valid.value(index); + // Writers usually leave null lists empty already. + if !(0..sizes.len()).any(|index| null(index) && sizes[index] != T::default()) { + return Ok(list); + } + PrimitiveArray::from_iter((0..sizes.len()).map(|index| { + if null(index) { + T::default() + } else { + sizes[index] + } + })) + }); + // SAFETY: only sizes changed, and only to 0, so every view stays within the elements. The new + // array is not marked zero-copy to a list: an emptied list leaves its elements behind as a gap, + // which the rebuild must drop. + Ok(unsafe { + ListViewArray::new_unchecked( + list.elements().clone(), + list.offsets().clone(), + sizes.into_array(), + validity, + ) + }) +} diff --git a/vortex-array/src/higher_order_fn/fns/list_transform/tests.rs b/vortex-array/src/higher_order_fn/fns/list_transform/tests.rs new file mode 100644 index 00000000000..c11a12e56d1 --- /dev/null +++ b/vortex-array/src/higher_order_fn/fns/list_transform/tests.rs @@ -0,0 +1,546 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::RecursiveCanonical; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::Dict; +use crate::arrays::DictArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::BoundExpression; +use crate::expr::Expression; +use crate::expr::Lambda; +use crate::expr::checked_add; +use crate::expr::col; +use crate::expr::eq; +use crate::expr::get_item; +use crate::expr::is_not_null; +use crate::expr::list_contains; +use crate::expr::list_transform; +use crate::expr::lit; +use crate::expr::proto::ExprSerializeProtoExt; +use crate::expr::root; +use crate::expr::transform::replace; +use crate::expr::var; +use crate::scalar::Scalar; +use crate::validity::Validity; + +/// `[[1, 2], null (hiding [i32::MAX]), [3, null]]` +fn lists() -> VortexResult { + list_of( + PrimitiveArray::from_option_iter([Some(1_i32), Some(2), Some(i32::MAX), Some(3), None]) + .into_array(), + &[0, 2, 3, 5], + Validity::from_iter([true, false, true]), + ) +} + +/// Rows `{l}`. +fn rows(l: ArrayRef) -> VortexResult { + let len = l.len(); + Ok( + StructArray::try_new(FieldNames::from(["l"]), vec![l], len, Validity::NonNullable)? + .into_array(), + ) +} + +fn list_of(values: ArrayRef, offsets: &[u32], validity: Validity) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + Ok(ListArray::try_new( + values, + PrimitiveArray::from_iter(offsets.iter().copied()).into_array(), + validity, + &mut ctx, + )? + .into_array()) +} + +/// The lists of [`lists`] with new elements: the null list keeps no elements. +fn transformed(values: ArrayRef) -> VortexResult { + list_of( + values, + &[0, 2, 2, 4], + Validity::from_iter([true, false, true]), + ) +} + +/// `x -> x + 1`, which fails on `i32::MAX`. +fn plus_one() -> VortexResult { + Lambda::new(["x"], checked_add(var("x"), lit(1_i32))) +} + +fn assert_applies(input: ArrayRef, expr: &Expression, expected: ArrayRef) -> VortexResult<()> { + let actual = input.apply(expr)?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn the_lambda_runs_on_every_element() -> VortexResult<()> { + let expr = list_transform(col("l"), plus_one()?); + assert_eq!( + expr.to_string(), + "vortex.list.transform($.l, (x) -> (@x + 1i32))" + ); + // The null list's `i32::MAX` would overflow, so it must not reach the lambda. + let expected = transformed( + PrimitiveArray::from_option_iter([Some(2_i32), Some(3), Some(4), None]).into_array(), + )?; + assert_applies(rows(lists()?)?, &expr, expected) +} + +#[test] +fn a_body_of_another_dtype() -> VortexResult<()> { + let expr = list_transform(root(), Lambda::new(["x"], eq(var("x"), lit(3_i32)))?); + let expected = transformed( + BoolArray::from_iter([Some(false), Some(false), Some(true), None]).into_array(), + )?; + assert_applies(lists()?, &expr, expected) +} + +#[test] +fn a_null_list_without_elements() -> VortexResult<()> { + // `[[1, 2], null, [3]]` with nothing under the null list. + let validity = || Validity::from_iter([true, false, true]); + let lists = list_of(buffer![1_i32, 2, 3].into_array(), &[0, 2, 2, 3], validity())?; + let expected = list_of(buffer![2_i32, 3, 4].into_array(), &[0, 2, 2, 3], validity())?; + assert_applies(lists, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn a_sliced_list_skips_a_null_lists_elements() -> VortexResult<()> { + // `[[1, 2], null (hiding [i32::MAX]), [3], [4]]` sliced to `[null, [3], [4]]`, whose offsets + // no longer start at 0. + let lists = list_of( + buffer![1_i32, 2, i32::MAX, 3, 4].into_array(), + &[0, 2, 3, 4, 5], + Validity::from_iter([true, false, true, true]), + )? + .slice(1..4)?; + let expected = list_of( + buffer![4_i32, 5].into_array(), + &[0, 0, 1, 2], + Validity::from_iter([false, true, true]), + )?; + assert_applies(lists, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn every_list_null() -> VortexResult<()> { + // `[null (hiding [i32::MAX]), null (hiding [i32::MAX])]` + let lists = list_of( + buffer![i32::MAX, i32::MAX].into_array(), + &[0, 1, 2], + Validity::from_iter([false, false]), + )?; + let expected = list_of( + PrimitiveArray::from_iter::<[i32; 0]>([]).into_array(), + &[0, 0, 0], + Validity::from_iter([false, false]), + )?; + assert_applies(lists, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn a_failure_in_a_valid_list_is_an_error() -> VortexResult<()> { + let lists = list_of( + buffer![1_i32, i32::MAX].into_array(), + &[0, 1, 2], + Validity::from_iter([false, true]), + )?; + let mut ctx = array_session().create_execution_ctx(); + let error = lists + .apply(&list_transform(root(), plus_one()?))? + .execute::(&mut ctx) + .map(drop) + .expect_err("i32::MAX + 1 overflows in a valid list"); + assert!(error.to_string().contains("overflow"), "{error}"); + Ok(()) +} + +#[test] +fn a_lambda_takes_one_parameter() -> VortexResult<()> { + let expr = list_transform(root(), Lambda::new(["x", "i"], var("x"))?); + let error = expr.bind(lists()?.dtype()).expect_err("two parameters"); + assert!( + error + .to_string() + .contains("takes a lambda of one parameter"), + "{error}" + ); + Ok(()) +} + +#[test] +fn struct_elements_read_their_fields() -> VortexResult<()> { + // `[[{x: 1}, {x: 2}], [{x: 3}]]` under `p -> p.x + 1` + let points = StructArray::try_new( + FieldNames::from(["x"]), + vec![buffer![1_i32, 2, 3].into_array()], + 3, + Validity::NonNullable, + )? + .into_array(); + let expr = list_transform( + root(), + Lambda::new(["p"], checked_add(get_item("x", var("p")), lit(1_i32)))?, + ); + let expected = list_of( + buffer![2_i32, 3, 4].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + assert_applies( + list_of(points, &[0, 2, 3], Validity::NonNullable)?, + &expr, + expected, + ) +} + +#[test] +fn a_body_can_test_membership_in_a_literal_set() -> VortexResult<()> { + // `[[1, 2], [3]]` under `x -> list_contains([2, 3], x)` + let set = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![2_i32.into(), 3_i32.into()], + Nullability::NonNullable, + ); + let expr = list_transform( + root(), + Lambda::new(["x"], list_contains(lit(set), var("x")))?, + ); + let lists = list_of( + buffer![1_i32, 2, 3].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + let expected = list_of( + BoolArray::from_iter([false, true, true]).into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + assert_applies(lists, &expr, expected) +} + +#[test] +fn a_nested_list_result_over_null_lists() -> VortexResult<()> { + // `[[[1], [2]], null, [[3]]]` under `xs -> list_transform(xs, y -> y + 1)`: the body returns + // a list, and a null list must not reach the inner body. + let inner = list_of( + buffer![1_i32, 2, i32::MAX, 3].into_array(), + &[0, 1, 2, 3, 4], + Validity::NonNullable, + )?; + let outer = list_of( + inner, + &[0, 2, 3, 4], + Validity::from_iter([true, false, true]), + )?; + let expr = list_transform( + root(), + Lambda::new(["xs"], list_transform(var("xs"), plus_one()?))?, + ); + assert_eq!( + expr.to_string(), + "vortex.list.transform($, (xs) -> vortex.list.transform(@xs, (x) -> (@x + 1i32)))" + ); + let inner_expected = list_of( + buffer![2_i32, 3, 4].into_array(), + &[0, 1, 2, 3], + Validity::NonNullable, + )?; + let expected = list_of( + inner_expected, + &[0, 2, 2, 3], + Validity::from_iter([true, false, true]), + )?; + assert_applies(outer, &expr, expected) +} + +#[test] +fn three_nested_lambdas() -> VortexResult<()> { + // `[[[[1, 2], [3]]], [[[4]]]]` under + // `x -> list_transform(x, y -> list_transform(y, z -> z + 1))` + let l = list_of( + list_of( + list_of( + buffer![1_i32, 2, 3, 4].into_array(), + &[0, 2, 3, 4], + Validity::NonNullable, + )?, + &[0, 2, 3], + Validity::NonNullable, + )?, + &[0, 1, 2], + Validity::NonNullable, + )?; + let innermost = Lambda::new(["z"], checked_add(var("z"), lit(1_i32)))?; + let middle = Lambda::new(["y"], list_transform(var("y"), innermost))?; + let expr = list_transform( + root(), + Lambda::new(["x"], list_transform(var("x"), middle))?, + ); + let expected = list_of( + list_of( + list_of( + buffer![2_i32, 3, 4, 5].into_array(), + &[0, 2, 3, 4], + Validity::NonNullable, + )?, + &[0, 2, 3], + Validity::NonNullable, + )?, + &[0, 1, 2], + Validity::NonNullable, + )?; + assert_applies(l, &expr, expected) +} + +#[test] +fn sliced_list_views() -> VortexResult<()> { + // `[[1, 2], [3], [4, 5]]` as a list view, sliced to `[[3], [4, 5]]`, which leaves unused + // elements. + let mut ctx = array_session().create_execution_ctx(); + let view = list_of( + buffer![1_i32, 2, 3, 4, 5].into_array(), + &[0, 2, 3, 5], + Validity::NonNullable, + )? + .execute::(&mut ctx)? + .into_array() + .slice(1..3)?; + let expected = list_of( + buffer![4_i32, 5, 6].into_array(), + &[0, 1, 3], + Validity::NonNullable, + )?; + assert_applies(view, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn overlapping_list_views() -> VortexResult<()> { + // `[[3, 4], [1], [2, 3]]`, sharing element 3. + let mut ctx = array_session().create_execution_ctx(); + let view = ListViewArray::try_new( + buffer![1_i32, 2, 3, 4].into_array(), + buffer![2_u32, 0, 1].into_array(), + buffer![2_u32, 1, 2].into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array(); + let expr = list_transform(root(), Lambda::new(["x"], eq(var("x"), lit(3_i32)))?); + let expected = list_of( + BoolArray::from_iter([true, false, false, false, true]).into_array(), + &[0, 2, 3, 5], + Validity::NonNullable, + )?; + assert_applies(view, &expr, expected) +} + +#[test] +fn fixed_size_lists_are_not_supported_yet() -> VortexResult<()> { + let fsl = FixedSizeListArray::try_new( + buffer![1_i32, 2, 3, 4].into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let error = list_transform(root(), plus_one()?) + .bind(fsl.dtype()) + .expect_err("a fixed-size list is not a list"); + assert!( + error.to_string().contains("list_transform() takes a List"), + "{error}" + ); + Ok(()) +} + +#[test] +fn a_constant_body() -> VortexResult<()> { + let expr = list_transform(root(), Lambda::new(["x"], lit(7_i32))?); + assert_applies( + lists()?, + &expr, + transformed(buffer![7_i32, 7, 7, 7].into_array())?, + ) +} + +#[test] +fn no_lists() -> VortexResult<()> { + let empty = list_of( + PrimitiveArray::from_iter::<[i32; 0]>([]).into_array(), + &[0], + Validity::NonNullable, + )?; + let expected = list_of( + PrimitiveArray::from_iter::<[i32; 0]>([]).into_array(), + &[0], + Validity::NonNullable, + )?; + assert_applies(empty, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn the_bound_node_holds_its_lambda() -> VortexResult<()> { + let dtype = rows(lists()?)?.dtype().clone(); + let expr = list_transform(col("l"), plus_one()?); + let bound = expr.bind(&dtype)?; + // The bound node shows its lambda as written. + assert_eq!( + bound.to_string(), + "vortex.list.transform($.l, (x) -> (@x + 1i32))" + ); + assert_eq!(bound.children(), [col("l").bind(&dtype)?]); + let nullable_i32 = DType::Primitive(PType::I32, Nullability::Nullable); + assert_eq!( + bound.dtype(), + &DType::List(Arc::new(nullable_i32.clone()), Nullability::Nullable) + ); + assert_eq!(&expr.return_dtype(&dtype)?, bound.dtype()); + assert_eq!(bound.validity()?, is_not_null(col("l")).bind(&dtype)?); + // Binding the node bound its lambda to the element's dtype. + let (_, lambdas) = bound.as_higher_order().expect("a higher-order node"); + assert!(lambdas[0].is_bound_to(&[nullable_i32])); + Ok(()) +} + +#[test] +fn applying_the_bound_node() -> VortexResult<()> { + // A scan binds the expression once and applies the bound form to every batch. + let input = rows(lists()?)?; + let bound = list_transform(col("l"), plus_one()?).bind(input.dtype())?; + let expected = transformed( + PrimitiveArray::from_option_iter([Some(2_i32), Some(3), Some(4), None]).into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(input.apply_bound(&bound)?, expected, &mut ctx); + Ok(()) +} + +#[test] +fn a_rebuilt_node_rebinds_its_lambda() -> VortexResult<()> { + let list = |elements| { + DType::List( + Arc::new(DType::Primitive(PType::I32, elements)), + Nullability::NonNullable, + ) + }; + let bound = list_transform(root(), plus_one()?).bind(&list(Nullability::NonNullable))?; + // The node keeps its written lambda, so rebuilding it over nullable elements rebinds it. + let rebuilt = bound + .clone() + .with_children([BoundExpression::new_root(list(Nullability::Nullable))])?; + let nullable_i32 = DType::Primitive(PType::I32, Nullability::Nullable); + let (_, lambdas) = rebuilt.as_higher_order().expect("a higher-order node"); + assert!(lambdas[0].is_bound_to(std::slice::from_ref(&nullable_i32))); + assert_eq!( + rebuilt.dtype(), + &DType::List(Arc::new(nullable_i32), Nullability::NonNullable) + ); + + // Asked for the dtype with lambdas bound to other dtypes, the function says so. + let (function, lambdas) = bound.as_higher_order().expect("a higher-order node"); + let error = function + .return_dtype(&[list(Nullability::Nullable)], lambdas) + .expect_err("bound to i32, given i32?"); + assert!( + error + .to_string() + .contains("is bound to element dtype i32, not i32?"), + "{error}" + ); + Ok(()) +} + +#[test] +fn replacing_the_root_rereads_the_list() -> VortexResult<()> { + // Rows `{s: {l}}`: re-rooting the transform at `$.s` changes where the list is read. The body + // is in the options, which rewrites of the expression leave alone. + let input = StructArray::try_new( + FieldNames::from(["s"]), + vec![rows(lists()?)?], + 3, + Validity::NonNullable, + )? + .into_array(); + let expr = list_transform(col("l"), plus_one()?); + assert_eq!(replace(expr.clone(), &var("x"), lit(0_i32)), expr); + let expr = replace(expr, &root(), col("s")); + assert_eq!( + expr.to_string(), + "vortex.list.transform($.s.l, (x) -> (@x + 1i32))" + ); + let expected = transformed( + PrimitiveArray::from_option_iter([Some(2_i32), Some(3), Some(4), None]).into_array(), + )?; + assert_applies(input, &expr, expected) +} + +#[test] +#[ignore = "prototype: dictionary pushdown only knows ScalarFnArray, not HigherOrderArray"] +fn an_infallible_transform_pushes_into_dictionary_values() -> VortexResult<()> { + let values = list_of( + buffer![1_i32, 2, 3].into_array(), + &[0, 2, 3], + Validity::NonNullable, + )?; + let dict = DictArray::try_new(buffer![0_u8, 1, 0, 1].into_array(), values)?.into_array(); + let expr = list_transform(root(), Lambda::new(["x"], eq(var("x"), lit(1_i32)))?); + + let actual = dict.apply(&expr)?; + assert!(actual.is::(), "{actual}"); + let expected = list_of( + BoolArray::from_iter([true, false, false, true, false, false]).into_array(), + &[0, 2, 3, 5, 6], + Validity::NonNullable, + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn a_fallible_transform_skips_unused_dictionary_values() -> VortexResult<()> { + // Dictionary values `[[1], [i32::MAX]]`, of which the codes use only the first. + let values = list_of( + buffer![1_i32, i32::MAX].into_array(), + &[0, 1, 2], + Validity::NonNullable, + )?; + let dict = DictArray::try_new(buffer![0_u8, 0].into_array(), values)?.into_array(); + let expected = list_of( + buffer![2_i32, 2].into_array(), + &[0, 1, 2], + Validity::NonNullable, + )?; + assert_applies(dict, &list_transform(root(), plus_one()?), expected) +} + +#[test] +fn lambdas_are_not_serializable_yet() -> VortexResult<()> { + let expr = list_transform(col("l"), Lambda::new(["x"], var("x"))?); + let error = expr + .serialize_proto() + .expect_err("lambdas are not serializable"); + assert!(error.to_string().contains("is not serializable"), "{error}"); + Ok(()) +} diff --git a/vortex-array/src/higher_order_fn/fns/mod.rs b/vortex-array/src/higher_order_fn/fns/mod.rs new file mode 100644 index 00000000000..e9bf58dba45 --- /dev/null +++ b/vortex-array/src/higher_order_fn/fns/mod.rs @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Built-in higher-order functions. + +pub mod list_transform; diff --git a/vortex-array/src/higher_order_fn/mod.rs b/vortex-array/src/higher_order_fn/mod.rs new file mode 100644 index 00000000000..0817bc1e65b --- /dev/null +++ b/vortex-array/src/higher_order_fn/mod.rs @@ -0,0 +1,401 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Higher-order functions: functions that take lambdas as well as values, such as +//! `list_transform`. +//! +//! A higher-order function is its own kind of node, parallel to a scalar function: an +//! [`Expression::HigherOrder`] holds its value children and its lambdas as written, a +//! [`BoundExpression::HigherOrder`](crate::expr::BoundExpression::HigherOrder) holds them bound, and a +//! [`HigherOrderArray`](crate::arrays::HigherOrderArray) runs it. The function decides how each +//! lambda's parameters are typed ([`HigherOrderFnVTable::bind_lambdas`]) and when the lambdas run +//! ([`HigherOrderFnVTable::execute`]). + +pub mod fns; + +use std::any::Any; +use std::fmt; +use std::fmt::Debug; +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use vortex_error::VortexResult; +use vortex_utils::debug_with::DebugWith; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::expr::BoundLambda; +use crate::expr::Expression; +use crate::expr::Lambda; +use crate::expr::display::ExprDisplay; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; + +/// A function that takes lambdas as well as values. +pub trait HigherOrderFnVTable: 'static + Sized + Clone + Send + Sync { + /// Options other than the lambdas. + type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; + + /// The function's id. + fn id(&self) -> ScalarFnId; + + /// How many value children the function takes. + fn arity(&self, options: &Self::Options) -> Arity; + + /// The name of a value child. + fn child_name(&self, options: &Self::Options, child_idx: usize) -> ChildName; + + /// How many lambdas the function takes. + fn lambda_count(&self, options: &Self::Options) -> usize; + + /// Bind the lambdas against the value children's dtypes. + /// + /// The function decides each lambda's parameter dtypes, in whatever order it needs: a fold + /// binds the lambda that merges an element into the accumulator before the one that finishes. + fn bind_lambdas( + &self, + options: &Self::Options, + args: &[DType], + lambdas: &[Lambda], + ) -> VortexResult>; + + /// The result dtype, given the value children's dtypes and the bound lambdas. + fn return_dtype( + &self, + options: &Self::Options, + args: &[DType], + lambdas: &[BoundLambda], + ) -> VortexResult; + + /// Run the function over its value children, with its lambdas bound. + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + lambdas: &[BoundLambda], + ctx: &mut ExecutionCtx, + ) -> VortexResult; + + /// Whether a null in any value child makes the result null. + fn is_strict(&self, options: &Self::Options) -> bool { + _ = options; + false + } + + /// Whether the function itself can never fail. Its lambdas are checked separately. + fn is_infallible(&self, options: &Self::Options) -> bool { + _ = options; + false + } + + /// Format the call as SQL, given a way to display the value children. + fn fmt_sql( + &self, + options: &Self::Options, + expr: &dyn ExprDisplay, + lambdas: &[Lambda], + f: &mut Formatter<'_>, + ) -> fmt::Result { + write!(f, "{}(", self.id())?; + for idx in 0..expr.display_children_count() { + write!(f, "{}, ", expr.display_child(idx))?; + } + for (idx, lambda) in lambdas.iter().enumerate() { + if idx > 0 { + write!(f, ", ")?; + } + write!(f, "{lambda}")?; + } + _ = options; + write!(f, ")") + } +} + +/// Construct higher-order function nodes. +pub trait HigherOrderFnVTableExt: HigherOrderFnVTable { + /// This function with these options, type-erased. + fn bind(&self, options: Self::Options) -> HigherOrderFnRef { + HigherOrderFnRef(Arc::new(Instance { + vtable: self.clone(), + options, + })) + } + + /// A written node applying this function to these value children and lambdas. + fn new_expr( + &self, + options: Self::Options, + children: impl IntoIterator, + lambdas: impl IntoIterator, + ) -> VortexResult { + Expression::try_new_higher_order(self.bind(options), children, lambdas) + } +} + +impl HigherOrderFnVTableExt for V {} + +struct Instance { + vtable: V, + options: V::Options, +} + +trait DynHigherOrderFn: 'static + Send + Sync { + fn as_any(&self) -> &dyn Any; + fn id(&self) -> ScalarFnId; + fn arity(&self) -> Arity; + fn child_name(&self, child_idx: usize) -> ChildName; + fn lambda_count(&self) -> usize; + fn bind_lambdas(&self, args: &[DType], lambdas: &[Lambda]) -> VortexResult>; + fn return_dtype(&self, args: &[DType], lambdas: &[BoundLambda]) -> VortexResult; + fn execute( + &self, + args: &dyn ExecutionArgs, + lambdas: &[BoundLambda], + ctx: &mut ExecutionCtx, + ) -> VortexResult; + fn is_strict(&self) -> bool; + fn is_infallible(&self) -> bool; + fn fmt_sql( + &self, + expr: &dyn ExprDisplay, + lambdas: &[Lambda], + f: &mut Formatter<'_>, + ) -> fmt::Result; + fn options_any(&self) -> &dyn Any; + fn options_eq(&self, other: &dyn Any) -> bool; + fn options_hash(&self, hasher: &mut dyn Hasher); + fn options_display(&self, f: &mut Formatter<'_>) -> fmt::Result; + fn options_debug(&self, f: &mut Formatter<'_>) -> fmt::Result; +} + +impl DynHigherOrderFn for Instance { + fn as_any(&self) -> &dyn Any { + self + } + + fn id(&self) -> ScalarFnId { + self.vtable.id() + } + + fn arity(&self) -> Arity { + self.vtable.arity(&self.options) + } + + fn child_name(&self, child_idx: usize) -> ChildName { + self.vtable.child_name(&self.options, child_idx) + } + + fn lambda_count(&self) -> usize { + self.vtable.lambda_count(&self.options) + } + + fn bind_lambdas(&self, args: &[DType], lambdas: &[Lambda]) -> VortexResult> { + self.vtable.bind_lambdas(&self.options, args, lambdas) + } + + fn return_dtype(&self, args: &[DType], lambdas: &[BoundLambda]) -> VortexResult { + self.vtable.return_dtype(&self.options, args, lambdas) + } + + fn execute( + &self, + args: &dyn ExecutionArgs, + lambdas: &[BoundLambda], + ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.vtable.execute(&self.options, args, lambdas, ctx) + } + + fn is_strict(&self) -> bool { + self.vtable.is_strict(&self.options) + } + + fn is_infallible(&self) -> bool { + self.vtable.is_infallible(&self.options) + } + + fn fmt_sql( + &self, + expr: &dyn ExprDisplay, + lambdas: &[Lambda], + f: &mut Formatter<'_>, + ) -> fmt::Result { + self.vtable.fmt_sql(&self.options, expr, lambdas, f) + } + + fn options_any(&self) -> &dyn Any { + &self.options + } + + fn options_eq(&self, other: &dyn Any) -> bool { + other + .downcast_ref::() + .is_some_and(|other| &self.options == other) + } + + fn options_hash(&self, mut hasher: &mut dyn Hasher) { + self.options.hash(&mut hasher); + } + + fn options_display(&self, f: &mut Formatter<'_>) -> fmt::Result { + Display::fmt(&self.options, f) + } + + fn options_debug(&self, f: &mut Formatter<'_>) -> fmt::Result { + Debug::fmt(&self.options, f) + } +} + +/// A type-erased higher-order function with its options, but not its lambdas: the node holds +/// those, as written or bound. +#[derive(Clone)] +pub struct HigherOrderFnRef(Arc); + +impl HigherOrderFnRef { + /// The function's id. + pub fn id(&self) -> ScalarFnId { + self.0.id() + } + + /// Whether this is the given function. + pub fn is(&self) -> bool { + self.0.as_any().is::>() + } + + /// The options, if this is the given function. + pub fn as_opt(&self) -> Option<&V::Options> { + self.0 + .as_any() + .downcast_ref::>() + .map(|instance| &instance.options) + } + + /// How many value children the function takes. + pub fn arity(&self) -> Arity { + self.0.arity() + } + + /// The name of a value child. + pub fn child_name(&self, child_idx: usize) -> ChildName { + self.0.child_name(child_idx) + } + + /// How many lambdas the function takes. + pub fn lambda_count(&self) -> usize { + self.0.lambda_count() + } + + /// Bind the lambdas against the value children's dtypes. See + /// [`HigherOrderFnVTable::bind_lambdas`]. + pub fn bind_lambdas( + &self, + args: &[DType], + lambdas: &[Lambda], + ) -> VortexResult> { + self.0.bind_lambdas(args, lambdas) + } + + /// The result dtype. + pub fn return_dtype(&self, args: &[DType], lambdas: &[BoundLambda]) -> VortexResult { + self.0.return_dtype(args, lambdas) + } + + /// Run the function. + pub fn execute( + &self, + args: &dyn ExecutionArgs, + lambdas: &[BoundLambda], + ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.0.execute(args, lambdas, ctx) + } + + /// Whether a null in any value child makes the result null. + pub fn is_strict(&self) -> bool { + self.0.is_strict() + } + + /// Whether neither the function nor any of these lambdas can fail. + pub fn is_infallible(&self, lambdas: &[BoundLambda]) -> bool { + self.0.is_infallible() && lambdas.iter().all(BoundLambda::is_infallible) + } + + /// Format the call as SQL. + pub fn fmt_sql( + &self, + expr: &dyn ExprDisplay, + lambdas: &[Lambda], + f: &mut Formatter<'_>, + ) -> fmt::Result { + self.0.fmt_sql(expr, lambdas, f) + } + + /// Bind a written node's lambdas and compute its dtype: the lambdas bound to these value + /// dtypes, and the result dtype. + pub(crate) fn bind_node( + &self, + args: &[DType], + lambdas: &[Lambda], + ) -> VortexResult<(Vec, DType)> { + let bound = self.bind_lambdas(args, lambdas)?; + let dtype = self.return_dtype(args, &bound)?; + Ok((bound, dtype)) + } +} + +impl Debug for HigherOrderFnRef { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("HigherOrderFnRef") + .field("id", &self.0.id()) + .field("options", &DebugWith(|fmt| self.0.options_debug(fmt))) + .finish() + } +} + +impl Display for HigherOrderFnRef { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{}(", self.0.id())?; + self.0.options_display(f)?; + write!(f, ")") + } +} + +impl PartialEq for HigherOrderFnRef { + fn eq(&self, other: &Self) -> bool { + self.0.id() == other.0.id() && self.0.options_eq(other.0.options_any()) + } +} + +impl Eq for HigherOrderFnRef {} + +impl Hash for HigherOrderFnRef { + fn hash(&self, state: &mut H) { + self.0.id().hash(state); + self.0.options_hash(state); + } +} + +/// Lambdas of a bound node, rebound if the value children's dtypes changed. Used when a bound +/// node is rebuilt with new children. +pub(crate) fn rebind_if_needed( + function: &HigherOrderFnRef, + args: &[DType], + written: &[Lambda], + bound: &Arc>, +) -> VortexResult<(Arc>, DType)> { + let rebound = function.bind_lambdas(args, written)?; + let rebound = if rebound.as_slice() == bound.as_slice() { + Arc::clone(bound) + } else { + Arc::new(rebound) + }; + // The result dtype can change even when the lambdas don't, such as a list's own nullability. + let dtype = function.return_dtype(args, &rebound)?; + Ok((rebound, dtype)) +} diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 4eb2620fe82..20731020d2c 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -121,6 +121,7 @@ mod expression; pub mod extension; pub mod flatbuffers; mod hash; +pub mod higher_order_fn; pub mod iter; pub mod kernel; pub mod mask; diff --git a/vortex-layout/src/layouts/row_idx/expr.rs b/vortex-layout/src/layouts/row_idx/expr.rs index 31dbbd4d47e..659eb1f69d7 100644 --- a/vortex-layout/src/layouts/row_idx/expr.rs +++ b/vortex-layout/src/layouts/row_idx/expr.rs @@ -71,3 +71,26 @@ impl ScalarFnVTable for RowIdx { pub fn row_idx() -> Expression { RowIdx.new_expr(EmptyOptions, []) } + +#[cfg(test)] +mod tests { + use vortex_array::expr::Lambda; + use vortex_array::expr::checked_add; + use vortex_array::expr::var; + + use super::row_idx; + + /// The scan replaces `row_idx()` before anything runs, but never looks inside a lambda's body, + /// so a lambda can't read it. `Lambda::new` recognizes it by its id. + #[test] + fn a_lambda_cannot_read_row_idx() { + let error = Lambda::new(["x"], checked_add(var("x"), row_idx())) + .expect_err("row_idx() is not a parameter"); + assert!( + error + .to_string() + .contains("lambda (x) can only read its parameters, not #row_idx"), + "{error}" + ); + } +}