diff --git a/vortex-array/src/arrays/higher_order_fn/array.rs b/vortex-array/src/arrays/higher_order_fn/array.rs new file mode 100644 index 00000000000..ef47ac9b0d5 --- /dev/null +++ b/vortex-array/src/arrays/higher_order_fn/array.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ArraySlots; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::TypedArrayRef; +use crate::arrays::HigherOrderFn; +use crate::higher_order_fn::HigherOrderFunctionRef; +use crate::higher_order_fn::LambdaClosure; + +/// Per-array data for [`HigherOrderFnArray`]. +#[derive(Clone, Debug)] +pub struct HigherOrderFnData { + pub(super) higher_order_fn: HigherOrderFunctionRef, + pub(super) arg_count: usize, + pub(super) lambdas: Box<[LambdaClosure]>, +} + +impl Display for HigherOrderFnData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "higher_order_fn: {}", self.higher_order_fn) + } +} + +pub trait HigherOrderFnArrayExt: TypedArrayRef { + /// The function represented by this array. + fn higher_order_fn(&self) -> &HigherOrderFunctionRef { + &self.higher_order_fn + } + + /// The number of ordinary function arguments, before capture slots. + fn arg_count(&self) -> usize { + self.arg_count + } + + /// The lambdas closed over the capture slots. + fn lambdas(&self) -> &[LambdaClosure] { + &self.lambdas + } + + /// The function's ordinary argument arrays. + fn args(&self) -> Vec { + self.as_ref().slots()[..self.arg_count()] + .iter() + .map(|slot| { + slot.as_ref() + .vortex_expect("HigherOrderFnArray argument slot") + .clone() + }) + .collect() + } + + /// Arrays captured by its lambdas, in closure-slot order. + fn captures(&self) -> Vec { + self.as_ref().slots()[self.arg_count()..] + .iter() + .map(|slot| { + slot.as_ref() + .vortex_expect("HigherOrderFnArray capture slot") + .clone() + }) + .collect() + } +} +impl> HigherOrderFnArrayExt for T {} + +impl Array { + /// Build a lazy higher-order-function array from ordinary arguments and lexical closures. + pub(crate) fn try_new_with_len( + higher_order_fn: HigherOrderFunctionRef, + args: Vec, + lambdas: Vec, + capture_slots: Vec, + len: usize, + ) -> VortexResult { + vortex_ensure!( + args.iter().all(|arg| arg.len() == len), + "HigherOrderFnArray arguments must have the array length" + ); + vortex_ensure!( + higher_order_fn.arity().matches(args.len()), + "{} takes {} ordinary arguments, got {}", + higher_order_fn, + higher_order_fn.arity(), + args.len() + ); + vortex_ensure!( + lambdas.len() == higher_order_fn.lambda_arity(), + "{} takes {} lambda arguments, got {}", + higher_order_fn, + higher_order_fn.lambda_arity(), + lambdas.len() + ); + + let arg_dtypes = args + .iter() + .map(|arg| arg.dtype().clone()) + .collect::>(); + let arg_count = args.len(); + let lambdas = lambdas.into_boxed_slice(); + let mut slots = args; + slots.extend(capture_slots); + let bound_lambdas = lambdas + .iter() + .map(|lambda| lambda.lambda().clone()) + .collect::>(); + let dtype = higher_order_fn.return_dtype(&arg_dtypes, &bound_lambdas)?; + let data = HigherOrderFnData { + higher_order_fn: higher_order_fn.clone(), + arg_count, + lambdas, + }; + let vtable = HigherOrderFn { + id: higher_order_fn.id(), + }; + Ok(unsafe { + Array::from_parts_unchecked( + ArrayParts::new(vtable, dtype, len, data) + .with_slots(slots.into_iter().map(Some).collect::()), + ) + }) + } +} diff --git a/vortex-array/src/arrays/higher_order_fn/mod.rs b/vortex-array/src/arrays/higher_order_fn/mod.rs new file mode 100644 index 00000000000..7884b18548d --- /dev/null +++ b/vortex-array/src/arrays/higher_order_fn/mod.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod array; +mod vtable; + +pub use array::HigherOrderFnArrayExt; +pub use vtable::*; diff --git a/vortex-array/src/arrays/higher_order_fn/vtable.rs b/vortex-array/src/arrays/higher_order_fn/vtable.rs new file mode 100644 index 00000000000..12c9e076a96 --- /dev/null +++ b/vortex-array/src/arrays/higher_order_fn/vtable.rs @@ -0,0 +1,221 @@ +// 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 vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; + +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayRef; +use crate::Canonical; +use crate::EqMode; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::arrays::higher_order_fn::array::HigherOrderFnArrayExt; +use crate::arrays::higher_order_fn::array::HigherOrderFnData; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::higher_order_fn::HigherOrderFunctionId; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// A lazy array produced by a higher-order function. +pub type HigherOrderFnArray = Array; + +/// The Vortex array vtable for a registered higher-order function. +#[derive(Clone, Debug)] +pub struct HigherOrderFn { + pub(super) id: HigherOrderFunctionId, +} + +impl Display for HigherOrderFn { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.id.fmt(f) + } +} + +impl ArrayHash for HigherOrderFnData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.higher_order_fn.hash(state); + self.arg_count.hash(state); + self.lambdas.hash(state); + } +} + +impl ArrayEq for HigherOrderFnData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.higher_order_fn == other.higher_order_fn + && self.arg_count == other.arg_count + && self.lambdas == other.lambdas + } +} + +impl VTable for HigherOrderFn { + type TypedArrayData = HigherOrderFnData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + self.id + } + + fn validate( + &self, + data: &HigherOrderFnData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + data.higher_order_fn.id() == self.id, + "HigherOrderFnArray data function does not match vtable" + ); + vortex_ensure!( + slots.len() >= data.arg_count, + "HigherOrderFnArray has fewer slots than ordinary arguments" + ); + vortex_ensure!( + data.higher_order_fn.arity().matches(data.arg_count), + "HigherOrderFnArray argument count does not match function arity" + ); + vortex_ensure!( + slots.iter().all(Option::is_some), + "HigherOrderFnArray slots must not be empty" + ); + vortex_ensure!( + slots.iter().flatten().all(|slot| slot.len() == len), + "HigherOrderFnArray slots must have the array length" + ); + vortex_ensure!( + data.lambdas.len() == data.higher_order_fn.lambda_arity(), + "HigherOrderFnArray lambda count does not match function arity" + ); + + let capture_count = slots.len() - data.arg_count; + for lambda in &data.lambdas { + lambda.validate(capture_count)?; + } + + let arg_dtypes = slots[..data.arg_count] + .iter() + .flatten() + .map(|slot| slot.dtype().clone()) + .collect::>(); + let lambdas = data + .lambdas + .iter() + .map(|lambda| lambda.lambda().clone()) + .collect::>(); + vortex_ensure!( + data.higher_order_fn.return_dtype(&arg_dtypes, &lambdas)? == *dtype, + "HigherOrderFnArray dtype does not match higher-order function return dtype" + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("HigherOrderFnArray 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>> { + // Runtime closures carry array references and are intentionally not serializable. + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("Deserialization of HigherOrderFnArray metadata is not supported") + } + + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + if idx < array.arg_count() { + array.higher_order_fn().child_name(idx).to_string() + } else { + format!("capture[{}]", idx - array.arg_count()) + } + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let args = array.args(); + let captures = array.captures(); + let lambdas = array + .lambdas() + .iter() + .map(|lambda| lambda.call(&captures)) + .collect::>(); + array + .higher_order_fn() + .execute(&args, &lambdas, ctx) + .map(ExecutionResult::done) + } +} + +impl OperationsVTable for HigherOrderFn { + fn scalar_at( + array: ArrayView<'_, HigherOrderFn>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .execute_scalar(index, ctx) + } +} + +impl ValidityVTable for HigherOrderFn { + fn validity(array: ArrayView<'_, HigherOrderFn>) -> VortexResult { + let args = array.args(); + let lambdas = array + .lambdas() + .iter() + .map(|lambda| lambda.lambda().clone()) + .collect::>(); + array.higher_order_fn().validity(&args, &lambdas) + } +} diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index f96eebead65..7ec3f64e84f 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -66,6 +66,10 @@ pub mod fixed_size_list; pub use fixed_size_list::FixedSizeList; pub use fixed_size_list::FixedSizeListArray; +pub mod higher_order_fn; +pub use higher_order_fn::HigherOrderFn; +pub use higher_order_fn::HigherOrderFnArray; + pub mod interleave; pub use interleave::Interleave; pub use interleave::InterleaveArray; diff --git a/vortex-array/src/expr/analysis/fallible.rs b/vortex-array/src/expr/analysis/fallible.rs index ff43d51603b..cde0b209111 100644 --- a/vortex-array/src/expr/analysis/fallible.rs +++ b/vortex-array/src/expr/analysis/fallible.rs @@ -10,8 +10,8 @@ pub fn label_is_fallible(expr: &Expression) -> BooleanLabels<'_> { expr, |expr| match expr { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_fallible(), - // The scope itself cannot fail. - Expression::Root => false, + Expression::HigherOrder { .. } => true, + Expression::Root | Expression::Variable(_) => false, }, |acc, &child| acc | child, ) diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 6c2e4975a92..bd4491af00a 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -67,7 +67,11 @@ pub fn make_bound_free_field_annotator( ) -> impl AnnotationFn { move |expr: &BoundExpression| { let Some(scalar_fn) = expr.as_scalar() else { - return scope.names().iter().cloned().collect(); + return if expr.is_root() { + scope.names().iter().cloned().collect() + } else { + vec![] + }; }; if let Some(selection) = scalar_fn.as_opt::