Skip to main content

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;
10
11use crate::ArrayEq;
12use crate::ArrayHash;
13use crate::ArrayRef;
14use crate::Canonical;
15use crate::ExecutionCtx;
16use crate::ExecutionResult;
17use crate::Precision;
18use crate::array::Array;
19use crate::array::ArrayId;
20use crate::array::ArrayView;
21use crate::array::OperationsVTable;
22use crate::array::VTable;
23use crate::array::ValidityVTable;
24use crate::arrays::shared::SharedArrayExt;
25use crate::arrays::shared::SharedData;
26use crate::arrays::shared::array::SLOT_NAMES;
27use crate::buffer::BufferHandle;
28use crate::dtype::DType;
29use crate::scalar::Scalar;
30use crate::validity::Validity;
31
32/// A [`Shared`]-encoded Vortex array.
33pub type SharedArray = Array<Shared>;
34
35// TODO(ngates): consider hooking Shared into the iterative execution model. Cache either the
36//  most executed, or after each iteration, and return a shared cache for each execution.
37#[derive(Clone, Debug)]
38pub struct Shared;
39
40impl Shared {
41    pub const ID: ArrayId = ArrayId::new_ref("vortex.shared");
42}
43
44impl ArrayHash for SharedData {
45    fn array_hash<H: Hasher>(&self, _state: &mut H, _precision: Precision) {}
46}
47
48impl ArrayEq for SharedData {
49    fn array_eq(&self, _other: &Self, _precision: Precision) -> bool {
50        true
51    }
52}
53
54impl VTable for Shared {
55    type ArrayData = SharedData;
56    type OperationsVTable = Self;
57    type ValidityVTable = Self;
58
59    fn id(&self) -> ArrayId {
60        Self::ID
61    }
62
63    fn validate(
64        &self,
65        _data: &SharedData,
66        dtype: &DType,
67        len: usize,
68        slots: &[Option<ArrayRef>],
69    ) -> VortexResult<()> {
70        let source = slots[0]
71            .as_ref()
72            .vortex_expect("SharedArray source slot must be present");
73        vortex_error::vortex_ensure!(source.dtype() == dtype, "SharedArray dtype mismatch");
74        vortex_error::vortex_ensure!(source.len() == len, "SharedArray len mismatch");
75        Ok(())
76    }
77
78    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
79        0
80    }
81
82    fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle {
83        vortex_panic!("SharedArray has no buffers")
84    }
85
86    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
87        None
88    }
89
90    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
91        SLOT_NAMES[idx].to_string()
92    }
93
94    fn serialize(
95        _array: ArrayView<'_, Self>,
96        _session: &VortexSession,
97    ) -> VortexResult<Option<Vec<u8>>> {
98        vortex_error::vortex_bail!("Shared array is not serializable")
99    }
100
101    fn deserialize(
102        &self,
103        _dtype: &DType,
104        _len: usize,
105        _metadata: &[u8],
106
107        _buffers: &[BufferHandle],
108        _children: &dyn crate::serde::ArrayChildren,
109        _session: &VortexSession,
110    ) -> VortexResult<crate::array::ArrayParts<Self>> {
111        vortex_error::vortex_bail!("Shared array is not serializable")
112    }
113
114    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
115        array
116            .get_or_compute(|source| source.clone().execute::<Canonical>(ctx))
117            .map(ExecutionResult::done)
118    }
119}
120impl OperationsVTable<Shared> for Shared {
121    fn scalar_at(
122        array: ArrayView<'_, Shared>,
123        index: usize,
124        _ctx: &mut ExecutionCtx,
125    ) -> VortexResult<Scalar> {
126        array.current_array_ref().scalar_at(index)
127    }
128}
129
130impl ValidityVTable<Shared> for Shared {
131    fn validity(array: ArrayView<'_, Shared>) -> VortexResult<Validity> {
132        array.current_array_ref().validity()
133    }
134}