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_eq;
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_eq!(
75            args.children.nchildren(),
76            2,
77            "DictLayout expects exactly 2 children"
78        );
79        let codes_nullable = metadata
80            .is_nullable_codes
81            .map(Nullability::from)
82            .unwrap_or_else(|| args.dtype.nullability());
83        let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable);
84        args.children.child(0, args.dtype)?;
85        let codes = args.children.child(1, &codes_dtype)?;
86        vortex_ensure_eq!(
87            codes.row_count(),
88            args.row_count,
89            "Dictionary codes row count does not match parent"
90        );
91        Ok(DictData {
92            codes_dtype,
93            all_values_referenced: metadata.all_values_referenced.unwrap_or(false),
94        })
95    }
96
97    fn child_dtype(layout: &Layout<Self>, idx: usize) -> VortexResult<DType> {
98        match idx {
99            0 => Ok(layout.dtype().clone()),
100            1 => Ok(layout.codes_dtype.clone()),
101            _ => vortex_bail!("Dict child index out of bounds: {idx}"),
102        }
103    }
104
105    fn child_type(_layout: &Layout<Self>, idx: usize) -> LayoutChildType {
106        match idx {
107            0 => LayoutChildType::Auxiliary("values".into()),
108            1 => LayoutChildType::Transparent("codes".into()),
109            _ => vortex_panic!("Dict child index out of bounds: {idx}"),
110        }
111    }
112
113    fn new_reader(
114        layout: &Layout<Self>,
115        name: Arc<str>,
116        segment_source: Arc<dyn SegmentSource>,
117        session: &VortexSession,
118        ctx: &LayoutReaderContext,
119    ) -> VortexResult<LayoutReaderRef> {
120        Ok(Arc::new(DictReader::try_new(
121            layout.clone(),
122            name,
123            segment_source,
124            session.clone(),
125            ctx.clone(),
126        )?))
127    }
128}
129
130impl Layout<Dict> {
131    pub(crate) fn new(values: LayoutRef, codes: LayoutRef) -> Self {
132        Self::new_with_all_values_referenced(values, codes, false)
133    }
134
135    fn new_with_all_values_referenced(
136        values: LayoutRef,
137        codes: LayoutRef,
138        all_values_referenced: bool,
139    ) -> Self {
140        let dtype = values.dtype().clone();
141        let row_count = codes.row_count();
142        let codes_dtype = codes.dtype().clone();
143        LayoutParts::new(
144            Dict,
145            dtype,
146            row_count,
147            Vec::new(),
148            OwnedLayoutChildren::layout_children(vec![values, codes]),
149            DictData {
150                codes_dtype,
151                all_values_referenced,
152            },
153        )
154        .into_typed()
155    }
156
157    /// Returns whether every dictionary value is known to be referenced.
158    pub fn has_all_values_referenced(&self) -> bool {
159        self.all_values_referenced
160    }
161}
162
163#[derive(prost::Message)]
164pub struct DictLayoutMetadata {
165    #[prost(enumeration = "PType", tag = "1")]
166    codes_ptype: i32,
167    #[prost(optional, bool, tag = "2")]
168    is_nullable_codes: Option<bool>,
169    #[prost(optional, bool, tag = "3")]
170    pub(crate) all_values_referenced: Option<bool>,
171}
172
173impl DictLayoutMetadata {
174    pub fn new(codes_ptype: PType) -> Self {
175        let mut metadata = Self::default();
176        metadata.set_codes_ptype(codes_ptype);
177        metadata
178    }
179}
vortex_layout/layouts/dict/mod.rs

vortex_layout/layouts/dict/
mod.rs
