Skip to main content

vortex_array/array/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6
7use crate::ExecutionCtx;
8use crate::array::ArrayView;
9use crate::array::VTable;
10use crate::array::probe::ProbeState;
11use crate::scalar::Scalar;
12use crate::vtable::NotSupported;
13
14/// Element-level operations for an array encoding.
15///
16/// This trait is separated from [`VTable`] so encodings can organize scalar
17/// access independently from traversal, serialization, and execution. The erased
18/// [`ArrayRef`](crate::ArrayRef)
19/// methods perform common checks before dispatching here.
20pub trait OperationsVTable<V: VTable> {
21    /// Encoding-specific state retained by repeated scalar access.
22    ///
23    /// Built once per repeated probe and never for a one-off read. Preparation belongs in
24    /// [`Self::probe_scalar`]. State owns its preparation and may hold shared buffer or array
25    /// handles. Use `()` when no state is needed.
26    type ProbeState: Default + 'static;
27
28    /// Read the scalar at `index` of the array in `state`, including its nullness.
29    ///
30    /// Caller must guarantee `index < len`. The row may be null, and a null row must come back
31    /// as `Scalar::null` of the array's dtype. Check the array's own validity with
32    /// [`ProbeState::is_valid`].
33    ///
34    /// `state` carries the typed view of the array and, for a read through a
35    /// [`RepeatedArrayProbe`](crate::RepeatedArrayProbe), the state that probe keeps. Read
36    /// children through [`ProbeState::slot`], which follows the read's policy without the
37    /// encoding having to know it. Take encoding state from [`ProbeState::retained`]. The scalar
38    /// must retain the source's logical dtype, including nullability.
39    ///
40    /// The default checks validity and then reads through [`Self::scalar_at`].
41    fn probe_scalar(
42        state: &mut ProbeState<'_, V>,
43        index: usize,
44        ctx: &mut ExecutionCtx,
45    ) -> VortexResult<Scalar> {
46        // FIXME: Remove this default once all encodings have migrated to probe_scalar.
47        if !state.is_valid(index, ctx)? {
48            return Ok(Scalar::null(state.array().dtype().clone()));
49        }
50        Self::scalar_at(state.array(), index, ctx)
51    }
52
53    // FIXME: Deprecate scalar_at once encodings have migrated to probe_scalar.
54    /// Fetch the scalar at the given index.
55    ///
56    /// ## Preconditions
57    ///
58    /// Bounds-checking has already been performed by the time this function is called,
59    /// and the index is guaranteed to be non-null. Implementations may assume `index < len`.
60    ///
61    /// ## Postconditions
62    ///
63    /// The returned [`Scalar`] must have the same logical dtype as the array's element dtype.
64    fn scalar_at(
65        array: ArrayView<'_, V>,
66        index: usize,
67        ctx: &mut ExecutionCtx,
68    ) -> VortexResult<Scalar>;
69}
70
71impl<V: VTable> OperationsVTable<V> for NotSupported {
72    type ProbeState = ();
73
74    fn probe_scalar(
75        state: &mut ProbeState<'_, V>,
76        index: usize,
77        ctx: &mut ExecutionCtx,
78    ) -> VortexResult<Scalar> {
79        Self::scalar_at(state.array(), index, ctx)
80    }
81
82    fn scalar_at(
83        array: ArrayView<'_, V>,
84        _index: usize,
85        _ctx: &mut ExecutionCtx,
86    ) -> VortexResult<Scalar> {
87        vortex_bail!(
88            "Legacy scalar_at operation is not supported for {} arrays",
89            array.encoding_id()
90        )
91    }
92}