vortex_array/vtable/
dyn_.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::mem::transmute;
6use std::sync::Arc;
7
8use arcref::ArcRef;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_ensure;
12use vortex_session::VortexSession;
13
14use crate::ArrayAdapter;
15use crate::ArrayRef;
16use crate::DynArray;
17use crate::ExecutionResult;
18use crate::ExecutionStep;
19use crate::IntoArray;
20use crate::buffer::BufferHandle;
21use crate::dtype::DType;
22use crate::executor::ExecutionCtx;
23use crate::serde::ArrayChildren;
24use crate::vtable::VTable;
25
26/// ArrayId is a globally unique name for the array's vtable.
27pub type ArrayId = ArcRef<str>;
28
29/// Reference-counted DynVTable
30pub type DynVTableRef = Arc<dyn DynVTable>;
31
32/// Dynamically typed vtable trait.
33///
34/// This trait contains the implementation API for Vortex arrays, allowing us to keep the public
35/// [`DynArray`] trait API to a minimum.
36pub trait DynVTable: 'static + Send + Sync + Debug {
37    /// Clone this vtable into a `Box<dyn DynVTable>`.
38    fn clone_boxed(&self) -> Box<dyn DynVTable>;
39
40    #[allow(clippy::too_many_arguments)]
41    fn build(
42        &self,
43        id: ArrayId,
44        dtype: &DType,
45        len: usize,
46        metadata: &[u8],
47        buffers: &[BufferHandle],
48        children: &dyn ArrayChildren,
49        session: &VortexSession,
50    ) -> VortexResult<ArrayRef>;
51    fn with_children(&self, array: &ArrayRef, children: Vec<ArrayRef>) -> VortexResult<ArrayRef>;
52
53    /// See [`VTable::reduce`]
54    fn reduce(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
55
56    /// See [`VTable::reduce_parent`]
57    fn reduce_parent(
58        &self,
59        array: &ArrayRef,
60        parent: &ArrayRef,
61        child_idx: usize,
62    ) -> VortexResult<Option<ArrayRef>>;
63
64    /// See [`VTable::execute`]
65    fn execute(&self, array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult>;
66
67    /// See [`VTable::execute_parent`]
68    fn execute_parent(
69        &self,
70        array: &ArrayRef,
71        parent: &ArrayRef,
72        child_idx: usize,
73        ctx: &mut ExecutionCtx,
74    ) -> VortexResult<Option<ArrayRef>>;
75}
76
77impl<V: VTable> DynVTable for V {
78    fn clone_boxed(&self) -> Box<dyn DynVTable> {
79        Box::new(self.clone())
80    }
81
82    fn build(
83        &self,
84        _id: ArrayId,
85        dtype: &DType,
86        len: usize,
87        metadata: &[u8],
88        buffers: &[BufferHandle],
89        children: &dyn ArrayChildren,
90        session: &VortexSession,
91    ) -> VortexResult<ArrayRef> {
92        let metadata = V::deserialize(metadata, dtype, len, buffers, session)?;
93        let array = V::build(dtype, len, &metadata, buffers, children)?;
94        assert_eq!(array.len(), len, "Array length mismatch after building");
95        assert_eq!(array.dtype(), dtype, "Array dtype mismatch after building");
96        Ok(array.into_array())
97    }
98
99    fn with_children(&self, array: &ArrayRef, children: Vec<ArrayRef>) -> VortexResult<ArrayRef> {
100        let mut array = array.as_::<V>().clone();
101        V::with_children(&mut array, children)?;
102        Ok(array.into_array())
103    }
104
105    fn reduce(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
106        let Some(reduced) = V::reduce(downcast::<V>(array))? else {
107            return Ok(None);
108        };
109        vortex_ensure!(
110            reduced.len() == array.len(),
111            "Reduced array length mismatch from {} to {}",
112            array.encoding_id(),
113            reduced.encoding_id()
114        );
115        vortex_ensure!(
116            reduced.dtype() == array.dtype(),
117            "Reduced array dtype mismatch from {} to {}",
118            array.encoding_id(),
119            reduced.encoding_id()
120        );
121        Ok(Some(reduced))
122    }
123
124    fn reduce_parent(
125        &self,
126        array: &ArrayRef,
127        parent: &ArrayRef,
128        child_idx: usize,
129    ) -> VortexResult<Option<ArrayRef>> {
130        let Some(reduced) = V::reduce_parent(downcast::<V>(array), parent, child_idx)? else {
131            return Ok(None);
132        };
133
134        vortex_ensure!(
135            reduced.len() == parent.len(),
136            "Reduced array length mismatch from {} to {}",
137            parent.encoding_id(),
138            reduced.encoding_id()
139        );
140        vortex_ensure!(
141            reduced.dtype() == parent.dtype(),
142            "Reduced array dtype mismatch from {} to {}",
143            parent.encoding_id(),
144            reduced.encoding_id()
145        );
146
147        Ok(Some(reduced))
148    }
149
150    fn execute(&self, array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
151        // Capture metadata before the move for post-validation and stats inheritance.
152        let len = array.len();
153        let dtype = array.dtype().clone();
154        let stats = array.statistics().to_owned();
155
156        let owned = downcast_owned::<V>(array);
157        let result = V::execute(owned, ctx)?;
158
159        if matches!(result.step(), ExecutionStep::Done) {
160            if cfg!(debug_assertions) {
161                vortex_ensure!(
162                    result.array().len() == len,
163                    "Result length mismatch for {:?}",
164                    self
165                );
166                vortex_ensure!(
167                    result.array().dtype() == &dtype,
168                    "Executed canonical dtype mismatch for {:?}",
169                    self
170                );
171            }
172
173            // TODO(ngates): do we want to do this on every execution? We used to in to_canonical.
174            result.array().statistics().set_iter(stats.into_iter());
175        }
176
177        Ok(result)
178    }
179
180    fn execute_parent(
181        &self,
182        array: &ArrayRef,
183        parent: &ArrayRef,
184        child_idx: usize,
185        ctx: &mut ExecutionCtx,
186    ) -> VortexResult<Option<ArrayRef>> {
187        let Some(result) = V::execute_parent(downcast::<V>(array), parent, child_idx, ctx)? else {
188            return Ok(None);
189        };
190
191        if cfg!(debug_assertions) {
192            vortex_ensure!(
193                result.as_ref().len() == parent.len(),
194                "Executed parent canonical length mismatch"
195            );
196            vortex_ensure!(
197                result.as_ref().dtype() == parent.dtype(),
198                "Executed parent canonical dtype mismatch"
199            );
200        }
201
202        Ok(Some(result))
203    }
204}
205
206fn downcast<V: VTable>(array: &ArrayRef) -> &V::Array {
207    array
208        .as_any()
209        .downcast_ref::<ArrayAdapter<V>>()
210        .vortex_expect("Failed to downcast array to expected encoding type")
211        .as_inner()
212}
213
214/// Downcast an `ArrayRef` into an `Arc<V::Array>` without cloning.
215///
216/// This is a zero-cost pointer cast leveraging the `#[repr(transparent)]` layout of
217/// [`ArrayAdapter`].
218fn downcast_owned<V: VTable>(array: ArrayRef) -> Arc<V::Array> {
219    let adapter: Arc<ArrayAdapter<V>> = array
220        .as_any_arc()
221        .downcast::<ArrayAdapter<V>>()
222        .ok()
223        .vortex_expect("Failed to downcast array to expected encoding type");
224    // SAFETY: ArrayAdapter<V> is #[repr(transparent)] over V::Array,
225    // so Arc<ArrayAdapter<V>> and Arc<V::Array> have identical layout.
226    unsafe { transmute::<Arc<ArrayAdapter<V>>, Arc<V::Array>>(adapter) }
227}
228
229/// Upcast an `Arc<V::Array>` into an `ArrayRef` without cloning.
230///
231/// This is a zero-cost pointer cast leveraging the `#[repr(transparent)]` layout of
232/// [`ArrayAdapter`]. It is the reverse of `downcast_owned`.
233pub(crate) fn upcast_array<V: VTable>(array: Arc<V::Array>) -> ArrayRef {
234    // SAFETY: ArrayAdapter<V> is #[repr(transparent)] over V::Array,
235    // so Arc<V::Array> and Arc<ArrayAdapter<V>> have identical layout.
236    unsafe { transmute::<Arc<V::Array>, Arc<ArrayAdapter<V>>>(array) }
237}
vortex_array/vtable/dyn_.rs

vortex_array/vtable/
dyn_.rs
