vortex_layout/
vtable.rs

1use std::collections::BTreeSet;
2use std::fmt::Debug;
3use std::ops::Deref;
4use std::sync::Arc;
5
6use vortex_array::{ArrayContext, DeserializeMetadata, SerializeMetadata};
7use vortex_dtype::{DType, FieldMask};
8use vortex_error::VortexResult;
9
10use crate::children::LayoutChildren;
11use crate::segments::{SegmentId, SegmentSource};
12use crate::{
13    IntoLayout, Layout, LayoutChildType, LayoutEncoding, LayoutEncodingRef, LayoutId,
14    LayoutReaderRef, LayoutRef,
15};
16
17pub trait VTable: 'static + Sized + Send + Sync + Debug {
18    type Layout: 'static + Send + Sync + Clone + Debug + Deref<Target = dyn Layout> + IntoLayout;
19    type Encoding: 'static + Send + Sync + Deref<Target = dyn LayoutEncoding>;
20    type Metadata: SerializeMetadata + DeserializeMetadata + Debug;
21
22    /// Returns the ID of the layout encoding.
23    fn id(encoding: &Self::Encoding) -> LayoutId;
24
25    /// Returns the encoding for the layout.
26    fn encoding(layout: &Self::Layout) -> LayoutEncodingRef;
27
28    /// Returns the row count for the layout reader.
29    fn row_count(layout: &Self::Layout) -> u64;
30
31    /// Returns the dtype for the layout reader.
32    fn dtype(layout: &Self::Layout) -> &DType;
33
34    /// Returns the metadata for the layout.
35    fn metadata(layout: &Self::Layout) -> Self::Metadata;
36
37    /// Returns the segment IDs for the layout.
38    fn segment_ids(layout: &Self::Layout) -> Vec<SegmentId>;
39
40    /// Returns the number of children for the layout.
41    fn nchildren(layout: &Self::Layout) -> usize;
42
43    /// Return the child at the given index.
44    fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef>;
45
46    /// Return the type of the child at the given index.
47    fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType;
48
49    /// Register row splits for the layout.
50    // TODO(ngates): this should be implemented with a visitor, but we need to fix the FieldMask
51    //  API first.
52    fn register_splits(
53        layout: &Self::Layout,
54        field_mask: &[FieldMask],
55        row_offset: u64,
56        splits: &mut BTreeSet<u64>,
57    ) -> VortexResult<()>;
58
59    /// Create a new reader for the layout.
60    fn new_reader(
61        layout: &Self::Layout,
62        name: &Arc<str>,
63        segment_source: &Arc<dyn SegmentSource>,
64        ctx: &ArrayContext,
65    ) -> VortexResult<LayoutReaderRef>;
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    ) -> VortexResult<Self::Layout>;
76}
77
78#[macro_export]
79macro_rules! vtable {
80    ($V:ident) => {
81        $crate::aliases::paste::paste! {
82            #[derive(Debug)]
83            pub struct [<$V VTable>];
84
85            impl AsRef<dyn $crate::Layout> for [<$V Layout>] {
86                fn as_ref(&self) -> &dyn $crate::Layout {
87                    // We can unsafe cast ourselves to a LayoutAdapter.
88                    unsafe { &*(self as *const [<$V Layout>] as *const $crate::LayoutAdapter<[<$V VTable>]>) }
89                }
90            }
91
92            impl std::ops::Deref for [<$V Layout>] {
93                type Target = dyn $crate::Layout;
94
95                fn deref(&self) -> &Self::Target {
96                    // We can unsafe cast ourselves to an LayoutAdapter.
97                    unsafe { &*(self as *const [<$V Layout>] as *const $crate::LayoutAdapter<[<$V VTable>]>) }
98                }
99            }
100
101            impl $crate::IntoLayout for [<$V Layout>] {
102                fn into_layout(self) -> $crate::LayoutRef {
103                    // We can unsafe transmute ourselves to an LayoutAdapter.
104                    std::sync::Arc::new(unsafe { std::mem::transmute::<[<$V Layout>], $crate::LayoutAdapter::<[<$V VTable>]>>(self) })
105                }
106            }
107
108            impl AsRef<dyn $crate::LayoutEncoding> for [<$V LayoutEncoding>] {
109                fn as_ref(&self) -> &dyn $crate::LayoutEncoding {
110                    // We can unsafe cast ourselves to an LayoutEncodingAdapter.
111                    unsafe { &*(self as *const [<$V LayoutEncoding>] as *const $crate::LayoutEncodingAdapter<[<$V VTable>]>) }
112                }
113            }
114
115            impl std::ops::Deref for [<$V LayoutEncoding>] {
116                type Target = dyn $crate::LayoutEncoding;
117
118                fn deref(&self) -> &Self::Target {
119                    // We can unsafe cast ourselves to an LayoutEncodingAdapter.
120                    unsafe { &*(self as *const [<$V LayoutEncoding>] as *const $crate::LayoutEncodingAdapter<[<$V VTable>]>) }
121                }
122            }
123        }
124    };
125}