|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! Scalar functions computed one row at a time. |
| 5 | +
|
| 6 | +use std::fmt::Debug; |
| 7 | +use std::fmt::Display; |
| 8 | +use std::hash::Hash; |
| 9 | + |
| 10 | +use vortex_error::VortexResult; |
| 11 | +use vortex_error::vortex_bail; |
| 12 | +use vortex_session::VortexSession; |
| 13 | + |
| 14 | +use super::visitor::RowVisitor; |
| 15 | +use crate::ArrayRef; |
| 16 | +use crate::ExecutionCtx; |
| 17 | +use crate::dtype::DType; |
| 18 | +use crate::scalar_fn::ScalarFnId; |
| 19 | +use crate::scalar_fn::unstable::row::RowExecution; |
| 20 | + |
| 21 | +/// A scalar function computed one row at a time. |
| 22 | +/// |
| 23 | +/// Declare the argument names and use [`dispatch`](Self::dispatch) to choose concrete element and |
| 24 | +/// sink types for each accepted dtype combination. |
| 25 | +/// |
| 26 | +/// Every `RowFn` receives the standard [`ScalarFnVTable`] implementation. A function that needs |
| 27 | +/// custom scalar-function hooks instead implements `ScalarFnVTable` on its public type and |
| 28 | +/// delegates row execution to a private `RowFn` kernel with [`row_fn_return_dtype`] and |
| 29 | +/// [`execute_rows`]. Implement only `ScalarFnVTable` when the natural kernel is columnar. |
| 30 | +/// |
| 31 | +/// [`ScalarFnVTable`]: crate::scalar_fn::ScalarFnVTable |
| 32 | +/// [`execute_rows`]: crate::scalar_fn::unstable::row::execute_rows |
| 33 | +/// [`row_fn_return_dtype`]: crate::scalar_fn::unstable::row::row_fn_return_dtype |
| 34 | +pub trait RowFn: 'static + Sized + Clone + Send + Sync { |
| 35 | + /// Options for this function, or [`EmptyOptions`](crate::scalar_fn::EmptyOptions) for none. |
| 36 | + type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; |
| 37 | + |
| 38 | + /// The arguments in display order. Its length is the function's exact arity. |
| 39 | + const ARG_NAMES: &'static [&'static str]; |
| 40 | + |
| 41 | + /// Whether any legal dispatch or encoding-aware execution can raise a semantic error. |
| 42 | + /// |
| 43 | + /// The framework checks this at compile time for every fallible dispatched element or result. |
| 44 | + /// A conservative `true` is allowed when only some dtype choices are fallible. |
| 45 | + /// |
| 46 | + /// Set this to `true` when [`reduce_encoded`](Self::reduce_encoded) can return a semantic |
| 47 | + /// error or [`RowExecution::DeferredError`]. Compile-time checks cover dispatched element and |
| 48 | + /// result types, but cannot inspect that hook. |
| 49 | + /// [`OutputSink`](crate::scalar_fn::unstable::row::OutputSink) lifecycle methods only report |
| 50 | + /// incidental execution failures. Semantic sink errors must come from the row callback. |
| 51 | + /// |
| 52 | + /// Semantic errors are defined by |
| 53 | + /// [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible). |
| 54 | + const FALLIBLE: bool = false; |
| 55 | + |
| 56 | + /// Returns the ID of the scalar function. |
| 57 | + fn id(&self) -> ScalarFnId; |
| 58 | + |
| 59 | + /// Serialize this function's options, or return `None` when the function is not serializable. |
| 60 | + fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> { |
| 61 | + _ = options; |
| 62 | + Ok(None) |
| 63 | + } |
| 64 | + |
| 65 | + /// Restore options written by [`serialize`](Self::serialize). |
| 66 | + fn deserialize( |
| 67 | + &self, |
| 68 | + _metadata: &[u8], |
| 69 | + _session: &VortexSession, |
| 70 | + ) -> VortexResult<Self::Options> { |
| 71 | + vortex_bail!("Expression {} is not deserializable", self.id()) |
| 72 | + } |
| 73 | + |
| 74 | + /// Choose element types for these input dtypes and visit the framework with them. |
| 75 | + /// |
| 76 | + /// Plan time and run time both call this method, so the choice **must** be a pure function of |
| 77 | + /// `options` and `args`. The framework rejects a change to the derived nullable execution |
| 78 | + /// policy before row execution. It cannot compare the remaining types, preparation values, or |
| 79 | + /// closure behavior, so those must also remain stable. Cross-argument dtype validation belongs |
| 80 | + /// here. |
| 81 | + fn dispatch<V: RowVisitor<Self::Options>>( |
| 82 | + &self, |
| 83 | + options: &Self::Options, |
| 84 | + args: &[DType], |
| 85 | + visitor: V, |
| 86 | + ) -> VortexResult<V::VisitResult>; |
| 87 | + |
| 88 | + /// Try an encoding-aware implementation before decoding the inputs into row elements. |
| 89 | + /// |
| 90 | + /// `None` continues to the dispatched row loop. [`Output`](RowExecution::Output) skips that |
| 91 | + /// loop and may remain encoded or lazy. [`DeferredError`](RowExecution::DeferredError) reruns |
| 92 | + /// only valid rows when null payloads may have caused the failure. For non-nullary functions, |
| 93 | + /// batch execution calls this hook at most once with the original, unfiltered arrays; slices |
| 94 | + /// and compacted retries do not reach it. |
| 95 | + /// |
| 96 | + /// Like a dense row closure, this hook must be total over every stored payload, including |
| 97 | + /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or |
| 98 | + /// retried through the row layer. |
| 99 | + /// |
| 100 | + /// # Requirements |
| 101 | + /// |
| 102 | + /// - `output.len()` **must** equal `args[0].len()`. |
| 103 | + /// - The output dtype **must** match the planned dtype when ignoring nullability. |
| 104 | + /// - The output **must not** introduce a null where every input is valid. |
| 105 | + /// |
| 106 | + /// The framework skips this hook for nullary functions. |
| 107 | + fn reduce_encoded( |
| 108 | + &self, |
| 109 | + options: &Self::Options, |
| 110 | + args: &[ArrayRef], |
| 111 | + ctx: &mut ExecutionCtx, |
| 112 | + ) -> VortexResult<Option<RowExecution>> { |
| 113 | + _ = (options, args, ctx); |
| 114 | + Ok(None) |
| 115 | + } |
| 116 | +} |
0 commit comments