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::sync::Arc;
6
7use arcref::ArcRef;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_ensure;
11use vortex_session::VortexSession;
12
13use crate::ArrayRef;
14use crate::DynArray;
15use crate::ExecutionResult;
16use crate::ExecutionStep;
17use crate::IntoArray;
18use crate::buffer::BufferHandle;
19use crate::dtype::DType;
20use crate::executor::ExecutionCtx;
21use crate::serde::ArrayChildren;
22use crate::stats::ArrayStats;
23use crate::vtable::Array;
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 inner = V::build(dtype, len, &metadata, buffers, children)?;
94        // Validate the inner array's properties before wrapping.
95        assert_eq!(V::len(&inner), len, "Array length mismatch after building");
96        assert_eq!(
97            V::dtype(&inner),
98            dtype,
99            "Array dtype mismatch after building"
100        );
101        // Wrap in Array<V> for safe downcasting.
102        // SAFETY: We just validated that V::len(&inner) == len and V::dtype(&inner) == dtype.
103        let array = unsafe {
104            Array::new_unchecked(
105                self.clone(),
106                dtype.clone(),
107                len,
108                inner,
109                ArrayStats::default(),
110            )
111        };
112        Ok(array.into_array())
113    }
114
115    fn with_children(&self, array: &ArrayRef, children: Vec<ArrayRef>) -> VortexResult<ArrayRef> {
116        let mut array = array.as_::<V>().clone();
117        V::with_children(&mut array, children)?;
118        Ok(array.into_array())
119    }
120
121    fn reduce(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
122        let Some(reduced) = V::reduce(downcast::<V>(array))? else {
123            return Ok(None);
124        };
125        vortex_ensure!(
126            reduced.len() == array.len(),
127            "Reduced array length mismatch from {} to {}",
128            array.encoding_id(),
129            reduced.encoding_id()
130        );
131        vortex_ensure!(
132            reduced.dtype() == array.dtype(),
133            "Reduced array dtype mismatch from {} to {}",
134            array.encoding_id(),
135            reduced.encoding_id()
136        );
137        Ok(Some(reduced))
138    }
139
140    fn reduce_parent(
141        &self,
142        array: &ArrayRef,
143        parent: &ArrayRef,
144        child_idx: usize,
145    ) -> VortexResult<Option<ArrayRef>> {
146        let Some(reduced) = V::reduce_parent(downcast::<V>(array), parent, child_idx)? else {
147            return Ok(None);
148        };
149
150        vortex_ensure!(
151            reduced.len() == parent.len(),
152            "Reduced array length mismatch from {} to {}",
153            parent.encoding_id(),
154            reduced.encoding_id()
155        );
156        vortex_ensure!(
157            reduced.dtype() == parent.dtype(),
158            "Reduced array dtype mismatch from {} to {}",
159            parent.encoding_id(),
160            reduced.encoding_id()
161        );
162
163        Ok(Some(reduced))
164    }
165
166    fn execute(&self, array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
167        // Capture metadata before the move for post-validation and stats inheritance.
168        let len = array.len();
169        let dtype = array.dtype().clone();
170        let stats = array.statistics().to_owned();
171
172        let owned = downcast_owned::<V>(array);
173        let result = V::execute(owned, ctx)?;
174
175        if matches!(result.step(), ExecutionStep::Done) {
176            if cfg!(debug_assertions) {
177                vortex_ensure!(
178                    result.array().len() == len,
179                    "Result length mismatch for {:?}",
180                    self
181                );
182                vortex_ensure!(
183                    result.array().dtype() == &dtype,
184                    "Executed canonical dtype mismatch for {:?}",
185                    self
186                );
187            }
188
189            // TODO(ngates): do we want to do this on every execution? We used to in to_canonical.
190            result.array().statistics().set_iter(stats.into_iter());
191        }
192
193        Ok(result)
194    }
195
196    fn execute_parent(
197        &self,
198        array: &ArrayRef,
199        parent: &ArrayRef,
200        child_idx: usize,
201        ctx: &mut ExecutionCtx,
202    ) -> VortexResult<Option<ArrayRef>> {
203        let Some(result) = V::execute_parent(downcast::<V>(array), parent, child_idx, ctx)? else {
204            return Ok(None);
205        };
206
207        if cfg!(debug_assertions) {
208            vortex_ensure!(
209                result.as_ref().len() == parent.len(),
210                "Executed parent canonical length mismatch"
211            );
212            vortex_ensure!(
213                result.as_ref().dtype() == parent.dtype(),
214                "Executed parent canonical dtype mismatch"
215            );
216        }
217
218        Ok(Some(result))
219    }
220}
221
222/// Borrow-downcast an `ArrayRef` to `&Array<V>`.
223fn downcast<V: VTable>(array: &ArrayRef) -> &Array<V> {
224    array
225        .as_any()
226        .downcast_ref::<Array<V>>()
227        .vortex_expect("Failed to downcast array to expected encoding type")
228}
229
230/// Downcast an `ArrayRef` into an `Arc<Array<V>>`.
231fn downcast_owned<V: VTable>(array: ArrayRef) -> Arc<Array<V>> {
232    let any_arc = array.as_any_arc();
233    any_arc
234        .downcast::<Array<V>>()
235        .ok()
236        .vortex_expect("Failed to downcast array to expected encoding type")
237}
vortex_array/vtable/dyn_.rs

vortex_array/vtable/
dyn_.rs
