vortex_array/scalar_fn/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5use std::fmt;
6use std::fmt::Debug;
7use std::fmt::Display;
8use std::fmt::Formatter;
9use std::hash::Hash;
10
11use arcref::ArcRef;
12use vortex_error::VortexExpect;
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_error::vortex_err;
16use vortex_session::VortexSession;
17use vortex_utils::iter::ReduceBalancedIterExt;
18
19use crate::ArrayRef;
20use crate::ExecutionCtx;
21use crate::IntoArray;
22use crate::arrays::ConstantArray;
23use crate::arrays::ScalarFn;
24use crate::arrays::ScalarFnArray;
25use crate::dtype::DType;
26use crate::expr::BoundExpression;
27use crate::expr::Expression;
28use crate::expr::bound;
29use crate::expr::display::ExprDisplay;
30use crate::scalar::Scalar;
31use crate::scalar_fn::ScalarFnId;
32use crate::scalar_fn::ScalarFnRef;
33use crate::scalar_fn::TypedScalarFnInstance;
34use crate::scalar_fn::fns::binary::Binary;
35use crate::scalar_fn::fns::is_not_null::IsNotNull;
36use crate::scalar_fn::fns::literal::Literal;
37use crate::scalar_fn::fns::operators::Operator;
38use crate::validity::Validity;
39
40// Here and beyond we use validity(x) interchangeably with is_not_null(x).
41#[derive(Clone)]
42pub enum ReduceNodeValidity<T: ReduceNode> {
43    /// Validity of T can be symbolically reduced (i.e. without evaluating T)
44    /// to a function over validities of this node's children. For an
45    /// expression reduce node, one example is byte_length(x).
46    /// validity(byte_length(x)) can be symbolically reduced to validity(x)
47    /// since byte_length doesn't change validity.
48    Reduced(T),
49    /// Validity of T can't be symbolically reduced to anything, and all
50    /// further reductions require evaluating T first. For an expression reduce
51    /// node, one example is list_contains(x, C) since you can't reduce
52    /// validity(list_contains(x, C)). list_contains([], null) is false, but to
53    /// know that, you need to evaluate x's offsets.
54    ///
55    /// This is also the default case.
56    Irreducible,
57}
58
59/// IsNotNull(child) as a reducible node
60pub(crate) fn is_not_null_node<T: ReduceNode>(child: &T) -> VortexResult<T> {
61    child.new_node(IsNotNull.bind(EmptyOptions), std::slice::from_ref(child))
62}
63
64/// "And" over "node's" non-nullable children
65pub fn union_child_validities<T: ReduceNode>(node: &T) -> VortexResult<T> {
66    let mut parts = Vec::with_capacity(node.child_count());
67    for i in 0..node.child_count() {
68        let child = node.child(i);
69        if child.node_dtype()?.is_nullable() {
70            parts.push(is_not_null_node(&child)?);
71        }
72    }
73    let parts = parts
74        .into_iter()
75        .try_reduce_balanced(|lhs, rhs| {
76            lhs.new_node(Binary.bind(Operator::And), &[lhs.clone(), rhs])
77        })?
78        .unwrap_or_else(|| node.new_constant(true.into()));
79    Ok(parts)
80}
81
82/// This trait defines the interface for scalar function vtables, including methods for
83/// serialization, deserialization, validation, child naming, return type computation,
84/// and evaluation.
85///
86/// This trait is non-object safe and allows the implementer to make use of associated types
87/// for improved type safety, while allowing Vortex to enforce runtime checks on the inputs and
88/// outputs of each function.
89///
90/// The [`ScalarFnVTable`] trait should be implemented for a struct that holds global data across
91/// all instances of the expression. In almost all cases, this struct will be an empty unit
92/// struct, since most expressions do not require any global state.
93pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync {
94    /// Options for this expression.
95    type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash;
96
97    /// Returns the ID of the scalar function vtable.
98    fn id(&self) -> ScalarFnId;
99
100    /// Serialize the options for this expression.
101    ///
102    /// Should return `Ok(None)` if the expression is not serializable, and `Ok(vec![])` if it is
103    /// serializable but has no metadata.
104    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
105        _ = options;
106        Ok(None)
107    }
108
109    /// Deserialize the options of this expression.
110    fn deserialize(
111        &self,
112        _metadata: &[u8],
113        _session: &VortexSession,
114    ) -> VortexResult<Self::Options> {
115        vortex_bail!("Expression {} is not deserializable", self.id());
116    }
117
118    /// Returns the arity of this expression.
119    fn arity(&self, options: &Self::Options) -> Arity;
120
121    /// Returns the name of the nth child of the expr.
122    fn child_name(&self, options: &Self::Options, child_idx: usize) -> ChildName;
123
124    /// Format an expression tree in a human-readable SQL-style format.
125    ///
126    /// The expression may be either an [`Expression`] or a
127    /// [`bound expression`](crate::expr::BoundExpression).
128    fn fmt_sql(
129        &self,
130        options: &Self::Options,
131        expr: &dyn ExprDisplay,
132        f: &mut Formatter<'_>,
133    ) -> fmt::Result {
134        write!(f, "{}(", self.id())?;
135        let nchildren = expr.display_children_count();
136        for i in 0..nchildren {
137            Display::fmt(expr.display_child(i), f)?;
138            if i + 1 < nchildren {
139                write!(f, ", ")?;
140            }
141        }
142        let opts = format!("{}", options);
143        if !opts.is_empty() {
144            write!(f, ", opts={}", opts)?;
145        }
146        write!(f, ")")
147    }
148
149    /// Compute the return [`DType`] of the expression if evaluated over the given input types.
150    ///
151    /// # Preconditions
152    ///
153    /// The length of `args` must match the [`Arity`] of this function. Callers are responsible
154    /// for validating this (e.g., [`Expression::try_new`] checks arity at construction time).
155    /// Implementations may assume correct arity and will panic or return nonsensical results if
156    /// violated.
157    ///
158    /// [`Expression::try_new`]: crate::expr::Expression::try_new
159    fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult<DType>;
160
161    /// Execute the expression over the input arguments.
162    ///
163    /// Implementations are encouraged to check their inputs for constant arrays to perform
164    /// more optimized execution.
165    ///
166    /// If the input arguments cannot be directly used for execution (for example, an expression
167    /// may require canonical input arrays), then the implementation should perform a single
168    /// child execution and return a new [`crate::arrays::ScalarFnArray`] wrapping up the new child.
169    ///
170    /// This provides maximum opportunities for array-level optimizations using execute_parent
171    /// kernels.
172    fn execute(
173        &self,
174        options: &Self::Options,
175        args: &dyn ExecutionArgs,
176        ctx: &mut ExecutionCtx,
177    ) -> VortexResult<ArrayRef>;
178
179    /// Implement an abstract reduction rule over a tree of scalar functions.
180    ///
181    /// The [`ReduceNode`] can be used to traverse children, inspect their types, and
182    /// construct the result via [`ReduceNode::new_node`]. The rule is generic over the node
183    /// type and is instantiated once per reducible tree kind (expressions and arrays).
184    ///
185    /// Return `Ok(None)` if no reduction is possible.
186    fn reduce<T: ReduceNode>(&self, options: &Self::Options, node: &T) -> VortexResult<Option<T>> {
187        _ = options;
188        _ = node;
189        Ok(None)
190    }
191
192    /// For node, returns node' which is exactly the result of evaluating
193    /// validity(node). Returned node' is either a lazy computation over
194    /// children of node, a constant, or Irreducible which means you need to
195    /// evaluate node to get its validity.
196    fn validity<T: ReduceNode>(
197        &self,
198        options: &Self::Options,
199        node: &T,
200    ) -> VortexResult<ReduceNodeValidity<T>> {
201        if !self.is_strict(options) {
202            return Ok(ReduceNodeValidity::Irreducible);
203        }
204
205        let mut dtypes = Vec::with_capacity(node.child_count());
206        for i in 0..node.child_count() {
207            let dtype = node.child(i).node_dtype()?;
208            if matches!(dtype, DType::Null) {
209                return Ok(ReduceNodeValidity::Irreducible);
210            }
211            dtypes.push(dtype.as_nonnullable());
212        }
213
214        let res = if let Ok(dtype) = self.return_dtype(options, &dtypes)
215            && !dtype.is_nullable()
216        {
217            ReduceNodeValidity::Reduced(union_child_validities(node)?)
218        } else {
219            ReduceNodeValidity::Irreducible
220        };
221        Ok(res)
222    }
223
224    /// Simplify the bound expression if possible.
225    ///
226    /// Every node of `expr` carries its dtype, so rules read types directly from the tree.
227    fn simplify(
228        &self,
229        options: &Self::Options,
230        expr: &BoundExpression,
231    ) -> VortexResult<Option<BoundExpression>> {
232        _ = options;
233        _ = expr;
234        Ok(None)
235    }
236
237    /// Returns whether this scalar function is strict.
238    ///
239    /// A strict function returns null for a row when any argument is null for that row. This
240    /// matches [PostgreSQL's `STRICT` convention](https://www.postgresql.org/docs/current/sql-createfunction.html)
241    /// for null propagation.
242    ///
243    /// Return `true` only when this holds for every argument. `add` is strict, but Kleene `AND`
244    /// is not because `false AND null` returns `false`. `is_null` is also not strict.
245    ///
246    /// Strictness does not require valid inputs to produce a valid output. For example,
247    /// [`crate::expr::list_sum`] returns null for a valid empty list. Implement
248    /// [`ScalarFnVTable::validity`] only when the output validity can be derived without
249    /// evaluation.
250    ///
251    /// [`ScalarFnVTable::return_dtype`] must return a nullable output dtype when any input dtype is
252    /// nullable. A `cast` that forces a non-nullable output dtype is therefore not strict.
253    ///
254    /// This property applies only to the scalar function, not its child expressions. Nullary
255    /// functions are vacuously strict. The default is conservatively `false`.
256    fn is_strict(&self, options: &Self::Options) -> bool {
257        _ = options;
258        false
259    }
260
261    /// Returns whether this scalar function can never raise a semantic error.
262    ///
263    /// Return `true` only when a well-typed call cannot error because of its values. `checked_add`
264    /// is fallible on integer overflow, and integer division is fallible when its divisor is zero.
265    /// A null result is not an error: [`crate::expr::list_sum`] is infallible for an empty list.
266    ///
267    /// Ignore incidental execution errors, such as canonicalization failures, allocation errors,
268    /// and encoding mismatches. They are not part of the function's semantics.
269    ///
270    /// Returning `true` permits optimizations that evaluate the function over values that no input
271    /// row references. Dictionary push-down, for example, evaluates every dictionary value, so a
272    /// fallible function could error on a value that row-wise evaluation would never reach.
273    ///
274    /// This applies only to the scalar function, not its child expressions, and only to inputs
275    /// accepted by [`ScalarFnVTable::return_dtype`]. The default is conservatively `false`.
276    fn is_infallible(&self, options: &Self::Options) -> bool {
277        _ = options;
278        false
279    }
280}
281
282/// A node used for implementing abstract reduction rules over a tree of scalar functions.
283///
284/// Reduction rules are generic over the node type, so a rule is written once and monomorphized
285/// per reducible tree kind: [`ExpressionReduceNode`] for bound expression trees and
286/// [`ArrayReduceNode`] for array trees. Nodes borrow from the tree being reduced, making
287/// traversal allocation-free, while nodes produced by [`ReduceNode::new_node`] own their
288/// freshly-built subtrees.
289pub trait ReduceNode: Clone {
290    /// Return the data type of this node.
291    fn node_dtype(&self) -> VortexResult<DType>;
292
293    /// Return this node's scalar function if it is indeed a scalar fn.
294    fn scalar_fn(&self) -> Option<&ScalarFnRef>;
295
296    /// Descend to the child of this node.
297    fn child(&self, idx: usize) -> Self;
298
299    /// Returns the number of children of this node.
300    fn child_count(&self) -> usize;
301
302    /// Create a new node from the given scalar function and children, inheriting this node's
303    /// reduction context (e.g. the expression scope, or the array row count).
304    fn new_node(&self, scalar_fn: ScalarFnRef, children: &[Self]) -> VortexResult<Self>;
305
306    /// Return a scalar value if this node is constant
307    fn as_constant(&self) -> Option<Scalar> {
308        None
309    }
310
311    /// Produce a new constant node in the same scope as "self"
312    fn new_constant(&self, value: Scalar) -> Self;
313
314    /// Symbolic validity of this node. Reduced() if you can get from node's
315    /// validity to validity of its children or a constant without evaluating
316    /// node.
317    fn validity(&self) -> VortexResult<ReduceNodeValidity<Self>>
318    where
319        Self: Sized;
320}
321
322/// A [`ReduceNode`] over a bound expression tree.
323#[derive(Clone)]
324pub struct ExpressionReduceNode<'a> {
325    expression: Cow<'a, BoundExpression>,
326}
327
328impl<'a> ExpressionReduceNode<'a> {
329    /// Creates a node borrowing the given bound expression.
330    pub fn new(expression: &'a BoundExpression) -> Self {
331        Self {
332            expression: Cow::Borrowed(expression),
333        }
334    }
335
336    /// Returns the bound expression backing this node.
337    pub fn expression(&self) -> &BoundExpression {
338        &self.expression
339    }
340
341    /// Consumes this node and returns the backing bound expression.
342    pub fn into_expression(self) -> BoundExpression {
343        self.expression.into_owned()
344    }
345}
346
347impl ReduceNode for ExpressionReduceNode<'_> {
348    fn node_dtype(&self) -> VortexResult<DType> {
349        Ok(self.expression.dtype().clone())
350    }
351
352    fn scalar_fn(&self) -> Option<&ScalarFnRef> {
353        self.expression.as_scalar()
354    }
355
356    fn child(&self, idx: usize) -> Self {
357        let expression = match &self.expression {
358            Cow::Borrowed(expression) => Cow::Borrowed(expression.child(idx)),
359            Cow::Owned(expression) => Cow::Owned(expression.child(idx).clone()),
360        };
361        Self { expression }
362    }
363
364    fn child_count(&self) -> usize {
365        self.expression.children().len()
366    }
367
368    fn new_node(&self, scalar_fn: ScalarFnRef, children: &[Self]) -> VortexResult<Self> {
369        let expression = BoundExpression::try_new(
370            scalar_fn,
371            children.iter().map(|c| c.expression.as_ref().clone()),
372        )?;
373        Ok(Self {
374            expression: Cow::Owned(expression),
375        })
376    }
377
378    fn as_constant(&self) -> Option<Scalar> {
379        self.expression.as_opt::<Literal>().cloned()
380    }
381
382    fn new_constant(&self, value: Scalar) -> Self {
383        Self {
384            expression: Cow::Owned(bound::lit(value)),
385        }
386    }
387
388    fn validity(&self) -> VortexResult<ReduceNodeValidity<Self>> {
389        match self.expression.as_scalar() {
390            Some(scalar_fn) => scalar_fn.validity_expression(self),
391            None => Ok(ReduceNodeValidity::Irreducible),
392        }
393    }
394}
395
396/// A [`ReduceNode`] over an array tree.
397#[derive(Clone)]
398pub struct ArrayReduceNode<'a> {
399    array: Cow<'a, ArrayRef>,
400}
401
402impl<'a> ArrayReduceNode<'a> {
403    /// Creates a node borrowing the given array.
404    pub fn new(array: &'a ArrayRef) -> Self {
405        Self {
406            array: Cow::Borrowed(array),
407        }
408    }
409
410    /// Returns the array backing this node.
411    pub fn array(&self) -> &ArrayRef {
412        &self.array
413    }
414
415    /// Consumes this node and returns the backing array.
416    pub fn into_array(self) -> ArrayRef {
417        self.array.into_owned()
418    }
419}
420
421impl ReduceNode for ArrayReduceNode<'_> {
422    fn node_dtype(&self) -> VortexResult<DType> {
423        Ok(self.array.dtype().clone())
424    }
425
426    fn scalar_fn(&self) -> Option<&ScalarFnRef> {
427        self.array
428            .as_opt::<ScalarFn>()
429            .map(|a| a.data().scalar_fn())
430    }
431
432    fn child(&self, idx: usize) -> Self {
433        let array = match &self.array {
434            Cow::Borrowed(array) => Cow::Borrowed(
435                array
436                    .nth_child(idx)
437                    .vortex_expect("child idx out of bounds"),
438            ),
439            Cow::Owned(array) => Cow::Owned(
440                array
441                    .nth_child(idx)
442                    .vortex_expect("child idx out of bounds")
443                    .clone(),
444            ),
445        };
446        Self { array }
447    }
448
449    fn child_count(&self) -> usize {
450        self.array.nchildren()
451    }
452
453    fn new_node(&self, scalar_fn: ScalarFnRef, children: &[Self]) -> VortexResult<Self> {
454        let array = ScalarFnArray::try_new_with_len(
455            scalar_fn,
456            children.iter().map(|c| c.array.as_ref().clone()).collect(),
457            self.array.len(),
458        )?;
459        Ok(Self {
460            array: Cow::Owned(array.into_array()),
461        })
462    }
463
464    fn as_constant(&self) -> Option<Scalar> {
465        self.array.as_constant()
466    }
467
468    fn new_constant(&self, value: Scalar) -> Self {
469        let array = ConstantArray::new(value, self.array.len());
470        Self {
471            array: Cow::Owned(array.into_array()),
472        }
473    }
474
475    fn validity(&self) -> VortexResult<ReduceNodeValidity<Self>> {
476        if let Some(scalar_fn) = self.array.as_opt::<ScalarFn>() {
477            return scalar_fn.data().scalar_fn().validity_array(self);
478        }
479        Ok(ReduceNodeValidity::Reduced(
480            match self.array.validity()? {
481                Validity::NonNullable | Validity::AllValid => self.new_constant(true.into()),
482                Validity::AllInvalid => self.new_constant(false.into()),
483                Validity::Array(array) => Self {
484                    array: Cow::Owned(array),
485                },
486            },
487        ))
488    }
489}
490
491/// The arity (number of arguments) of a function.
492#[derive(Clone, Copy, Debug, PartialEq, Eq)]
493pub enum Arity {
494    Exact(usize),
495    Variadic { min: usize, max: Option<usize> },
496}
497
498impl Display for Arity {
499    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
500        match self {
501            Arity::Exact(n) => write!(f, "{}", n),
502            Arity::Variadic { min, max } => match max {
503                Some(max) if min == max => write!(f, "{}", min),
504                Some(max) => write!(f, "{}..{}", min, max),
505                None => write!(f, "{}+", min),
506            },
507        }
508    }
509}
510
511impl Arity {
512    /// Whether the given argument count matches this arity.
513    pub fn matches(&self, arg_count: usize) -> bool {
514        match self {
515            Arity::Exact(m) => *m == arg_count,
516            Arity::Variadic { min, max } => {
517                if arg_count < *min {
518                    return false;
519                }
520                if let Some(max) = max
521                    && arg_count > *max
522                {
523                    return false;
524                }
525                true
526            }
527        }
528    }
529}
530
531/// Arguments for expression execution.
532pub trait ExecutionArgs {
533    /// Returns the input array at the given index.
534    fn get(&self, index: usize) -> VortexResult<ArrayRef>;
535
536    /// Returns the number of inputs.
537    fn num_inputs(&self) -> usize;
538
539    /// Returns the row count of the execution scope.
540    fn row_count(&self) -> usize;
541}
542
543/// A concrete [`ExecutionArgs`] backed by a `Vec<ArrayRef>`.
544pub struct VecExecutionArgs {
545    inputs: Vec<ArrayRef>,
546    row_count: usize,
547}
548
549impl VecExecutionArgs {
550    /// Create a new `VecExecutionArgs`.
551    pub fn new(inputs: Vec<ArrayRef>, row_count: usize) -> Self {
552        Self { inputs, row_count }
553    }
554}
555
556impl ExecutionArgs for VecExecutionArgs {
557    fn get(&self, index: usize) -> VortexResult<ArrayRef> {
558        self.inputs.get(index).cloned().ok_or_else(|| {
559            vortex_err!(
560                "Input index {} out of bounds (num_inputs={})",
561                index,
562                self.inputs.len()
563            )
564        })
565    }
566
567    fn num_inputs(&self) -> usize {
568        self.inputs.len()
569    }
570
571    fn row_count(&self) -> usize {
572        self.row_count
573    }
574}
575
576#[derive(Clone, Debug, PartialEq, Eq, Hash)]
577pub struct EmptyOptions;
578impl Display for EmptyOptions {
579    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
580        write!(f, "")
581    }
582}
583
584/// Factory functions for vtables.
585pub trait ScalarFnVTableExt: ScalarFnVTable {
586    /// Bind this vtable with the given options into a [`ScalarFnRef`].
587    fn bind(&self, options: Self::Options) -> ScalarFnRef {
588        TypedScalarFnInstance::new(self.clone(), options).erased()
589    }
590
591    /// Create a new expression with this vtable and the given options and children.
592    fn new_expr(
593        &self,
594        options: Self::Options,
595        children: impl IntoIterator<Item = Expression>,
596    ) -> Expression {
597        Self::try_new_expr(self, options, children).vortex_expect("Failed to create expression")
598    }
599
600    /// Try to create a new expression with this vtable and the given options and children.
601    fn try_new_expr(
602        &self,
603        options: Self::Options,
604        children: impl IntoIterator<Item = Expression>,
605    ) -> VortexResult<Expression> {
606        Expression::try_new(self.bind(options), children)
607    }
608
609    /// Try to create a bound expression with this vtable, the given options, and bound children.
610    fn try_new_bound_expr(
611        &self,
612        options: Self::Options,
613        children: impl IntoIterator<Item = BoundExpression>,
614    ) -> VortexResult<BoundExpression> {
615        BoundExpression::try_new(self.bind(options), children)
616    }
617}
618impl<V: ScalarFnVTable> ScalarFnVTableExt for V {}
619
620/// A reference to the name of a child expression.
621pub type ChildName = ArcRef<str>;
vortex_array/scalar_fn/vtable.rs

vortex_array/scalar_fn/
vtable.rs
