Skip to main content

vortex_layout/
encoding.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::fmt::Debug;
6use std::fmt::Display;
7use std::fmt::Formatter;
8
9use arcref::ArcRef;
10use vortex_array::DeserializeMetadata;
11use vortex_array::dtype::DType;
12use vortex_error::VortexExpect;
13use vortex_error::VortexResult;
14use vortex_error::vortex_panic;
15use vortex_session::registry::Id;
16use vortex_session::registry::ReadContext;
17
18use crate::IntoLayout;
19use crate::LayoutChildren;
20use crate::LayoutRef;
21use crate::VTable;
22use crate::segments::SegmentId;
23
24/// A unique identifier for a layout encoding.
25pub type LayoutEncodingId = Id;
26pub type LayoutEncodingRef = ArcRef<dyn LayoutEncoding>;
27
28pub trait LayoutEncoding: 'static + Send + Sync + Debug + private::Sealed {
29    fn as_any(&self) -> &dyn Any;
30
31    fn id(&self) -> LayoutEncodingId;
32
33    fn build(
34        &self,
35        dtype: &DType,
36        row_count: u64,
37        metadata: &[u8],
38        segment_ids: Vec<SegmentId>,
39        children: &dyn LayoutChildren,
40        ctx: &ReadContext,
41    ) -> VortexResult<LayoutRef>;
42}
43
44#[repr(transparent)]
45pub struct LayoutEncodingAdapter<V: VTable>(V::Encoding);
46
47impl<V: VTable> LayoutEncoding for LayoutEncodingAdapter<V> {
48    fn as_any(&self) -> &dyn Any {
49        self
50    }
51
52    fn id(&self) -> LayoutEncodingId {
53        V::id(&self.0)
54    }
55
56    fn build(
57        &self,
58        dtype: &DType,
59        row_count: u64,
60        metadata: &[u8],
61        segment_ids: Vec<SegmentId>,
62        children: &dyn LayoutChildren,
63        ctx: &ReadContext,
64    ) -> VortexResult<LayoutRef> {
65        let metadata = <V::Metadata as DeserializeMetadata>::deserialize(metadata)?;
66        let layout = V::build(
67            &self.0,
68            dtype,
69            row_count,
70            &metadata,
71            segment_ids,
72            children,
73            ctx,
74        )?;
75
76        // Validate that the builder function returned the expected values.
77        if layout.row_count() != row_count {
78            vortex_panic!(
79                "Layout row count mismatch: {} != {}",
80                layout.row_count(),
81                row_count
82            );
83        }
84        if layout.dtype() != dtype {
85            vortex_panic!("Layout dtype mismatch: {} != {}", layout.dtype(), dtype);
86        }
87
88        Ok(layout.into_layout())
89    }
90}
91
92impl<V: VTable> Debug for LayoutEncodingAdapter<V> {
93    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
94        f.debug_struct("LayoutEncoding")
95            .field("id", &self.id())
96            .finish()
97    }
98}
99
100impl Display for dyn LayoutEncoding + '_ {
101    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
102        write!(f, "{}", self.id())
103    }
104}
105
106impl PartialEq for dyn LayoutEncoding + '_ {
107    fn eq(&self, other: &Self) -> bool {
108        self.id() == other.id()
109    }
110}
111
112impl Eq for dyn LayoutEncoding + '_ {}
113
114impl dyn LayoutEncoding + '_ {
115    pub fn is<V: VTable>(&self) -> bool {
116        self.as_opt::<V>().is_some()
117    }
118
119    pub fn as_<V: VTable>(&self) -> &V::Encoding {
120        self.as_opt::<V>()
121            .vortex_expect("LayoutEncoding is not of the expected type")
122    }
123
124    pub fn as_opt<V: VTable>(&self) -> Option<&V::Encoding> {
125        self.as_any()
126            .downcast_ref::<LayoutEncodingAdapter<V>>()
127            .map(|e| &e.0)
128    }
129}
130
131mod private {
132    use super::*;
133    use crate::layouts::foreign::ForeignLayoutEncoding;
134
135    pub trait Sealed {}
136
137    impl<V: VTable> Sealed for LayoutEncodingAdapter<V> {}
138    impl Sealed for ForeignLayoutEncoding {}
139}