vortex_array/vtable/
dyn_.rs

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

vortex_array/vtable/
dyn_.rs
