vortex_layout/layouts/dict/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod reader;
5pub mod writer;
6
7use std::sync::Arc;
8
9use reader::DictReader;
10use vortex_array::ProstMetadata;
11use vortex_array::dtype::DType;
12use vortex_array::dtype::Nullability;
13use vortex_array::dtype::PType;
14use vortex_error::VortexExpect;
15use vortex_error::VortexResult;
16use vortex_error::vortex_bail;
17use vortex_error::vortex_ensure;
18use vortex_error::vortex_panic;
19use vortex_session::VortexSession;
20use vortex_session::registry::CachedId;
21
22use crate::Layout;
23use crate::LayoutChildType;
24use crate::LayoutDeserializeArgs;
25use crate::LayoutId;
26use crate::LayoutParts;
27use crate::LayoutReaderContext;
28use crate::LayoutReaderRef;
29use crate::LayoutRef;
30use crate::VTable;
31use crate::children::OwnedLayoutChildren;
32use crate::segments::SegmentSource;
33
34/// Dictionary layout vtable.
35#[derive(Clone, Debug)]
36pub struct Dict;
37
38/// Backwards-compatible name for the dictionary layout plugin.
39pub use Dict as DictLayoutEncoding;
40
41/// Dictionary-layout-specific data.
42#[derive(Clone, Debug)]
43pub struct DictData {
44    codes_dtype: DType,
45    all_values_referenced: bool,
46}
47
48/// A dictionary values child paired with a codes child.
49pub type DictLayout = Layout<Dict>;
50
51impl VTable for Dict {
52    type LayoutData = DictData;
53    type Metadata = ProstMetadata<DictLayoutMetadata>;
54
55    fn id(&self) -> LayoutId {
56        static ID: CachedId = CachedId::new("vortex.dict");
57        *ID
58    }
59
60    fn metadata(layout: &Layout<Self>) -> Self::Metadata {
61        let mut metadata = DictLayoutMetadata::new(
62            PType::try_from(&layout.codes_dtype).vortex_expect("codes ptype"),
63        );
64        metadata.is_nullable_codes = Some(layout.codes_dtype.is_nullable());
65        metadata.all_values_referenced = Some(layout.all_values_referenced);
66        ProstMetadata(metadata)
67    }
68
69    fn deserialize(
70        &self,
71        args: &LayoutDeserializeArgs<'_>,
72        metadata: &DictLayoutMetadata,
73    ) -> VortexResult<Self::LayoutData> {
74        vortex_ensure!(
75            args.children.nchildren() == 2,
76            "DictLayout expects exactly 2 children"
77        );
78        let codes_nullable = metadata
79            .is_nullable_codes
80            .map(Nullability::from)
81            .unwrap_or_else(|| args.dtype.nullability());
82        let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable);
83        args.children.child(0, args.dtype)?;
84        let codes = args.children.child(1, &codes_dtype)?;
85        vortex_ensure!(
86            codes.row_count() == args.row_count,
87            "Dictionary codes row count does not match parent"
88        );
89        Ok(DictData {
90            codes_dtype,
91            all_values_referenced: metadata.all_values_referenced.unwrap_or(false),
92        })
93    }
94
95    fn child_dtype(layout: &Layout<Self>, idx: usize) -> VortexResult<DType> {
96        match idx {
97            0 => Ok(layout.dtype().clone()),
98            1 => Ok(layout.codes_dtype.clone()),
99            _ => vortex_bail!("Dict child index out of bounds: {idx}"),
100        }
101    }
102
103    fn child_type(_layout: &Layout<Self>, idx: usize) -> LayoutChildType {
104        match idx {
105            0 => LayoutChildType::Auxiliary("values".into()),
106            1 => LayoutChildType::Transparent("codes".into()),
107            _ => vortex_panic!("Dict child index out of bounds: {idx}"),
108        }
109    }
110
111    fn new_reader(
112        layout: &Layout<Self>,
113        name: Arc<str>,
114        segment_source: Arc<dyn SegmentSource>,
115        session: &VortexSession,
116        ctx: &LayoutReaderContext,
117    ) -> VortexResult<LayoutReaderRef> {
118        Ok(Arc::new(DictReader::try_new(
119            layout.clone(),
120            name,
121            segment_source,
122            session.clone(),
123            ctx.clone(),
124        )?))
125    }
126}
127
128impl Layout<Dict> {
129    pub(crate) fn new(values: LayoutRef, codes: LayoutRef) -> Self {
130        Self::new_with_all_values_referenced(values, codes, false)
131    }
132
133    fn new_with_all_values_referenced(
134        values: LayoutRef,
135        codes: LayoutRef,
136        all_values_referenced: bool,
137    ) -> Self {
138        let dtype = values.dtype().clone();
139        let row_count = codes.row_count();
140        let codes_dtype = codes.dtype().clone();
141        LayoutParts::new(
142            Dict,
143            dtype,
144            row_count,
145            Vec::new(),
146            OwnedLayoutChildren::layout_children(vec![values, codes]),
147            DictData {
148                codes_dtype,
149                all_values_referenced,
150            },
151        )
152        .into_typed()
153    }
154
155    /// Returns whether every dictionary value is known to be referenced.
156    pub fn has_all_values_referenced(&self) -> bool {
157        self.all_values_referenced
158    }
159}
160
161#[derive(prost::Message)]
162pub struct DictLayoutMetadata {
163    #[prost(enumeration = "PType", tag = "1")]
164    codes_ptype: i32,
165    #[prost(optional, bool, tag = "2")]
166    is_nullable_codes: Option<bool>,
167    #[prost(optional, bool, tag = "3")]
168    pub(crate) all_values_referenced: Option<bool>,
169}
170
171impl DictLayoutMetadata {
172    pub fn new(codes_ptype: PType) -> Self {
173        let mut metadata = Self::default();
174        metadata.set_codes_ptype(codes_ptype);
175        metadata
176    }
177}
vortex_layout/layouts/dict/mod.rs

vortex_layout/layouts/dict/
mod.rs
