vortex_array/arrays/shared/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hasher;
5
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_error::vortex_panic;
9use vortex_session::VortexSession;
10use vortex_session::registry::CachedId;
11
12use crate::ArrayEq;
13use crate::ArrayHash;
14use crate::ArrayParts;
15use crate::ArrayRef;
16use crate::Canonical;
17use crate::EqMode;
18use crate::ExecutionCtx;
19use crate::ExecutionResult;
20use crate::array::Array;
21use crate::array::ArrayId;
22use crate::array::ArrayView;
23use crate::array::OperationsVTable;
24use crate::array::ProbeState;
25use crate::array::RepeatedArrayProbe;
26use crate::array::VTable;
27use crate::array::ValidityVTable;
28use crate::array::with_empty_buffers;
29use crate::arrays::shared::SharedArrayExt;
30use crate::arrays::shared::SharedData;
31use crate::arrays::shared::SharedSlots;
32use crate::buffer::BufferHandle;
33use crate::dtype::DType;
34use crate::scalar::Scalar;
35use crate::validity::Validity;
36
37/// A [`Shared`]-encoded Vortex array.
38pub type SharedArray = Array<Shared>;
39
40// TODO(ngates): consider hooking Shared into the iterative execution model. Cache either the
41//  most executed, or after each iteration, and return a shared cache for each execution.
42#[derive(Clone, Debug)]
43pub struct Shared;
44
45impl ArrayHash for SharedData {
46    fn array_hash<H: Hasher>(&self, _state: &mut H, _accuracy: EqMode) {}
47}
48
49impl ArrayEq for SharedData {
50    fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool {
51        true
52    }
53}
54
55impl VTable for Shared {
56    type TypedArrayData = SharedData;
57    type OperationsVTable = Self;
58    type ValidityVTable = Self;
59    fn id(&self) -> ArrayId {
60        static ID: CachedId = CachedId::new("vortex.shared");
61        *ID
62    }
63
64    fn validate(
65        &self,
66        _data: &SharedData,
67        dtype: &DType,
68        len: usize,
69        slots: &[Option<ArrayRef>],
70    ) -> VortexResult<()> {
71        let source = slots[SharedSlots::SOURCE]
72            .as_ref()
73            .vortex_expect("SharedArray source slot must be present");
74        vortex_error::vortex_ensure_eq!(source.dtype(), dtype, "SharedArray dtype mismatch");
75        vortex_error::vortex_ensure_eq!(source.len(), len, "SharedArray len mismatch");
76        Ok(())
77    }
78
79    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
80        0
81    }
82
83    fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle {
84        vortex_panic!("SharedArray has no buffers")
85    }
86
87    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
88        None
89    }
90
91    fn with_buffers(
92        &self,
93        array: ArrayView<'_, Self>,
94        buffers: &[BufferHandle],
95    ) -> VortexResult<ArrayParts<Self>> {
96        with_empty_buffers(self, array, buffers)
97    }
98
99    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
100        SharedSlots::NAMES[idx].to_string()
101    }
102
103    fn serialize(
104        _array: ArrayView<'_, Self>,
105        _session: &VortexSession,
106    ) -> VortexResult<Option<Vec<u8>>> {
107        vortex_error::vortex_bail!("Shared array is not serializable")
108    }
109
110    fn deserialize(
111        &self,
112        _dtype: &DType,
113        _len: usize,
114        _metadata: &[u8],
115
116        _buffers: &[BufferHandle],
117        _children: &dyn crate::serde::ArrayChildren,
118        _session: &VortexSession,
119    ) -> VortexResult<ArrayParts<Self>> {
120        vortex_error::vortex_bail!("Shared array is not serializable")
121    }
122
123    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
124        array
125            .get_or_compute(|source| source.clone().execute::<Canonical>(ctx))
126            .map(ExecutionResult::done)
127    }
128}
129
130/// A probe over whichever array [`SharedArrayExt::current_array_ref`] resolves to. The source
131/// can be replaced by the computed result between reads, so the retained probe is rebuilt when
132/// the array it was built over is no longer the current one.
133#[derive(Default)]
134pub struct SharedProbeState {
135    current: Option<RepeatedArrayProbe>,
136}
137
138impl OperationsVTable<Shared> for Shared {
139    type ProbeState = SharedProbeState;
140
141    fn probe_scalar(
142        state: &mut ProbeState<'_, Shared>,
143        index: usize,
144        ctx: &mut ExecutionCtx,
145    ) -> VortexResult<Scalar> {
146        let array = state.array();
147        let current = array.current_array_ref();
148        let Some(retained) = state.retained() else {
149            return current.probe().execute_scalar(index, ctx);
150        };
151        let probe = match &mut retained.current {
152            Some(probe) if ArrayRef::ptr_eq(probe.array(), current) => probe,
153            slot => slot.insert(RepeatedArrayProbe::new(current.clone())),
154        };
155        probe.execute_scalar(index, ctx)
156    }
157
158    fn scalar_at(
159        array: ArrayView<'_, Shared>,
160        index: usize,
161        ctx: &mut ExecutionCtx,
162    ) -> VortexResult<Scalar> {
163        Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
164    }
165}
166
167impl ValidityVTable<Shared> for Shared {
168    fn validity(array: ArrayView<'_, Shared>) -> VortexResult<Validity> {
169        array.current_array_ref().validity()
170    }
171}
vortex_array/arrays/shared/vtable.rs

vortex_array/arrays/shared/
vtable.rs
