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