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::Hash;
5use std::sync::Arc;
6
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_panic;
10use vortex_session::VortexSession;
11
12use crate::ArrayRef;
13use crate::Canonical;
14use crate::EmptyMetadata;
15use crate::ExecutionCtx;
16use crate::ExecutionResult;
17use crate::Precision;
18use crate::arrays::SharedArray;
19use crate::buffer::BufferHandle;
20use crate::dtype::DType;
21use crate::hash::ArrayEq;
22use crate::hash::ArrayHash;
23use crate::scalar::Scalar;
24use crate::stats::StatsSetRef;
25use crate::validity::Validity;
26use crate::vtable;
27use crate::vtable::Array;
28use crate::vtable::ArrayId;
29use crate::vtable::OperationsVTable;
30use crate::vtable::VTable;
31use crate::vtable::ValidityVTable;
32
33vtable!(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 VTable for Shared {
45    type Array = SharedArray;
46    type Metadata = EmptyMetadata;
47    type OperationsVTable = Self;
48    type ValidityVTable = Self;
49    fn vtable(_array: &Self::Array) -> &Self {
50        &Shared
51    }
52
53    fn id(&self) -> ArrayId {
54        Self::ID
55    }
56
57    fn len(array: &SharedArray) -> usize {
58        array.current_array_ref().len()
59    }
60
61    fn dtype(array: &SharedArray) -> &DType {
62        &array.dtype
63    }
64
65    fn stats(array: &SharedArray) -> StatsSetRef<'_> {
66        array.stats.to_ref(array.as_ref())
67    }
68
69    fn array_hash<H: std::hash::Hasher>(array: &SharedArray, state: &mut H, precision: Precision) {
70        let current = array.current_array_ref();
71        current.array_hash(state, precision);
72        array.dtype.hash(state);
73    }
74
75    fn array_eq(array: &SharedArray, other: &SharedArray, precision: Precision) -> bool {
76        let current = array.current_array_ref();
77        let other_current = other.current_array_ref();
78        current.array_eq(other_current, precision) && array.dtype == other.dtype
79    }
80
81    fn nbuffers(_array: &Self::Array) -> usize {
82        0
83    }
84
85    fn buffer(_array: &Self::Array, _idx: usize) -> BufferHandle {
86        vortex_panic!("SharedArray has no buffers")
87    }
88
89    fn buffer_name(_array: &Self::Array, _idx: usize) -> Option<String> {
90        None
91    }
92
93    fn nchildren(_array: &Self::Array) -> usize {
94        1
95    }
96
97    fn child(array: &Self::Array, idx: usize) -> ArrayRef {
98        match idx {
99            0 => array.current_array_ref().clone(),
100            _ => vortex_panic!("SharedArray child index {idx} out of bounds"),
101        }
102    }
103
104    fn child_name(_array: &Self::Array, idx: usize) -> String {
105        match idx {
106            0 => "source".to_string(),
107            _ => vortex_panic!("SharedArray child_name index {idx} out of bounds"),
108        }
109    }
110
111    fn metadata(_array: &Self::Array) -> VortexResult<Self::Metadata> {
112        Ok(EmptyMetadata)
113    }
114
115    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
116        vortex_error::vortex_bail!("Shared array is not serializable")
117    }
118
119    fn deserialize(
120        _bytes: &[u8],
121        _dtype: &DType,
122        _len: usize,
123        _buffers: &[BufferHandle],
124        _session: &VortexSession,
125    ) -> VortexResult<Self::Metadata> {
126        vortex_error::vortex_bail!("Shared array is not serializable")
127    }
128
129    fn build(
130        dtype: &DType,
131        len: usize,
132        _metadata: &Self::Metadata,
133        _buffers: &[BufferHandle],
134        children: &dyn crate::serde::ArrayChildren,
135    ) -> VortexResult<SharedArray> {
136        let child = children.get(0, dtype, len)?;
137        Ok(SharedArray::new(child))
138    }
139
140    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
141        vortex_error::vortex_ensure!(
142            children.len() == 1,
143            "SharedArray expects exactly 1 child, got {}",
144            children.len()
145        );
146        let child = children
147            .into_iter()
148            .next()
149            .vortex_expect("children length already validated");
150        array.set_source(child);
151        Ok(())
152    }
153
154    fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
155        array
156            .get_or_compute(|source| source.clone().execute::<Canonical>(ctx))
157            .map(ExecutionResult::done)
158    }
159}
160impl OperationsVTable<Shared> for Shared {
161    fn scalar_at(
162        array: &SharedArray,
163        index: usize,
164        _ctx: &mut ExecutionCtx,
165    ) -> VortexResult<Scalar> {
166        array.current_array_ref().scalar_at(index)
167    }
168}
169
170impl ValidityVTable<Shared> for Shared {
171    fn validity(array: &SharedArray) -> VortexResult<Validity> {
172        array.current_array_ref().validity()
173    }
174}