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