vortex_layout/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::ops::Deref;
6use std::sync::Arc;
7
8use vortex_array::{ArrayContext, DeserializeMetadata, SerializeMetadata};
9use vortex_dtype::DType;
10use vortex_error::VortexResult;
11
12use crate::children::LayoutChildren;
13use crate::segments::{SegmentId, SegmentSource};
14use crate::{
15    IntoLayout, Layout, LayoutChildType, LayoutEncoding, LayoutEncodingRef, LayoutId,
16    LayoutReaderRef, LayoutRef,
17};
18
19pub trait VTable: 'static + Sized + Send + Sync + Debug {
20    type Layout: 'static + Send + Sync + Clone + Debug + Deref<Target = dyn Layout> + IntoLayout;
21    type Encoding: 'static + Send + Sync + Deref<Target = dyn LayoutEncoding>;
22    type Metadata: SerializeMetadata + DeserializeMetadata + Debug;
23
24    /// Returns the ID of the layout encoding.
25    fn id(encoding: &Self::Encoding) -> LayoutId;
26
27    /// Returns the encoding for the layout.
28    fn encoding(layout: &Self::Layout) -> LayoutEncodingRef;
29
30    /// Returns the row count for the layout reader.
31    fn row_count(layout: &Self::Layout) -> u64;
32
33    /// Returns the dtype for the layout reader.
34    fn dtype(layout: &Self::Layout) -> &DType;
35
36    /// Returns the metadata for the layout.
37    fn metadata(layout: &Self::Layout) -> Self::Metadata;
38
39    /// Returns the segment IDs for the layout.
40    fn segment_ids(layout: &Self::Layout) -> Vec<SegmentId>;
41
42    /// Returns the number of children for the layout.
43    fn nchildren(layout: &Self::Layout) -> usize;
44
45    /// Return the child at the given index.
46    fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef>;
47
48    /// Return the type of the child at the given index.
49    fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType;
50
51    /// Create a new reader for the layout.
52    fn new_reader(
53        layout: &Self::Layout,
54        name: Arc<str>,
55        segment_source: Arc<dyn SegmentSource>,
56    ) -> VortexResult<LayoutReaderRef>;
57
58    #[cfg(gpu_unstable)]
59    /// Create a new reader for the layout that uses a gpu device
60    fn new_gpu_reader(
61        layout: &Self::Layout,
62        name: Arc<str>,
63        segment_source: Arc<dyn SegmentSource>,
64        ctx: Arc<cudarc::driver::CudaContext>,
65    ) -> VortexResult<crate::gpu::GpuLayoutReaderRef>;
66
67    /// Construct a new [`Layout`] from the provided parts.
68    fn build(
69        encoding: &Self::Encoding,
70        dtype: &DType,
71        row_count: u64,
72        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
73        segment_ids: Vec<SegmentId>,
74        children: &dyn LayoutChildren,
75        ctx: ArrayContext,
76    ) -> VortexResult<Self::Layout>;
77}
78
79#[macro_export]
80macro_rules! vtable {
81    ($V:ident) => {
82        $crate::aliases::paste::paste! {
83            #[derive(Debug)]
84            pub struct [<$V VTable>];
85
86            impl AsRef<dyn $crate::Layout> for [<$V Layout>] {
87                fn as_ref(&self) -> &dyn $crate::Layout {
88                    // SAFETY: LayoutAdapter is #[repr(transparent)] over the Layout type,
89                    // which guarantees identical memory layout. This cast is safe because
90                    // we're only changing the type metadata, not the actual data.
91                    unsafe { &*(self as *const [<$V Layout>] as *const $crate::LayoutAdapter<[<$V VTable>]>) }
92                }
93            }
94
95            impl std::ops::Deref for [<$V Layout>] {
96                type Target = dyn $crate::Layout;
97
98                fn deref(&self) -> &Self::Target {
99                    // SAFETY: LayoutAdapter is #[repr(transparent)] over the Layout type,
100                    // which guarantees identical memory layout. This cast is safe because
101                    // we're only changing the type metadata, not the actual data.
102                    unsafe { &*(self as *const [<$V Layout>] as *const $crate::LayoutAdapter<[<$V VTable>]>) }
103                }
104            }
105
106            impl $crate::IntoLayout for [<$V Layout>] {
107                fn into_layout(self) -> $crate::LayoutRef {
108                    // SAFETY: LayoutAdapter is #[repr(transparent)] over the Layout type,
109                    // guaranteeing identical memory layout and alignment. The transmute is safe
110                    // because both types have the same size and representation.
111                    std::sync::Arc::new(unsafe { std::mem::transmute::<[<$V Layout>], $crate::LayoutAdapter::<[<$V VTable>]>>(self) })
112                }
113            }
114
115            impl From<[<$V Layout>]> for $crate::LayoutRef {
116                fn from(value: [<$V Layout>]) -> $crate::LayoutRef {
117                    use $crate::IntoLayout;
118                    value.into_layout()
119                }
120            }
121
122            impl AsRef<dyn $crate::LayoutEncoding> for [<$V LayoutEncoding>] {
123                fn as_ref(&self) -> &dyn $crate::LayoutEncoding {
124                    // SAFETY: LayoutEncodingAdapter is #[repr(transparent)] over the LayoutEncoding type,
125                    // which guarantees identical memory layout. This cast is safe because
126                    // we're only changing the type metadata, not the actual data.
127                    unsafe { &*(self as *const [<$V LayoutEncoding>] as *const $crate::LayoutEncodingAdapter<[<$V VTable>]>) }
128                }
129            }
130
131            impl std::ops::Deref for [<$V LayoutEncoding>] {
132                type Target = dyn $crate::LayoutEncoding;
133
134                fn deref(&self) -> &Self::Target {
135                    // SAFETY: LayoutEncodingAdapter is #[repr(transparent)] over the LayoutEncoding type,
136                    // which guarantees identical memory layout. This cast is safe because
137                    // we're only changing the type metadata, not the actual data.
138                    unsafe { &*(self as *const [<$V LayoutEncoding>] as *const $crate::LayoutEncodingAdapter<[<$V VTable>]>) }
139                }
140            }
141        }
142    };
143}