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::DeserializeMetadata;
11use vortex_array::ProstMetadata;
12use vortex_array::dtype::DType;
13use vortex_array::dtype::Nullability;
14use vortex_array::dtype::PType;
15use vortex_error::VortexExpect;
16use vortex_error::VortexResult;
17use vortex_error::vortex_bail;
18use vortex_error::vortex_ensure;
19use vortex_error::vortex_err;
20use vortex_error::vortex_panic;
21use vortex_session::VortexSession;
22use vortex_session::registry::ReadContext;
23
24use crate::LayoutChildType;
25use crate::LayoutEncodingRef;
26use crate::LayoutId;
27use crate::LayoutReaderRef;
28use crate::LayoutRef;
29use crate::VTable;
30use crate::children::LayoutChildren;
31use crate::segments::SegmentId;
32use crate::segments::SegmentSource;
33use crate::vtable;
34
35vtable!(Dict);
36
37impl VTable for Dict {
38    type Layout = DictLayout;
39    type Encoding = DictLayoutEncoding;
40    type Metadata = ProstMetadata<DictLayoutMetadata>;
41
42    fn id(_encoding: &Self::Encoding) -> LayoutId {
43        LayoutId::new("vortex.dict")
44    }
45
46    fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef {
47        LayoutEncodingRef::new_ref(DictLayoutEncoding.as_ref())
48    }
49
50    fn row_count(layout: &Self::Layout) -> u64 {
51        layout.codes.row_count()
52    }
53
54    fn dtype(layout: &Self::Layout) -> &DType {
55        layout.values.dtype()
56    }
57
58    fn metadata(layout: &Self::Layout) -> Self::Metadata {
59        let mut metadata =
60            DictLayoutMetadata::new(PType::try_from(layout.codes.dtype()).vortex_expect("ptype"));
61        metadata.is_nullable_codes = Some(layout.codes.dtype().is_nullable());
62        metadata.all_values_referenced = Some(layout.all_values_referenced);
63        ProstMetadata(metadata)
64    }
65
66    fn segment_ids(_layout: &Self::Layout) -> Vec<SegmentId> {
67        vec![]
68    }
69
70    fn nchildren(_layout: &Self::Layout) -> usize {
71        2
72    }
73
74    fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef> {
75        match idx {
76            0 => Ok(Arc::clone(&layout.values)),
77            1 => Ok(Arc::clone(&layout.codes)),
78            _ => vortex_bail!("Unreachable child index: {}", idx),
79        }
80    }
81
82    fn child_type(_layout: &Self::Layout, idx: usize) -> LayoutChildType {
83        match idx {
84            0 => LayoutChildType::Auxiliary("values".into()),
85            1 => LayoutChildType::Transparent("codes".into()),
86            _ => vortex_panic!("Unreachable child index: {}", idx),
87        }
88    }
89
90    fn new_reader(
91        layout: &Self::Layout,
92        name: Arc<str>,
93        segment_source: Arc<dyn SegmentSource>,
94        session: &VortexSession,
95        ctx: &crate::LayoutReaderContext,
96    ) -> VortexResult<LayoutReaderRef> {
97        Ok(Arc::new(DictReader::try_new(
98            layout.clone(),
99            name,
100            segment_source,
101            session.clone(),
102            ctx.clone(),
103        )?))
104    }
105
106    fn build(
107        _encoding: &Self::Encoding,
108        dtype: &DType,
109        _row_count: u64,
110        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
111        _segment_ids: Vec<SegmentId>,
112        children: &dyn LayoutChildren,
113        _ctx: &ReadContext,
114    ) -> VortexResult<Self::Layout> {
115        let values = children.child(0, dtype)?;
116        let codes_nullable = metadata
117            .is_nullable_codes
118            .map(Nullability::from)
119            // The old behaviour (without `is_nullable_codes` metadata) used the nullability
120            // of the values (and whole array).
121            // see [`SerdeVTable<Dict>::build`].
122            .unwrap_or_else(|| dtype.nullability());
123        let codes = children.child(1, &DType::Primitive(metadata.codes_ptype(), codes_nullable))?;
124        Ok(unsafe {
125            DictLayout::new(values, codes)
126                .set_all_values_referenced(metadata.all_values_referenced.unwrap_or(false))
127        })
128    }
129
130    fn with_children(layout: &mut Self::Layout, children: Vec<LayoutRef>) -> VortexResult<()> {
131        vortex_ensure!(
132            children.len() == 2,
133            "DictLayout expects exactly 2 children (values, codes), got {}",
134            children.len()
135        );
136        let mut children_iter = children.into_iter();
137        layout.values = children_iter
138            .next()
139            .ok_or_else(|| vortex_err!("Missing values child"))?;
140        layout.codes = children_iter
141            .next()
142            .ok_or_else(|| vortex_err!("Missing codes child"))?;
143        Ok(())
144    }
145}
146
147#[derive(Debug)]
148pub struct DictLayoutEncoding;
149
150/// Stores a shared dictionary of values alongside compact integer codes that index into it.
151///
152/// Useful for columns with many repeated values, where storing each value once and
153/// referencing it by index saves significant space.
154#[derive(Clone, Debug)]
155pub struct DictLayout {
156    values: LayoutRef,
157    codes: LayoutRef,
158    /// Indicates whether all dictionary values are definitely referenced by at least one code.
159    /// `true` = all values are referenced (computed during encoding).
160    /// `false` = unknown/might have unreferenced values.
161    all_values_referenced: bool,
162}
163
164impl DictLayout {
165    pub(crate) fn new(values: LayoutRef, codes: LayoutRef) -> Self {
166        Self {
167            values,
168            codes,
169            all_values_referenced: false,
170        }
171    }
172
173    /// Set whether all dictionary values are definitely referenced.
174    ///
175    /// # Safety
176    /// The caller must ensure that when setting `all_values_referenced = true`, ALL dictionary
177    /// values are actually referenced by at least one valid code. Setting this incorrectly can
178    /// lead to incorrect query results in operations like min/max.
179    ///
180    /// This is typically only set to `true` during dictionary encoding when we know for certain
181    /// that all values are referenced.
182    /// See `DictArray::set_all_values_referenced`.
183    pub unsafe fn set_all_values_referenced(mut self, all_values_referenced: bool) -> Self {
184        self.all_values_referenced = all_values_referenced;
185        self
186    }
187
188    pub fn has_all_values_referenced(&self) -> bool {
189        self.all_values_referenced
190    }
191}
192
193#[derive(prost::Message)]
194pub struct DictLayoutMetadata {
195    #[prost(enumeration = "PType", tag = "1")]
196    // i32 is required for proto, use the generated getter to read this field.
197    codes_ptype: i32,
198    // nullable codes are optional since they were added after stabilisation
199    #[prost(optional, bool, tag = "2")]
200    is_nullable_codes: Option<bool>,
201    // all_values_referenced is optional for backward compatibility
202    // true = all dictionary values are definitely referenced by at least one code
203    // false/None = unknown whether all values are referenced (conservative default)
204    // see `DictArray::all_values_referenced`
205    #[prost(optional, bool, tag = "3")]
206    pub(crate) all_values_referenced: Option<bool>,
207}
208
209impl DictLayoutMetadata {
210    pub fn new(codes_ptype: PType) -> Self {
211        let mut metadata = Self::default();
212        metadata.set_codes_ptype(codes_ptype);
213        metadata
214    }
215}
vortex_layout/layouts/dict/mod.rs

vortex_layout/layouts/dict/
mod.rs
