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::ArrayContext;
11use vortex_array::DeserializeMetadata;
12use vortex_array::ProstMetadata;
13use vortex_dtype::DType;
14use vortex_dtype::Nullability;
15use vortex_dtype::PType;
16use vortex_error::VortexExpect;
17use vortex_error::VortexResult;
18use vortex_error::vortex_bail;
19use vortex_error::vortex_ensure;
20use vortex_error::vortex_err;
21use vortex_error::vortex_panic;
22use vortex_session::VortexSession;
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 DictVTable {
38    type Layout = DictLayout;
39    type Encoding = DictLayoutEncoding;
40    type Metadata = ProstMetadata<DictLayoutMetadata>;
41
42    fn id(_encoding: &Self::Encoding) -> LayoutId {
43        LayoutId::new_ref("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(layout.values.clone()),
77            1 => Ok(layout.codes.clone()),
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    ) -> VortexResult<LayoutReaderRef> {
96        Ok(Arc::new(DictReader::try_new(
97            layout.clone(),
98            name,
99            segment_source,
100            session.clone(),
101        )?))
102    }
103
104    fn build(
105        _encoding: &Self::Encoding,
106        dtype: &DType,
107        _row_count: u64,
108        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
109        _segment_ids: Vec<SegmentId>,
110        children: &dyn LayoutChildren,
111        _ctx: &ArrayContext,
112    ) -> VortexResult<Self::Layout> {
113        let values = children.child(0, dtype)?;
114        let codes_nullable = metadata
115            .is_nullable_codes
116            .map(Nullability::from)
117            // The old behaviour (without `is_nullable_codes` metadata) used the nullability
118            // of the values (and whole array).
119            // see [`SerdeVTable<DictVTable>::build`].
120            .unwrap_or_else(|| dtype.nullability());
121        let codes = children.child(1, &DType::Primitive(metadata.codes_ptype(), codes_nullable))?;
122        Ok(unsafe {
123            DictLayout::new(values, codes)
124                .set_all_values_referenced(metadata.all_values_referenced.unwrap_or(false))
125        })
126    }
127
128    fn with_children(layout: &mut Self::Layout, children: Vec<LayoutRef>) -> VortexResult<()> {
129        vortex_ensure!(
130            children.len() == 2,
131            "DictLayout expects exactly 2 children (values, codes), got {}",
132            children.len()
133        );
134        let mut children_iter = children.into_iter();
135        layout.values = children_iter
136            .next()
137            .ok_or_else(|| vortex_err!("Missing values child"))?;
138        layout.codes = children_iter
139            .next()
140            .ok_or_else(|| vortex_err!("Missing codes child"))?;
141        Ok(())
142    }
143}
144
145#[derive(Debug)]
146pub struct DictLayoutEncoding;
147
148#[derive(Clone, Debug)]
149pub struct DictLayout {
150    values: LayoutRef,
151    codes: LayoutRef,
152    /// Indicates whether all dictionary values are definitely referenced by at least one code.
153    /// `true` = all values are referenced (computed during encoding).
154    /// `false` = unknown/might have unreferenced values.
155    all_values_referenced: bool,
156}
157
158impl DictLayout {
159    pub(crate) fn new(values: LayoutRef, codes: LayoutRef) -> Self {
160        Self {
161            values,
162            codes,
163            all_values_referenced: false,
164        }
165    }
166
167    /// Set whether all dictionary values are definitely referenced.
168    ///
169    /// # Safety
170    /// The caller must ensure that when setting `all_values_referenced = true`, ALL dictionary
171    /// values are actually referenced by at least one valid code. Setting this incorrectly can
172    /// lead to incorrect query results in operations like min/max.
173    ///
174    /// This is typically only set to `true` during dictionary encoding when we know for certain
175    /// that all values are referenced.
176    /// See `DictArray::set_all_values_referenced`.
177    pub unsafe fn set_all_values_referenced(mut self, all_values_referenced: bool) -> Self {
178        self.all_values_referenced = all_values_referenced;
179        self
180    }
181
182    pub fn has_all_values_referenced(&self) -> bool {
183        self.all_values_referenced
184    }
185}
186
187#[derive(prost::Message)]
188pub struct DictLayoutMetadata {
189    #[prost(enumeration = "PType", tag = "1")]
190    // i32 is required for proto, use the generated getter to read this field.
191    codes_ptype: i32,
192    // nullable codes are optional since they were added after stabilisation
193    #[prost(optional, bool, tag = "2")]
194    is_nullable_codes: Option<bool>,
195    // all_values_referenced is optional for backward compatibility
196    // true = all dictionary values are definitely referenced by at least one code
197    // false/None = unknown whether all values are referenced (conservative default)
198    // see `DictArray::all_values_referenced`
199    #[prost(optional, bool, tag = "3")]
200    pub(crate) all_values_referenced: Option<bool>,
201}
202
203impl DictLayoutMetadata {
204    pub fn new(codes_ptype: PType) -> Self {
205        let mut metadata = Self::default();
206        metadata.set_codes_ptype(codes_ptype);
207        metadata
208    }
209}
vortex_layout/layouts/dict/mod.rs

vortex_layout/layouts/dict/
mod.rs
