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