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