zeppelin_embed/meta/
dict.rs

1//! Segment-local UTF-8 dictionaries and automatically widened row codes.
2
3use std::collections::HashMap;
4
5/// A typed failure while constructing a string dictionary.
6#[derive(Clone, Copy, Debug, Eq, PartialEq)]
7pub enum DictionaryError {
8    /// The dictionary cardinality exceeded the u32 code space.
9    CardinalityOverflow,
10    /// UTF-8 bytes or offsets exceeded the u32 side-array representation.
11    StringStorageOverflow,
12}
13
14impl std::fmt::Display for DictionaryError {
15    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        match self {
17            Self::CardinalityOverflow => formatter.write_str("dictionary cardinality exceeds u32"),
18            Self::StringStorageOverflow => formatter.write_str("string storage exceeds u32"),
19        }
20    }
21}
22
23impl std::error::Error for DictionaryError {}
24
25/// The physical width selected for a dictionary-code array.
26#[derive(Clone, Copy, Debug, Eq, PartialEq)]
27pub enum CodeWidth {
28    /// Sixteen-bit codes for at most `u16::MAX` distinct values.
29    U16,
30    /// Thirty-two-bit codes for larger dictionaries.
31    U32,
32}
33
34/// A contiguous dictionary-code array selected at segment seal time.
35#[derive(Clone, Debug, Eq, PartialEq)]
36pub enum DictionaryCodes {
37    /// Sixteen-bit row codes.
38    U16(Vec<u16>),
39    /// Thirty-two-bit row codes.
40    U32(Vec<u32>),
41}
42
43impl DictionaryCodes {
44    pub(crate) fn resident_bytes(&self) -> Option<usize> {
45        match self {
46            Self::U16(codes) => codes.capacity().checked_mul(std::mem::size_of::<u16>()),
47            Self::U32(codes) => codes.capacity().checked_mul(std::mem::size_of::<u32>()),
48        }
49    }
50
51    /// Returns the physical code width.
52    #[must_use]
53    pub const fn width(&self) -> CodeWidth {
54        match self {
55            Self::U16(_) => CodeWidth::U16,
56            Self::U32(_) => CodeWidth::U32,
57        }
58    }
59
60    /// Returns the number of row codes.
61    #[must_use]
62    pub fn len(&self) -> usize {
63        match self {
64            Self::U16(codes) => codes.len(),
65            Self::U32(codes) => codes.len(),
66        }
67    }
68
69    /// Returns whether there are no row codes.
70    #[must_use]
71    pub fn is_empty(&self) -> bool {
72        self.len() == 0
73    }
74
75    /// Reads one row code without exposing its physical width.
76    #[must_use]
77    pub fn get(&self, row: usize) -> Option<u32> {
78        match self {
79            Self::U16(codes) => codes.get(row).copied().map(u32::from),
80            Self::U32(codes) => codes.get(row).copied(),
81        }
82    }
83
84    pub(crate) fn iter(&self) -> impl Iterator<Item = u32> + '_ {
85        enum Iter<'a> {
86            U16(std::slice::Iter<'a, u16>),
87            U32(std::slice::Iter<'a, u32>),
88        }
89
90        impl Iterator for Iter<'_> {
91            type Item = u32;
92
93            fn next(&mut self) -> Option<Self::Item> {
94                match self {
95                    Iter::U16(values) => values.next().copied().map(u32::from),
96                    Iter::U32(values) => values.next().copied(),
97                }
98            }
99        }
100
101        match self {
102            Self::U16(values) => Iter::U16(values.iter()),
103            Self::U32(values) => Iter::U32(values.iter()),
104        }
105    }
106
107    pub(crate) fn from_u32(codes: Vec<u32>, cardinality: u64) -> Result<Self, DictionaryError> {
108        match code_width_for_cardinality(cardinality)? {
109            CodeWidth::U16 => {
110                let narrowed = codes
111                    .into_iter()
112                    .map(u16::try_from)
113                    .collect::<Result<Vec<_>, _>>()
114                    .map_err(|_| DictionaryError::CardinalityOverflow)?;
115                Ok(Self::U16(narrowed))
116            }
117            CodeWidth::U32 => Ok(Self::U32(codes)),
118        }
119    }
120}
121
122pub(crate) fn code_width_for_cardinality(cardinality: u64) -> Result<CodeWidth, DictionaryError> {
123    if cardinality <= u64::from(u16::MAX) {
124        Ok(CodeWidth::U16)
125    } else if cardinality <= u64::from(u32::MAX) {
126        Ok(CodeWidth::U32)
127    } else {
128        Err(DictionaryError::CardinalityOverflow)
129    }
130}
131
132#[derive(Clone, Debug, Eq, PartialEq)]
133#[repr(C)]
134pub(crate) struct StringStorage {
135    offsets: Vec<u32>,
136    bytes: Vec<u8>,
137}
138
139impl StringStorage {
140    pub(crate) fn resident_bytes(&self) -> Option<usize> {
141        self.offsets
142            .capacity()
143            .checked_mul(std::mem::size_of::<u32>())?
144            .checked_add(self.bytes.capacity())
145    }
146    pub(crate) fn new() -> Self {
147        Self {
148            offsets: vec![0],
149            bytes: Vec::new(),
150        }
151    }
152
153    pub(crate) fn can_push(&self, value: &str) -> Result<(), DictionaryError> {
154        let new_length = self
155            .bytes
156            .len()
157            .checked_add(value.len())
158            .ok_or(DictionaryError::StringStorageOverflow)?;
159        u32::try_from(new_length)
160            .map(|_| ())
161            .map_err(|_| DictionaryError::StringStorageOverflow)
162    }
163
164    pub(crate) fn push(&mut self, value: &str) -> Result<(), DictionaryError> {
165        self.can_push(value)?;
166        self.bytes.extend_from_slice(value.as_bytes());
167        let offset =
168            u32::try_from(self.bytes.len()).map_err(|_| DictionaryError::StringStorageOverflow)?;
169        self.offsets.push(offset);
170        Ok(())
171    }
172
173    pub(crate) fn len(&self) -> usize {
174        self.offsets.len().saturating_sub(1)
175    }
176
177    #[cfg(any(test, feature = "test-seams"))]
178    pub(crate) fn byte_len(&self) -> usize {
179        self.bytes.len()
180    }
181
182    pub(crate) fn get(&self, position: usize) -> Option<&str> {
183        let start = self
184            .offsets
185            .get(position)
186            .copied()
187            .map(|value| value as usize)?;
188        let end = self
189            .offsets
190            .get(position.saturating_add(1))
191            .copied()
192            .map(|value| value as usize)?;
193        let bytes = self.bytes.get(start..end)?;
194        std::str::from_utf8(bytes).ok()
195    }
196}
197
198/// An immutable segment-local UTF-8 dictionary.
199#[derive(Clone, Debug, Eq, PartialEq)]
200#[repr(C)]
201pub struct StringDictionary {
202    values: StringStorage,
203}
204
205impl StringDictionary {
206    pub(crate) fn resident_bytes(&self) -> Option<usize> {
207        self.values.resident_bytes()
208    }
209
210    /// Returns the number of distinct values.
211    #[must_use]
212    pub fn len(&self) -> usize {
213        self.values.len()
214    }
215
216    /// Returns whether the dictionary is empty.
217    #[must_use]
218    pub fn is_empty(&self) -> bool {
219        self.len() == 0
220    }
221
222    /// Decodes one dictionary code.
223    #[must_use]
224    pub fn get(&self, code: u32) -> Option<&str> {
225        let position = usize::try_from(code).ok()?;
226        self.values.get(position)
227    }
228
229    /// Finds the code assigned to a string value.
230    #[must_use]
231    pub fn find(&self, needle: &str) -> Option<u32> {
232        self.values
233            .offsets
234            .windows(2)
235            .enumerate()
236            .find_map(|(position, _)| {
237                (self.values.get(position) == Some(needle))
238                    .then(|| u32::try_from(position).ok())
239                    .flatten()
240            })
241    }
242}
243
244pub(crate) struct DictionaryBuilder {
245    lookup: HashMap<String, u32>,
246    values: StringStorage,
247}
248
249impl DictionaryBuilder {
250    pub(crate) fn new() -> Self {
251        Self {
252            lookup: HashMap::new(),
253            values: StringStorage::new(),
254        }
255    }
256
257    pub(crate) fn can_intern(&self, value: &str) -> Result<(), DictionaryError> {
258        if self.lookup.contains_key(value) {
259            return Ok(());
260        }
261        let cardinality = u64::try_from(self.lookup.len())
262            .unwrap_or(u64::MAX)
263            .saturating_add(1);
264        code_width_for_cardinality(cardinality)?;
265        self.values.can_push(value)
266    }
267
268    pub(crate) fn intern(&mut self, value: &str) -> Result<u32, DictionaryError> {
269        if let Some(code) = self.lookup.get(value).copied() {
270            return Ok(code);
271        }
272        self.can_intern(value)?;
273        let code =
274            u32::try_from(self.lookup.len()).map_err(|_| DictionaryError::CardinalityOverflow)?;
275        self.values.push(value)?;
276        self.lookup.insert(value.to_owned(), code);
277        Ok(code)
278    }
279
280    pub(crate) fn cardinality(&self) -> u64 {
281        u64::try_from(self.lookup.len()).unwrap_or(u64::MAX)
282    }
283
284    #[cfg(any(test, feature = "test-seams"))]
285    pub(crate) fn contains(&self, value: &str) -> bool {
286        self.lookup.contains_key(value)
287    }
288
289    #[cfg(any(test, feature = "test-seams"))]
290    pub(crate) fn string_bytes(&self) -> usize {
291        self.values.byte_len()
292    }
293
294    pub(crate) fn finish(self) -> StringDictionary {
295        StringDictionary {
296            values: self.values,
297        }
298    }
299}
300
301#[allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)]
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use proptest::prelude::*;
306    use proptest::test_runner::{Config, RngSeed, TestRunner};
307    use rand::RngCore;
308
309    #[test]
310    fn prop_dict_roundtrip() {
311        let name = "meta::dict::prop_dict_roundtrip";
312        let mut seeded = crate::test_support::seeded_rng(name);
313        let config = Config {
314            rng_seed: RngSeed::Fixed(seeded.next_u64()),
315            ..Config::default()
316        };
317        let mut runner = TestRunner::new(config);
318        let values = prop::collection::vec("[a-z]{0,16}", 0..=256);
319        let result = runner.run(&values, |values| {
320            let mut builder = DictionaryBuilder::new();
321            let codes = values
322                .iter()
323                .map(|value| builder.intern(value).expect("bounded dictionary"))
324                .collect::<Vec<_>>();
325            let dictionary = builder.finish();
326            let decoded = codes
327                .into_iter()
328                .map(|code| dictionary.get(code).map(str::to_owned))
329                .collect::<Vec<_>>();
330            let expected = values.into_iter().map(Some).collect::<Vec<_>>();
331            prop_assert_eq!(decoded, expected);
332            Ok(())
333        });
334        assert!(result.is_ok(), "property result: {result:?}");
335    }
336
337    #[test]
338    fn dictionary_auto_widens_u16_to_u32() {
339        let cardinality = u64::from(u16::MAX).saturating_add(1);
340        let codes = (0..=u32::from(u16::MAX)).collect::<Vec<_>>();
341        let encoded = DictionaryCodes::from_u32(codes, cardinality).expect("u32 cardinality");
342        assert_eq!(encoded.width(), CodeWidth::U32);
343    }
344
345    #[test]
346    fn dictionary_reports_u32_cardinality_overflow() {
347        assert_eq!(
348            code_width_for_cardinality(u64::from(u32::MAX).saturating_add(1)),
349            Err(DictionaryError::CardinalityOverflow)
350        );
351    }
352}
zeppelin_embed/meta/dict.rs

zeppelin_embed/meta/
dict.rs
