HOF and ListTransform

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<HigherOrderFn> {
+    /// 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<ArrayRef> {
+        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<ArrayRef> {
+        self.as_ref().slots()[self.arg_count()..]
+            .iter()
+            .map(|slot| {
+                slot.as_ref()
+                    .vortex_expect("HigherOrderFnArray capture slot")
+                    .clone()
+            })
+            .collect()
+    }
+}
+impl<T: TypedArrayRef<HigherOrderFn>> HigherOrderFnArrayExt for T {}
+
+impl Array<HigherOrderFn> {
+    /// 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<ArrayRef>,
+        lambdas: Vec<LambdaClosure>,
+        capture_slots: Vec<ArrayRef>,
+        len: usize,
+    ) -> VortexResult<Self> {
+        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::<Vec<_>>();
+        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::<Vec<_>>();
+        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::<ArraySlots>()),
+            )
+        })
+    }
+}

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::*;

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<HigherOrderFn>;
+
+/// 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<H: Hasher>(&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<ArrayRef>],
+    ) -> 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::<Vec<_>>();
+        let lambdas = data
+            .lambdas
+            .iter()
+            .map(|lambda| lambda.lambda().clone())
+            .collect::<Vec<_>>();
+        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<String> {
+        None
+    }
+
+    fn with_buffers(
+        &self,
+        array: ArrayView<'_, Self>,
+        buffers: &[BufferHandle],
+    ) -> VortexResult<ArrayParts<Self>> {
+        with_empty_buffers(self, array, buffers)
+    }
+
+    fn serialize(
+        _array: ArrayView<'_, Self>,
+        _session: &VortexSession,
+    ) -> VortexResult<Option<Vec<u8>>> {
+        // 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<ArrayParts<Self>> {
+        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<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        let args = array.args();
+        let captures = array.captures();
+        let lambdas = array
+            .lambdas()
+            .iter()
+            .map(|lambda| lambda.call(&captures))
+            .collect::<Vec<_>>();
+        array
+            .higher_order_fn()
+            .execute(&args, &lambdas, ctx)
+            .map(ExecutionResult::done)
+    }
+}
+
+impl OperationsVTable<HigherOrderFn> for HigherOrderFn {
+    fn scalar_at(
+        array: ArrayView<'_, HigherOrderFn>,
+        index: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<crate::scalar::Scalar> {
+        array
+            .array()
+            .clone()
+            .execute::<Canonical>(ctx)?
+            .into_array()
+            .execute_scalar(index, ctx)
+    }
+}
+
+impl ValidityVTable<HigherOrderFn> for HigherOrderFn {
+    fn validity(array: ArrayView<'_, HigherOrderFn>) -> VortexResult<Validity> {
+        let args = array.args();
+        let lambdas = array
+            .lambdas()
+            .iter()
+            .map(|lambda| lambda.lambda().clone())
+            .collect::<Vec<_>>();
+        array.higher_order_fn().validity(&args, &lambdas)
+    }
+}

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;

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,
     )