akar_storage/
string_dictionary.rs

1//! String dictionary encoding for efficient string storage.
2//!
3//! Dictionary encoding maps strings to integer IDs, storing each unique
4//! string only once in a dictionary. Repeated strings reference the dict
5//! by ID, reducing storage for low-cardinality string columns.
6
7use std::collections::HashMap;
8
9/// A dictionary-encoded string column.
10#[derive(Debug, Clone)]
11pub struct StringDictionary {
12    /// The dictionary: string_id -> string value.
13    strings: Vec<String>,
14    /// Reverse lookup: string value -> string_id.
15    lookup: HashMap<String, u32>,
16}
17
18impl StringDictionary {
19    /// Create a new empty dictionary.
20    pub fn new() -> Self {
21        Self {
22            strings: Vec::new(),
23            lookup: HashMap::new(),
24        }
25    }
26
27    /// Encode a batch of strings, returning (dictionary, encoded_ids).
28    ///
29    /// The dictionary contains unique strings. `encoded_ids` is a `Vec<u32>`
30    /// where each entry is the dictionary ID for the corresponding input string.
31    /// Unknown/NULL strings get ID `u32::MAX`.
32    pub fn encode(strings: &[Option<&str>]) -> (Self, Vec<u32>) {
33        let mut dict = Self::new();
34        let mut ids = Vec::with_capacity(strings.len());
35        for s in strings {
36            let id = match s {
37                None => u32::MAX,
38                Some(val) => dict.intern(val),
39            };
40            ids.push(id);
41        }
42        (dict, ids)
43    }
44
45    /// Insert a single string and return its ID.
46    pub fn intern(&mut self, s: &str) -> u32 {
47        if let Some(&id) = self.lookup.get(s) {
48            return id;
49        }
50        let id = self.strings.len() as u32;
51        self.strings.push(s.to_string());
52        self.lookup.insert(s.to_string(), id);
53        id
54    }
55
56    /// Look up a string by ID. Returns `None` if ID is out of range.
57    pub fn lookup(&self, id: u32) -> Option<&str> {
58        self.strings.get(id as usize).map(|s| s.as_str())
59    }
60
61    /// Look up a string value and return its ID. Returns `None` if not found.
62    pub fn lookup_id(&self, s: &str) -> Option<u32> {
63        self.lookup.get(s).copied()
64    }
65
66    /// Return the number of unique strings in the dictionary.
67    pub fn len(&self) -> usize {
68        self.strings.len()
69    }
70
71    /// Check if the dictionary is empty.
72    pub fn is_empty(&self) -> bool {
73        self.strings.is_empty()
74    }
75
76    /// Serialize the dictionary to bytes.
77    ///
78    /// Format: `[num_strings: u32][for each string: [len: u32][bytes...]]`
79    pub fn serialize(&self) -> Vec<u8> {
80        let mut buf = Vec::new();
81        buf.extend_from_slice(&(self.strings.len() as u32).to_le_bytes());
82        for s in &self.strings {
83            buf.extend_from_slice(&(s.len() as u32).to_le_bytes());
84            buf.extend_from_slice(s.as_bytes());
85        }
86        buf
87    }
88
89    /// Deserialize a dictionary from bytes.
90    pub fn deserialize(data: &[u8]) -> std::io::Result<Self> {
91        if data.len() < 4 {
92            return Err(std::io::Error::new(
93                std::io::ErrorKind::UnexpectedEof,
94                "data too short for dictionary header",
95            ));
96        }
97        let num_strings = u32::from_le_bytes(data[..4].try_into().unwrap()) as usize;
98        let mut dict = Self::with_capacity(num_strings);
99        let mut offset = 4usize;
100        for _ in 0..num_strings {
101            if offset + 4 > data.len() {
102                return Err(std::io::Error::new(
103                    std::io::ErrorKind::UnexpectedEof,
104                    "data too short for string length",
105                ));
106            }
107            let len = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap()) as usize;
108            offset += 4;
109            if offset + len > data.len() {
110                return Err(std::io::Error::new(
111                    std::io::ErrorKind::UnexpectedEof,
112                    "data too short for string content",
113                ));
114            }
115            let s = String::from_utf8(data[offset..offset + len].to_vec())
116                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
117            offset += len;
118            let id = dict.strings.len() as u32;
119            dict.strings.push(s.clone());
120            dict.lookup.insert(s, id);
121        }
122        Ok(dict)
123    }
124
125    /// Compute the memory usage in bytes.
126    pub fn memory_usage(&self) -> usize {
127        let mut total = std::mem::size_of::<Self>();
128        total += self.strings.capacity() * std::mem::size_of::<String>();
129        for s in &self.strings {
130            total += s.capacity();
131        }
132        total += self.lookup.capacity() * (std::mem::size_of::<String>() + std::mem::size_of::<u32>());
133        for k in self.lookup.keys() {
134            total += k.capacity();
135        }
136        total
137    }
138
139    fn with_capacity(cap: usize) -> Self {
140        Self {
141            strings: Vec::with_capacity(cap),
142            lookup: HashMap::with_capacity(cap),
143        }
144    }
145}
146
147impl Default for StringDictionary {
148    fn default() -> Self {
149        Self::new()
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn test_encode_decode_roundtrip() {
159        let input = [
160            Some("apple"),
161            Some("banana"),
162            Some("apple"),
163            Some("cherry"),
164            Some("banana"),
165        ];
166        let (dict, ids) = StringDictionary::encode(&input);
167        assert_eq!(dict.len(), 3);
168        for (i, s) in input.iter().enumerate() {
169            let expected = ids[i];
170            if let Some(val) = s {
171                assert_eq!(dict.lookup(expected), Some(*val));
172            }
173        }
174    }
175
176    #[test]
177    fn test_intern_dedup() {
178        let mut dict = StringDictionary::new();
179        let id1 = dict.intern("hello");
180        let id2 = dict.intern("hello");
181        assert_eq!(id1, id2);
182        let id3 = dict.intern("world");
183        assert_ne!(id1, id3);
184        assert_eq!(dict.len(), 2);
185    }
186
187    #[test]
188    fn test_serialize_roundtrip() {
189        let mut dict = StringDictionary::new();
190        dict.intern("alpha");
191        dict.intern("beta");
192        dict.intern("gamma");
193        let bytes = dict.serialize();
194        let deserialized = StringDictionary::deserialize(&bytes).unwrap();
195        assert_eq!(deserialized.len(), dict.len());
196        assert_eq!(deserialized.lookup_id("alpha"), Some(0));
197        assert_eq!(deserialized.lookup_id("beta"), Some(1));
198        assert_eq!(deserialized.lookup_id("gamma"), Some(2));
199        assert_eq!(deserialized.lookup_id("delta"), None);
200    }
201
202    #[test]
203    fn test_empty_input() {
204        let input: &[Option<&str>] = &[];
205        let (dict, ids) = StringDictionary::encode(input);
206        assert!(dict.is_empty());
207        assert!(ids.is_empty());
208    }
209
210    #[test]
211    fn test_null_handling() {
212        let input = [Some("a"), None, Some("b"), None];
213        let (dict, ids) = StringDictionary::encode(&input);
214        assert_eq!(dict.len(), 2);
215        assert_eq!(ids[0], 0);
216        assert_eq!(ids[1], u32::MAX);
217        assert_eq!(ids[2], 1);
218        assert_eq!(ids[3], u32::MAX);
219    }
220
221    #[test]
222    fn test_lookup_miss() {
223        let mut dict = StringDictionary::new();
224        dict.intern("foo");
225        assert_eq!(dict.lookup(0), Some("foo"));
226        assert_eq!(dict.lookup(1), None);
227        assert_eq!(dict.lookup(u32::MAX), None);
228        assert_eq!(dict.lookup_id("bar"), None);
229    }
230
231    #[test]
232    fn test_memory_usage() {
233        let mut dict = StringDictionary::new();
234        dict.intern("short");
235        dict.intern("a longer string value");
236        let usage = dict.memory_usage();
237        assert!(usage > 0);
238        assert!(usage > std::mem::size_of::<StringDictionary>());
239    }
240
241    #[test]
242    fn test_deserialize_empty() {
243        let bytes = 0u32.to_le_bytes().to_vec();
244        let dict = StringDictionary::deserialize(&bytes).unwrap();
245        assert!(dict.is_empty());
246        assert_eq!(dict.len(), 0);
247    }
248
249    #[test]
250    fn test_deserialize_truncated() {
251        let result = StringDictionary::deserialize(&[1, 0, 0, 0]);
252        assert!(result.is_err());
253    }
254
255    #[test]
256    fn test_deserialize_invalid_utf8() {
257        let bytes = [1u8, 0, 0, 0, 1, 0, 0, 0, 0xFF];
258        let result = StringDictionary::deserialize(&bytes);
259        assert!(result.is_err());
260    }
261
262    #[test]
263    fn test_empty_string() {
264        let mut dict = StringDictionary::new();
265        let id = dict.intern("");
266        assert_eq!(dict.lookup(id), Some(""));
267        assert_eq!(dict.lookup_id(""), Some(id));
268    }
269
270    #[test]
271    fn test_compression_integration() {
272        let dict = StringDictionary::new();
273        let serialized = dict.serialize();
274        let chunk = crate::compression::compress(akar_common::enums::CompressionType::StringDictionary, &serialized, 0);
275        let decompressed = crate::compression::decompress(&chunk, serialized.len());
276        let deserialized = StringDictionary::deserialize(&decompressed).unwrap();
277        assert!(deserialized.is_empty());
278    }
279}
akar_storage/string_dictionary.rs

akar_storage/
string_dictionary.rs
