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}