Skip to main content

nodedb_columnar/memtable/column_data/
dict_encode.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Dictionary encoding conversion for `ColumnData`.
4
5use super::types::ColumnData;
6
7impl ColumnData {
8    /// Attempt to convert a `String` column to `DictEncoded`.
9    pub fn try_dict_encode(col: &ColumnData, max_cardinality: u32) -> Option<ColumnData> {
10        let (data, offsets, valid) = match col {
11            ColumnData::String {
12                data,
13                offsets,
14                valid,
15            } => (data, offsets, valid),
16            _ => return None,
17        };
18
19        let row_count = col.len();
20        let mut dictionary: Vec<String> = Vec::new();
21        let mut reverse: std::collections::HashMap<String, u32> = std::collections::HashMap::new();
22        let mut ids: Vec<u32> = Vec::with_capacity(row_count);
23
24        for i in 0..row_count {
25            if valid.as_ref().is_some_and(|v| !v[i]) {
26                ids.push(0);
27                continue;
28            }
29            let start = offsets[i] as usize;
30            let end = offsets[i + 1] as usize;
31            let s = match std::str::from_utf8(&data[start..end]) {
32                Ok(s) => s,
33                Err(_) => return None,
34            };
35            let id = if let Some(&existing) = reverse.get(s) {
36                existing
37            } else {
38                if dictionary.len() as u32 >= max_cardinality {
39                    return None;
40                }
41                let new_id = dictionary.len() as u32;
42                dictionary.push(s.to_string());
43                reverse.insert(s.to_string(), new_id);
44                new_id
45            };
46            ids.push(id);
47        }
48
49        Some(ColumnData::DictEncoded {
50            ids,
51            dictionary,
52            reverse,
53            valid: valid.clone(),
54        })
55    }
56}