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