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::ArrayId;
28use crate::vtable::OperationsVTable;
29use crate::vtable::VTable;
30use crate::vtable::ValidityVTable;
31
32vtable!(Shared);
33
34// TODO(ngates): consider hooking Shared into the iterative execution model. Cache either the
35//  most executed, or after each iteration, and return a shared cache for each execution.
36#[derive(Clone, Debug)]
37pub struct Shared;
38
39impl Shared {
40    pub const ID: ArrayId = ArrayId::new_ref("vortex.shared");
41}
42
43impl VTable for Shared {
44    type Array = SharedArray;
45    type Metadata = EmptyMetadata;
46    type OperationsVTable = Self;
47    type ValidityVTable = Self;
48    fn vtable(_array: &Self::Array) -> &Self {
49        &Shared
50    }
51
52    fn id(&self) -> ArrayId {
53        Self::ID
54    }
55
56    fn len(array: &SharedArray) -> usize {
57        array.current_array_ref().len()
58    }
59
60    fn dtype(array: &SharedArray) -> &DType {
61        &array.dtype
62    }
63
64    fn stats(array: &SharedArray) -> StatsSetRef<'_> {
65        array.stats.to_ref(array.as_ref())
66    }
67
68    fn array_hash<H: std::hash::Hasher>(array: &SharedArray, state: &mut H, precision: Precision) {
69        let current = array.current_array_ref();
70        current.array_hash(state, precision);
71        array.dtype.hash(state);
72    }
73
74    fn array_eq(array: &SharedArray, other: &SharedArray, precision: Precision) -> bool {
75        let current = array.current_array_ref();
76        let other_current = other.current_array_ref();
77        current.array_eq(other_current, precision) && array.dtype == other.dtype
78    }
79
80    fn nbuffers(_array: &Self::Array) -> usize {
81        0
82    }
83
84    fn buffer(_array: &Self::Array, _idx: usize) -> BufferHandle {
85        vortex_panic!("SharedArray has no buffers")
86    }
87
88    fn buffer_name(_array: &Self::Array, _idx: usize) -> Option<String> {
89        None
90    }
91
92    fn nchildren(_array: &Self::Array) -> usize {
93        1
94    }
95
96    fn child(array: &Self::Array, idx: usize) -> ArrayRef {
97        match idx {
98            0 => array.current_array_ref().clone(),
99            _ => vortex_panic!("SharedArray child index {idx} out of bounds"),
100        }
101    }
102
103    fn child_name(_array: &Self::Array, idx: usize) -> String {
104        match idx {
105            0 => "source".to_string(),
106            _ => vortex_panic!("SharedArray child_name index {idx} out of bounds"),
107        }
108    }
109
110    fn metadata(_array: &Self::Array) -> VortexResult<Self::Metadata> {
111        Ok(EmptyMetadata)
112    }
113
114    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
115        vortex_error::vortex_bail!("Shared array is not serializable")
116    }
117
118    fn deserialize(
119        _bytes: &[u8],
120        _dtype: &DType,
121        _len: usize,
122        _buffers: &[BufferHandle],
123        _session: &VortexSession,
124    ) -> VortexResult<Self::Metadata> {
125        vortex_error::vortex_bail!("Shared array is not serializable")
126    }
127
128    fn build(
129        dtype: &DType,
130        len: usize,
131        _metadata: &Self::Metadata,
132        _buffers: &[BufferHandle],
133        children: &dyn crate::serde::ArrayChildren,
134    ) -> VortexResult<SharedArray> {
135        let child = children.get(0, dtype, len)?;
136        Ok(SharedArray::new(child))
137    }
138
139    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
140        vortex_error::vortex_ensure!(
141            children.len() == 1,
142            "SharedArray expects exactly 1 child, got {}",
143            children.len()
144        );
145        let child = children
146            .into_iter()
147            .next()
148            .vortex_expect("children length already validated");
149        array.set_source(child);
150        Ok(())
151    }
152
153    fn execute(array: Arc<Self::Array>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
154        array
155            .get_or_compute(|source| source.clone().execute::<Canonical>(ctx))
156            .map(ExecutionResult::done)
157    }
158}
159impl OperationsVTable<Shared> for Shared {
160    fn scalar_at(array: &SharedArray, index: usize) -> VortexResult<Scalar> {
161        array.current_array_ref().scalar_at(index)
162    }
163}
164
165impl ValidityVTable<Shared> for Shared {
166    fn validity(array: &SharedArray) -> VortexResult<Validity> {
167        array.current_array_ref().validity()
168    }
169}