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

vortex_layout/layouts/dict/
mod.rs
