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_session::VortexSession;
9
10use crate::ArrayBufferVisitor;
11use crate::ArrayChildVisitor;
12use crate::ArrayRef;
13use crate::Canonical;
14use crate::EmptyMetadata;
15use crate::ExecutionCtx;
16use crate::Precision;
17use crate::arrays::shared::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::BaseArrayVTable;
28use crate::vtable::OperationsVTable;
29use crate::vtable::VTable;
30use crate::vtable::ValidityVTable;
31use crate::vtable::VisitorVTable;
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(Debug)]
38pub struct SharedVTable;
39
40impl SharedVTable {
41    pub const ID: ArrayId = ArrayId::new_ref("vortex.shared");
42}
43
44impl VTable for SharedVTable {
45    type Array = SharedArray;
46    type Metadata = EmptyMetadata;
47
48    type ArrayVTable = Self;
49    type OperationsVTable = Self;
50    type ValidityVTable = Self;
51    type VisitorVTable = Self;
52
53    fn id(_array: &Self::Array) -> ArrayId {
54        Self::ID
55    }
56
57    fn metadata(_array: &Self::Array) -> VortexResult<Self::Metadata> {
58        Ok(EmptyMetadata)
59    }
60
61    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
62        vortex_error::vortex_bail!("Shared array is not serializable")
63    }
64
65    fn deserialize(
66        _bytes: &[u8],
67        _dtype: &DType,
68        _len: usize,
69        _buffers: &[BufferHandle],
70        _session: &VortexSession,
71    ) -> VortexResult<Self::Metadata> {
72        vortex_error::vortex_bail!("Shared array is not serializable")
73    }
74
75    fn build(
76        dtype: &DType,
77        len: usize,
78        _metadata: &Self::Metadata,
79        _buffers: &[BufferHandle],
80        children: &dyn crate::serde::ArrayChildren,
81    ) -> VortexResult<SharedArray> {
82        let child = children.get(0, dtype, len)?;
83        Ok(SharedArray::new(child))
84    }
85
86    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
87        vortex_error::vortex_ensure!(
88            children.len() == 1,
89            "SharedArray expects exactly 1 child, got {}",
90            children.len()
91        );
92        let child = children
93            .into_iter()
94            .next()
95            .vortex_expect("children length already validated");
96        array.set_source(child);
97        Ok(())
98    }
99
100    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
101        array.get_or_compute(|source| source.clone().execute::<Canonical>(ctx))
102    }
103}
104
105impl BaseArrayVTable<SharedVTable> for SharedVTable {
106    fn len(array: &SharedArray) -> usize {
107        array.current_array_ref().len()
108    }
109
110    fn dtype(array: &SharedArray) -> &DType {
111        &array.dtype
112    }
113
114    fn stats(array: &SharedArray) -> StatsSetRef<'_> {
115        array.stats.to_ref(array.as_ref())
116    }
117
118    fn array_hash<H: std::hash::Hasher>(array: &SharedArray, state: &mut H, precision: Precision) {
119        let current = array.current_array_ref();
120        current.array_hash(state, precision);
121        array.dtype.hash(state);
122    }
123
124    fn array_eq(array: &SharedArray, other: &SharedArray, precision: Precision) -> bool {
125        let current = array.current_array_ref();
126        let other_current = other.current_array_ref();
127        current.array_eq(other_current, precision) && array.dtype == other.dtype
128    }
129}
130
131impl OperationsVTable<SharedVTable> for SharedVTable {
132    fn scalar_at(array: &SharedArray, index: usize) -> VortexResult<Scalar> {
133        array.current_array_ref().scalar_at(index)
134    }
135}
136
137impl ValidityVTable<SharedVTable> for SharedVTable {
138    fn validity(array: &SharedArray) -> VortexResult<Validity> {
139        array.current_array_ref().validity()
140    }
141}
142
143impl VisitorVTable<SharedVTable> for SharedVTable {
144    fn visit_buffers(_array: &SharedArray, _visitor: &mut dyn ArrayBufferVisitor) {}
145
146    fn visit_children(array: &SharedArray, visitor: &mut dyn ArrayChildVisitor) {
147        visitor.visit_child("source", array.current_array_ref());
148    }
149
150    fn nchildren(_array: &SharedArray) -> usize {
151        1
152    }
153
154    fn nth_child(array: &SharedArray, idx: usize) -> Option<ArrayRef> {
155        match idx {
156            0 => Some(array.current_array_ref().clone()),
157            _ => None,
158        }
159    }
160}