scirs2_io/columnar/
dictionary.rs

1//! Dictionary encoding for columnar data.
2//!
3//! Replaces repeated string (or generic) values with small integer codes,
4//! achieving significant compression for low-cardinality columns.
5//!
6//! ## Design
7//!
8//! - Vocabulary is built from training data sorted by descending frequency.
9//! - Codes are `u16`; code `u16::MAX` (65535) is reserved as the "null" sentinel.
10//! - Maximum vocabulary size is `u16::MAX - 1` = 65534 entries.
11//!
12//! ## Example
13//!
14//! ```rust
15//! use scirs2_io::columnar::dictionary::{DictionaryEncoder, UnknownPolicy};
16//!
17//! let data = vec!["red", "blue", "red", "green", "red", "blue"];
18//! let enc: DictionaryEncoder<String> =
19//!     DictionaryEncoder::fit(
20//!         &data.iter().map(|s| s.to_string()).collect::<Vec<_>>(),
21//!         UnknownPolicy::Error,
22//!     ).expect("fit failed");
23//!
24//! let codes = enc.encode(
25//!     &data.iter().map(|s| s.to_string()).collect::<Vec<_>>()
26//! ).expect("encode failed");
27//! assert_eq!(codes.len(), 6);
28//! ```
29
30use std::collections::HashMap;
31use std::hash::Hash;
32
33use crate::error::IoError;
34
35// ---------------------------------------------------------------------------
36// Policy
37// ---------------------------------------------------------------------------
38
39/// How to handle values encountered during encoding that are not in the
40/// vocabulary built by [`DictionaryEncoder::fit`].
41#[non_exhaustive]
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum UnknownPolicy {
44    /// Return an [`IoError`] when an unknown value is encountered.
45    Error,
46    /// Map unknown values to `u16::MAX` (the null sentinel code).
47    UseNullCode,
48    /// Dynamically extend the vocabulary; errors if capacity is exhausted.
49    AddToVocab,
50}
51
52// ---------------------------------------------------------------------------
53// DictionaryEncoder
54// ---------------------------------------------------------------------------
55
56/// Dictionary encoder for a generic value type `T`.
57///
58/// `T` must be `Eq + Hash + Clone + Ord`; the `Ord` bound is used only for
59/// tie-breaking when two values have the same frequency.
60#[derive(Debug, Clone)]
61pub struct DictionaryEncoder<T: Eq + Hash + Clone> {
62    /// Vocabulary: `vocab[code]` → value.  Sorted by descending frequency.
63    pub vocab: Vec<T>,
64    /// Reverse map: value → code.
65    code_map: HashMap<T, u16>,
66    /// Policy for values not present in the vocabulary.
67    pub unknown_policy: UnknownPolicy,
68}
69
70impl<T: Eq + Hash + Clone + Ord> DictionaryEncoder<T> {
71    /// Create an empty encoder with the given unknown policy.
72    pub fn new(unknown_policy: UnknownPolicy) -> Self {
73        Self {
74            vocab: Vec::new(),
75            code_map: HashMap::new(),
76            unknown_policy,
77        }
78    }
79
80    /// Build a `DictionaryEncoder` by scanning `data` and counting frequencies.
81    ///
82    /// The vocabulary is sorted by descending frequency (ties broken by
83    /// value ordering).  Codes `0..vocab.len()` are assigned in that order.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error if the number of unique values exceeds `u16::MAX - 1`
88    /// (= 65534).
89    pub fn fit(data: &[T], unknown_policy: UnknownPolicy) -> IoResult<Self> {
90        let mut freq: HashMap<T, u64> = HashMap::new();
91        for item in data {
92            *freq.entry(item.clone()).or_insert(0) += 1;
93        }
94
95        const MAX_VOCAB: usize = u16::MAX as usize - 1; // 65534
96        if freq.len() > MAX_VOCAB {
97            return Err(IoError::FormatError(format!(
98                "dictionary encoding: too many unique values ({} > {MAX_VOCAB})",
99                freq.len()
100            )));
101        }
102
103        // Sort by descending frequency; break ties by ascending value order
104        let mut entries: Vec<(T, u64)> = freq.into_iter().collect();
105        entries.sort_by(|(va, fa), (vb, fb)| fb.cmp(fa).then_with(|| va.cmp(vb)));
106
107        let mut vocab = Vec::with_capacity(entries.len());
108        let mut code_map: HashMap<T, u16> = HashMap::with_capacity(entries.len());
109
110        for (code, (value, _freq)) in entries.into_iter().enumerate() {
111            vocab.push(value.clone());
112            code_map.insert(value, code as u16);
113        }
114
115        Ok(Self {
116            vocab,
117            code_map,
118            unknown_policy,
119        })
120    }
121
122    /// Encode a slice of values to codes.
123    ///
124    /// Behaviour for unknown values is governed by `self.unknown_policy`.
125    pub fn encode(&self, data: &[T]) -> IoResult<Vec<u16>>
126    where
127        T: std::fmt::Debug,
128    {
129        let mut codes = Vec::with_capacity(data.len());
130        for item in data {
131            let code = self.encode_single(item)?;
132            codes.push(code);
133        }
134        Ok(codes)
135    }
136
137    /// Decode a slice of codes back to values.
138    ///
139    /// Returns an error if any code equals `u16::MAX` (null sentinel) or is
140    /// out of bounds.
141    pub fn decode(&self, codes: &[u16]) -> IoResult<Vec<T>> {
142        let mut out = Vec::with_capacity(codes.len());
143        for &code in codes {
144            out.push(self.decode_single(code)?.clone());
145        }
146        Ok(out)
147    }
148
149    /// Return the current vocabulary size.
150    pub fn vocab_size(&self) -> usize {
151        self.vocab.len()
152    }
153
154    /// Encode a single value to its code.
155    pub fn encode_single(&self, value: &T) -> IoResult<u16>
156    where
157        T: std::fmt::Debug,
158    {
159        match self.code_map.get(value) {
160            Some(&code) => Ok(code),
161            None => match self.unknown_policy {
162                UnknownPolicy::Error => Err(IoError::FormatError(format!(
163                    "dictionary encoding: unknown value {:?}",
164                    value
165                ))),
166                UnknownPolicy::UseNullCode => Ok(u16::MAX),
167                UnknownPolicy::AddToVocab => Err(IoError::FormatError(
168                    "dictionary encoding: AddToVocab requires a mutable encoder; \
169                         use encode_single_mut instead"
170                        .to_string(),
171                )),
172            },
173        }
174    }
175
176    /// Encode a single value, potentially adding it to the vocabulary
177    /// (only meaningful when `unknown_policy == AddToVocab`).
178    pub fn encode_single_mut(&mut self, value: T) -> IoResult<u16>
179    where
180        T: std::fmt::Debug,
181    {
182        if let Some(&code) = self.code_map.get(&value) {
183            return Ok(code);
184        }
185
186        match self.unknown_policy {
187            UnknownPolicy::Error => Err(IoError::FormatError(format!(
188                "dictionary encoding: unknown value {:?}",
189                value
190            ))),
191            UnknownPolicy::UseNullCode => Ok(u16::MAX),
192            UnknownPolicy::AddToVocab => {
193                const MAX_VOCAB: usize = u16::MAX as usize - 1;
194                if self.vocab.len() >= MAX_VOCAB {
195                    return Err(IoError::FormatError(
196                        "dictionary encoding: vocabulary capacity exhausted".to_string(),
197                    ));
198                }
199                let code = self.vocab.len() as u16;
200                self.vocab.push(value.clone());
201                self.code_map.insert(value, code);
202                Ok(code)
203            }
204        }
205    }
206
207    /// Decode a single code to a reference to the corresponding value.
208    pub fn decode_single(&self, code: u16) -> IoResult<&T> {
209        if code == u16::MAX {
210            return Err(IoError::FormatError(
211                "dictionary encoding: null sentinel code (u16::MAX) encountered during decode"
212                    .to_string(),
213            ));
214        }
215        self.vocab.get(code as usize).ok_or_else(|| {
216            IoError::FormatError(format!(
217                "dictionary encoding: code {} out of range (vocab size {})",
218                code,
219                self.vocab.len()
220            ))
221        })
222    }
223}
224
225/// Convenience type alias for string dictionary encoding.
226pub type StringDictionaryEncoder = DictionaryEncoder<String>;
227
228// ---------------------------------------------------------------------------
229// IoResult alias (for convenience)
230// ---------------------------------------------------------------------------
231
232type IoResult<T> = crate::error::Result<T>;
233
234// ---------------------------------------------------------------------------
235// Tests
236// ---------------------------------------------------------------------------
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn test_dict_encoder_fit_all_unique() {
244        let data: Vec<String> = (0..100).map(|i| format!("val_{i}")).collect();
245        let enc = DictionaryEncoder::fit(&data, UnknownPolicy::Error).expect("fit failed");
246        assert_eq!(enc.vocab_size(), 100);
247    }
248
249    #[test]
250    fn test_dict_encoder_encode_decode_roundtrip() {
251        let data: Vec<String> = vec!["a", "b", "a", "c", "b", "a"]
252            .into_iter()
253            .map(|s| s.to_string())
254            .collect();
255
256        let enc = DictionaryEncoder::fit(&data, UnknownPolicy::Error).expect("fit failed");
257
258        // "a" has freq 3 → code 0; "b" freq 2 → code 1; "c" freq 1 → code 2
259        assert_eq!(enc.vocab[0], "a");
260
261        let codes = enc.encode(&data).expect("encode failed");
262        assert_eq!(codes.len(), 6);
263
264        let decoded = enc.decode(&codes).expect("decode failed");
265        assert_eq!(decoded, data);
266    }
267
268    #[test]
269    fn test_dict_encoder_unknown_error_policy() {
270        let train: Vec<String> = vec!["x", "y"].into_iter().map(|s| s.to_string()).collect();
271        let enc = DictionaryEncoder::fit(&train, UnknownPolicy::Error).expect("fit failed");
272
273        let unknown = vec!["z".to_string()];
274        let result = enc.encode(&unknown);
275        assert!(result.is_err(), "expected error for unknown value");
276    }
277
278    #[test]
279    fn test_dict_encoder_null_code_policy() {
280        let train: Vec<String> = vec!["a".to_string()];
281        let enc = DictionaryEncoder::fit(&train, UnknownPolicy::UseNullCode).expect("fit failed");
282
283        let data = vec!["a".to_string(), "unknown".to_string()];
284        let codes = enc.encode(&data).expect("encode failed");
285        assert_eq!(codes[0], 0);
286        assert_eq!(codes[1], u16::MAX);
287    }
288
289    #[test]
290    fn test_dict_encoder_add_to_vocab_policy() {
291        let train: Vec<String> = vec!["a".to_string()];
292        let mut enc =
293            DictionaryEncoder::fit(&train, UnknownPolicy::AddToVocab).expect("fit failed");
294
295        let code = enc
296            .encode_single_mut("b".to_string())
297            .expect("encode failed");
298        assert_eq!(enc.vocab_size(), 2);
299        assert_eq!(enc.vocab[code as usize], "b");
300    }
301
302    #[test]
303    fn test_dict_encoder_decode_null_sentinel_error() {
304        let train: Vec<String> = vec!["a".to_string()];
305        let enc = DictionaryEncoder::fit(&train, UnknownPolicy::Error).expect("fit failed");
306
307        let result = enc.decode(&[u16::MAX]);
308        assert!(result.is_err(), "decoding null sentinel should error");
309    }
310
311    #[test]
312    fn test_dict_encoder_too_many_unique() {
313        // Build a dataset with exactly 65535 unique strings (exceeds limit)
314        let data: Vec<String> = (0u32..65535).map(|i| format!("v{i}")).collect();
315        let result = DictionaryEncoder::fit(&data, UnknownPolicy::Error);
316        assert!(result.is_err(), "should error on > 65534 unique values");
317    }
318}
scirs2_io/columnar/dictionary.rs

scirs2_io/columnar/
dictionary.rs
