Skip to main content

polars_utils/
aliases.rs

1use std::hash::{Hash, Hasher};
2
3use foldhash::SharedSeed;
4
5pub type PlRandomState = foldhash::quality::RandomState;
6pub type PlSeedableRandomStateQuality = foldhash::quality::SeedableRandomState;
7pub type PlRandomStateQuality = foldhash::quality::RandomState;
8pub type PlFixedStateQuality = foldhash::quality::FixedState;
9
10pub type PlHashMap<K, V> = hashbrown::HashMap<K, V, PlRandomState>;
11pub type PlHashSet<V> = hashbrown::HashSet<V, PlRandomState>;
12pub type PlIndexMap<K, V> = indexmap::IndexMap<K, V, PlRandomState>;
13pub type PlIndexSet<K> = indexmap::IndexSet<K, PlRandomState>;
14
15/// HashMap container with a getter that clears the HashMap.
16#[derive(Default)]
17pub struct ScratchHashMap<K, V>(PlHashMap<K, V>);
18
19impl<K, V> ScratchHashMap<K, V> {
20    /// Clear the HashMap and return a mutable reference to it.
21    pub fn get(&mut self) -> &mut PlHashMap<K, V> {
22        self.0.clear();
23        &mut self.0
24    }
25}
26
27/// HashSet container with a getter that clears the HashSet.
28#[derive(Default)]
29pub struct ScratchHashSet<K>(PlHashSet<K>);
30
31impl<K> ScratchHashSet<K> {
32    /// Clear the HashSet and return a mutable reference to it.
33    pub fn get(&mut self) -> &mut PlHashSet<K> {
34        self.0.clear();
35        &mut self.0
36    }
37}
38
39#[derive(Default)]
40pub struct ScratchIndexSet<K>(PlIndexSet<K>);
41
42impl<K> ScratchIndexSet<K> {
43    /// Clear the IndexSet and return a mutable reference to it.
44    pub fn get(&mut self) -> &mut PlIndexSet<K> {
45        self.0.clear();
46        &mut self.0
47    }
48}
49
50#[derive(Default)]
51pub struct ScratchIndexMap<K, V>(PlIndexMap<K, V>);
52
53impl<K, V> ScratchIndexMap<K, V> {
54    /// Clear the IndexMap and return a mutable reference to it.
55    pub fn get(&mut self) -> &mut PlIndexMap<K, V> {
56        self.0.clear();
57        &mut self.0
58    }
59}
60
61pub trait SeedableFromU64SeedExt {
62    fn seed_from_u64(seed: u64) -> Self;
63}
64
65impl SeedableFromU64SeedExt for PlSeedableRandomStateQuality {
66    fn seed_from_u64(seed: u64) -> Self {
67        PlSeedableRandomStateQuality::with_seed(seed, SharedSeed::global_fixed())
68    }
69}
70
71pub trait InitHashMaps {
72    type HashMap;
73
74    fn new() -> Self::HashMap;
75
76    fn with_capacity(capacity: usize) -> Self::HashMap;
77}
78
79impl<K, V> InitHashMaps for PlHashMap<K, V> {
80    type HashMap = Self;
81
82    #[inline]
83    fn new() -> Self::HashMap {
84        Self::with_capacity_and_hasher(0, Default::default())
85    }
86
87    #[inline]
88    fn with_capacity(capacity: usize) -> Self {
89        Self::with_capacity_and_hasher(capacity, Default::default())
90    }
91}
92impl<K> InitHashMaps for PlHashSet<K> {
93    type HashMap = Self;
94
95    #[inline]
96    fn new() -> Self::HashMap {
97        Self::with_capacity_and_hasher(0, Default::default())
98    }
99
100    #[inline]
101    fn with_capacity(capacity: usize) -> Self {
102        Self::with_capacity_and_hasher(capacity, Default::default())
103    }
104}
105
106impl<K> InitHashMaps for PlIndexSet<K> {
107    type HashMap = Self;
108
109    #[inline]
110    fn new() -> Self::HashMap {
111        Self::with_capacity_and_hasher(0, Default::default())
112    }
113
114    #[inline]
115    fn with_capacity(capacity: usize) -> Self::HashMap {
116        Self::with_capacity_and_hasher(capacity, Default::default())
117    }
118}
119
120impl<K, V> InitHashMaps for PlIndexMap<K, V> {
121    type HashMap = Self;
122
123    #[inline]
124    fn new() -> Self::HashMap {
125        Self::with_capacity_and_hasher(0, Default::default())
126    }
127
128    #[inline]
129    fn with_capacity(capacity: usize) -> Self::HashMap {
130        Self::with_capacity_and_hasher(capacity, Default::default())
131    }
132}
133
134#[derive(Debug, Clone)]
135#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
136#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
137#[repr(transparent)]
138pub struct PlIndexMapHashable<K: Eq + Hash, V>(pub PlIndexMap<K, V>);
139
140impl<K: Eq + Hash + PartialEq, V: PartialEq> PartialEq for PlIndexMapHashable<K, V> {
141    fn eq(&self, other: &Self) -> bool {
142        PartialEq::eq(&self.0, &other.0)
143    }
144}
145
146impl<K: Eq + Hash + PartialEq, V: PartialEq> Eq for PlIndexMapHashable<K, V> {}
147
148impl<K: Eq + Hash, V> Hash for PlIndexMapHashable<K, V> {
149    fn hash<H: Hasher>(&self, state: &mut H) {
150        self.keys().len().hash(state);
151        for key in self.keys() {
152            key.hash(state);
153        }
154    }
155}
156
157impl<K: Eq + Hash, V> std::ops::Deref for PlIndexMapHashable<K, V> {
158    type Target = PlIndexMap<K, V>;
159
160    fn deref(&self) -> &Self::Target {
161        &self.0
162    }
163}
164
165impl<K: Eq + Hash, V> std::ops::DerefMut for PlIndexMapHashable<K, V> {
166    fn deref_mut(&mut self) -> &mut Self::Target {
167        &mut self.0
168    }
169}