ferritin_core/columns/
dictionary.rs

1//! Dictionary encoding for string columns.
2
3use super::error::{BoundaryError, to_u32};
4use std::collections::HashMap;
5
6/// A dictionary-encoded string column.
7///
8/// Each row holds a `u32` code into a table of unique values. The value table
9/// is stored Arrow-style as one UTF-8 byte buffer plus `n_values() + 1`
10/// offsets, so both halves cross an FFI or wasm-bindgen boundary as plain
11/// numeric buffers (`Uint32Array` codes, `Uint8Array` bytes, `Uint32Array`
12/// offsets) and the consumer decodes each unique value once.
13///
14/// Codes are assigned in first-occurrence order, so encoding the same input
15/// always yields the same codes.
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct StringDictionary {
18    codes: Vec<u32>,
19    value_bytes: Vec<u8>,
20    value_offsets: Vec<u32>,
21}
22
23impl StringDictionary {
24    /// Encode a sequence of strings.
25    ///
26    /// Fails with [`CapacityExceeded`](super::BoundaryErrorKind::CapacityExceeded)
27    /// if the row count or total value bytes do not fit in `u32`.
28    pub fn encode<'s>(values: impl IntoIterator<Item = &'s str>) -> Result<Self, BoundaryError> {
29        let mut lookup: HashMap<&'s str, u32> = HashMap::new();
30        let mut codes = Vec::new();
31        let mut value_bytes = Vec::new();
32        let mut value_offsets = vec![0u32];
33
34        for value in values {
35            let code = match lookup.get(value) {
36                Some(&code) => code,
37                None => {
38                    let code = to_u32("dictionary value count", lookup.len())?;
39                    value_bytes.extend_from_slice(value.as_bytes());
40                    value_offsets.push(to_u32("dictionary value bytes", value_bytes.len())?);
41                    lookup.insert(value, code);
42                    code
43                }
44            };
45            codes.push(code);
46        }
47        to_u32("dictionary row count", codes.len())?;
48
49        Ok(Self {
50            codes,
51            value_bytes,
52            value_offsets,
53        })
54    }
55
56    /// Per-row codes into the value table.
57    pub fn codes(&self) -> &[u32] {
58        &self.codes
59    }
60
61    /// Concatenated UTF-8 bytes of all unique values.
62    pub fn value_bytes(&self) -> &[u8] {
63        &self.value_bytes
64    }
65
66    /// Offsets into [`Self::value_bytes`]; value `c` spans
67    /// `value_offsets[c]..value_offsets[c + 1]`.
68    pub fn value_offsets(&self) -> &[u32] {
69        &self.value_offsets
70    }
71
72    /// Number of rows.
73    pub fn len(&self) -> usize {
74        self.codes.len()
75    }
76
77    /// Returns `true` if there are no rows.
78    pub fn is_empty(&self) -> bool {
79        self.codes.is_empty()
80    }
81
82    /// Number of unique values.
83    pub fn n_values(&self) -> usize {
84        self.value_offsets.len() - 1
85    }
86
87    /// The unique value for `code`, or `None` if the code is out of range.
88    pub fn value(&self, code: u32) -> Option<&str> {
89        let code = code as usize;
90        let start = *self.value_offsets.get(code)? as usize;
91        let end = *self.value_offsets.get(code + 1)? as usize;
92        // Bytes were copied from `&str` values and split on their boundaries.
93        std::str::from_utf8(&self.value_bytes[start..end]).ok()
94    }
95
96    /// The decoded value of row `row`, or `None` if the row is out of range.
97    pub fn get(&self, row: usize) -> Option<&str> {
98        self.value(*self.codes.get(row)?)
99    }
100
101    /// Iterate over the unique values in code order.
102    pub fn values(&self) -> impl Iterator<Item = &str> + '_ {
103        (0..self.n_values() as u32).filter_map(|code| self.value(code))
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn test_dictionary_first_occurrence_codes() {
113        let dict = StringDictionary::encode(["CA", "N", "CA", "C", "N"]).unwrap();
114        assert_eq!(dict.codes(), &[0, 1, 0, 2, 1]);
115        assert_eq!(dict.n_values(), 3);
116        assert_eq!(dict.value_bytes(), b"CANC");
117        assert_eq!(dict.value_offsets(), &[0, 2, 3, 4]);
118        assert_eq!(dict.values().collect::<Vec<_>>(), vec!["CA", "N", "C"]);
119    }
120
121    #[test]
122    fn test_dictionary_roundtrip_and_bounds() {
123        let input = ["ALA", "", "HOH", "ALA", "Ω"];
124        let dict = StringDictionary::encode(input).unwrap();
125        let decoded: Vec<&str> = (0..dict.len()).map(|i| dict.get(i).unwrap()).collect();
126        assert_eq!(decoded, input);
127        assert_eq!(dict.get(input.len()), None);
128        assert_eq!(dict.value(dict.n_values() as u32), None);
129    }
130
131    #[test]
132    fn test_dictionary_empty() {
133        let dict = StringDictionary::encode(std::iter::empty()).unwrap();
134        assert!(dict.is_empty());
135        assert_eq!(dict.n_values(), 0);
136        assert_eq!(dict.value_offsets(), &[0]);
137    }
138}
ferritin_core/columns/dictionary.rs

ferritin_core/columns/
dictionary.rs
