polars_utils/idx_map/
bytes_idx_map.rs

1use hashbrown::hash_table::{
2    Entry as TEntry, HashTable, OccupiedEntry as TOccupiedEntry, VacantEntry as TVacantEntry,
3};
4
5use crate::IdxSize;
6
7const BASE_KEY_DATA_CAPACITY: usize = 1024;
8
9struct Key {
10    key_hash: u64,
11    key_buffer: u32,
12    key_offset: usize,
13    key_length: u32,
14}
15
16impl Key {
17    #[inline]
18    unsafe fn get<'k>(&self, key_data: &'k [Vec<u8>]) -> &'k [u8] {
19        let buf = unsafe { key_data.get_unchecked(self.key_buffer as usize) };
20        unsafe { buf.get_unchecked(self.key_offset..self.key_offset + self.key_length as usize) }
21    }
22}
23
24/// An IndexMap where the keys are always [u8] slices which are pre-hashed.
25pub struct BytesIndexMap<V> {
26    table: HashTable<IdxSize>,
27    tuples: Vec<(Key, V)>,
28    key_data: Vec<Vec<u8>>,
29
30    // Internal random seed used to keep hash iteration order decorrelated.
31    // We simply store a random odd number and multiply the canonical hash by it.
32    seed: u64,
33}
34
35impl<V> Default for BytesIndexMap<V> {
36    fn default() -> Self {
37        Self {
38            table: HashTable::new(),
39            tuples: Vec::new(),
40            key_data: vec![Vec::with_capacity(BASE_KEY_DATA_CAPACITY)],
41            seed: rand::random::<u64>() | 1,
42        }
43    }
44}
45
46impl<V> BytesIndexMap<V> {
47    pub fn new() -> Self {
48        Self::default()
49    }
50
51    pub fn reserve(&mut self, additional: usize) {
52        self.table.reserve(additional, |i| unsafe {
53            let tuple = self.tuples.get_unchecked(*i as usize);
54            tuple.0.key_hash.wrapping_mul(self.seed)
55        });
56        self.tuples.reserve(additional);
57    }
58
59    #[inline]
60    pub fn len(&self) -> IdxSize {
61        self.tuples.len() as IdxSize
62    }
63
64    #[inline]
65    pub fn is_empty(&self) -> bool {
66        self.tuples.is_empty()
67    }
68
69    pub fn get(&self, hash: u64, key: &[u8]) -> Option<&V> {
70        let idx = self.get_index_of(hash, key)?;
71        unsafe { Some(&self.tuples.get_unchecked(idx as usize).1) }
72    }
73
74    pub fn contains_key(&self, hash: u64, key: &[u8]) -> bool {
75        self.get_index_of(hash, key).is_some()
76    }
77
78    /// Gets the index by insertion order of the given key.
79    pub fn get_index_of(&self, hash: u64, key: &[u8]) -> Option<IdxSize> {
80        self.table
81            .find(hash.wrapping_mul(self.seed), |i| unsafe {
82                let t = self.tuples.get_unchecked(*i as usize);
83                hash == t.0.key_hash && key == t.0.get(&self.key_data)
84            })
85            .copied()
86    }
87
88    pub fn entry<'k>(&mut self, hash: u64, key: &'k [u8]) -> Entry<'_, 'k, V> {
89        let entry = self.table.entry(
90            hash.wrapping_mul(self.seed),
91            |i| unsafe {
92                let t = self.tuples.get_unchecked(*i as usize);
93                hash == t.0.key_hash && key == t.0.get(&self.key_data)
94            },
95            |i| unsafe {
96                let t = self.tuples.get_unchecked(*i as usize);
97                t.0.key_hash.wrapping_mul(self.seed)
98            },
99        );
100
101        match entry {
102            TEntry::Occupied(o) => Entry::Occupied(OccupiedEntry {
103                entry: o,
104                tuples: &mut self.tuples,
105            }),
106            TEntry::Vacant(v) => Entry::Vacant(VacantEntry {
107                key,
108                hash,
109                entry: v,
110                tuples: &mut self.tuples,
111                key_data: &mut self.key_data,
112            }),
113        }
114    }
115
116    /// Gets the hash, key and value at the given index by insertion order.
117    #[inline(always)]
118    pub fn get_index(&self, idx: IdxSize) -> Option<(u64, &[u8], &V)> {
119        let t = self.tuples.get(idx as usize)?;
120        Some((t.0.key_hash, unsafe { t.0.get(&self.key_data) }, &t.1))
121    }
122
123    /// Gets the hash, key and value at the given index by insertion order.
124    ///
125    /// # Safety
126    /// The index must be less than len().
127    #[inline(always)]
128    pub unsafe fn get_index_unchecked(&self, idx: IdxSize) -> (u64, &[u8], &V) {
129        let t = unsafe { self.tuples.get_unchecked(idx as usize) };
130        unsafe { (t.0.key_hash, t.0.get(&self.key_data), &t.1) }
131    }
132
133    /// Iterates over the (hash, key) pairs in insertion order, where each key slice runs to the
134    /// end of the buffer that holds it.
135    pub fn iter_hash_keys_to_buffer_end(&self) -> impl Iterator<Item = (u64, &[u8])> {
136        self.tuples.iter().map(|t| unsafe {
137            let buf = self.key_data.get_unchecked(t.0.key_buffer as usize);
138            (t.0.key_hash, buf.get_unchecked(t.0.key_offset..))
139        })
140    }
141
142    /// Iterates over the values in insertion order.
143    pub fn iter_values(&self) -> impl Iterator<Item = &V> {
144        self.tuples.iter().map(|t| &t.1)
145    }
146}
147
148pub enum Entry<'a, 'k, V> {
149    Occupied(OccupiedEntry<'a, V>),
150    Vacant(VacantEntry<'a, 'k, V>),
151}
152
153pub struct OccupiedEntry<'a, V> {
154    entry: TOccupiedEntry<'a, IdxSize>,
155    tuples: &'a mut Vec<(Key, V)>,
156}
157
158impl<'a, V> OccupiedEntry<'a, V> {
159    pub fn index(&self) -> IdxSize {
160        *self.entry.get()
161    }
162
163    pub fn into_mut(self) -> &'a mut V {
164        let idx = self.index();
165        unsafe { &mut self.tuples.get_unchecked_mut(idx as usize).1 }
166    }
167}
168
169pub struct VacantEntry<'a, 'k, V> {
170    hash: u64,
171    key: &'k [u8],
172    entry: TVacantEntry<'a, IdxSize>,
173    tuples: &'a mut Vec<(Key, V)>,
174    key_data: &'a mut Vec<Vec<u8>>,
175}
176
177#[allow(clippy::needless_lifetimes)]
178impl<'a, 'k, V> VacantEntry<'a, 'k, V> {
179    pub fn index(&self) -> IdxSize {
180        self.tuples.len() as IdxSize
181    }
182
183    pub fn insert(self, value: V) -> &'a mut V {
184        unsafe {
185            let tuple_idx: IdxSize = self.tuples.len().try_into().unwrap();
186
187            let mut num_buffers = self.key_data.len() as u32;
188            let mut active_buf = self.key_data.last_mut().unwrap_unchecked();
189            let key_len = self.key.len();
190            if active_buf.len() + key_len > active_buf.capacity() {
191                let ideal_next_cap = BASE_KEY_DATA_CAPACITY.checked_shl(num_buffers).unwrap();
192                let next_capacity = std::cmp::max(ideal_next_cap, key_len);
193                self.key_data.push(Vec::with_capacity(next_capacity));
194                active_buf = self.key_data.last_mut().unwrap_unchecked();
195                num_buffers += 1;
196            }
197
198            let tuple_key = Key {
199                key_hash: self.hash,
200                key_buffer: num_buffers - 1,
201                key_offset: active_buf.len(),
202                key_length: self.key.len().try_into().unwrap(),
203            };
204            self.tuples.push((tuple_key, value));
205            active_buf.extend_from_slice(self.key);
206            self.entry.insert(tuple_idx);
207            &mut self.tuples.last_mut().unwrap_unchecked().1
208        }
209    }
210}
polars_utils/idx_map/bytes_idx_map.rs

polars_utils/idx_map/
bytes_idx_map.rs
