polars_utils/
pl_str.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2
3#[macro_export]
4macro_rules! format_pl_smallstr {
5    ($($arg:tt)*) => {{
6        use std::fmt::Write;
7
8        let mut string = $crate::pl_str::PlSmallStr::EMPTY;
9        write!(string, $($arg)*).unwrap();
10        string
11    }}
12}
13
14type Inner = compact_str::CompactString;
15
16/// String type that inlines small strings.
17#[derive(Clone, Eq, Hash, PartialOrd, Ord)]
18#[cfg_attr(
19    feature = "serde",
20    derive(serde::Serialize, serde::Deserialize),
21    serde(transparent)
22)]
23pub struct PlSmallStr(Inner);
24
25impl PlSmallStr {
26    pub const EMPTY: Self = Self::from_static("");
27    pub const EMPTY_REF: &'static Self = &Self::from_static("");
28
29    #[inline(always)]
30    pub const fn from_static(s: &'static str) -> Self {
31        Self(Inner::const_new(s))
32    }
33
34    #[inline(always)]
35    #[allow(clippy::should_implement_trait)]
36    pub fn from_str(s: &str) -> Self {
37        Self(Inner::from(s))
38    }
39
40    #[inline(always)]
41    pub fn from_string(s: String) -> Self {
42        Self(Inner::from(s))
43    }
44
45    #[inline(always)]
46    pub fn as_str(&self) -> &str {
47        self.0.as_str()
48    }
49
50    #[inline(always)]
51    pub fn as_mut_str(&mut self) -> &mut str {
52        self.0.as_mut_str()
53    }
54
55    #[inline(always)]
56    pub fn into_string(self) -> String {
57        self.0.into_string()
58    }
59}
60
61impl Default for PlSmallStr {
62    #[inline(always)]
63    fn default() -> Self {
64        Self::EMPTY
65    }
66}
67
68// AsRef, Deref and Borrow impls to &str
69
70impl AsRef<str> for PlSmallStr {
71    #[inline(always)]
72    fn as_ref(&self) -> &str {
73        self.as_str()
74    }
75}
76
77impl core::ops::Deref for PlSmallStr {
78    type Target = str;
79
80    #[inline(always)]
81    fn deref(&self) -> &Self::Target {
82        self.as_str()
83    }
84}
85
86impl core::ops::DerefMut for PlSmallStr {
87    #[inline(always)]
88    fn deref_mut(&mut self) -> &mut Self::Target {
89        self.as_mut_str()
90    }
91}
92
93impl core::borrow::Borrow<str> for PlSmallStr {
94    #[inline(always)]
95    fn borrow(&self) -> &str {
96        self.as_str()
97    }
98}
99
100// AsRef impls for other types
101
102impl AsRef<std::path::Path> for PlSmallStr {
103    #[inline(always)]
104    fn as_ref(&self) -> &std::path::Path {
105        self.as_str().as_ref()
106    }
107}
108
109impl AsRef<[u8]> for PlSmallStr {
110    #[inline(always)]
111    fn as_ref(&self) -> &[u8] {
112        self.as_str().as_bytes()
113    }
114}
115
116impl AsRef<std::ffi::OsStr> for PlSmallStr {
117    #[inline(always)]
118    fn as_ref(&self) -> &std::ffi::OsStr {
119        self.as_str().as_ref()
120    }
121}
122
123// From impls
124
125impl From<&str> for PlSmallStr {
126    #[inline(always)]
127    fn from(value: &str) -> Self {
128        Self::from_str(value)
129    }
130}
131
132impl From<String> for PlSmallStr {
133    #[inline(always)]
134    fn from(value: String) -> Self {
135        Self::from_string(value)
136    }
137}
138
139impl From<&String> for PlSmallStr {
140    #[inline(always)]
141    fn from(value: &String) -> Self {
142        Self::from_str(value.as_str())
143    }
144}
145
146impl From<Inner> for PlSmallStr {
147    #[inline(always)]
148    fn from(value: Inner) -> Self {
149        Self(value)
150    }
151}
152
153// FromIterator impls
154
155impl FromIterator<PlSmallStr> for PlSmallStr {
156    #[inline(always)]
157    fn from_iter<T: IntoIterator<Item = PlSmallStr>>(iter: T) -> Self {
158        Self(Inner::from_iter(iter.into_iter().map(|x| x.0)))
159    }
160}
161
162impl<'a> FromIterator<&'a PlSmallStr> for PlSmallStr {
163    #[inline(always)]
164    fn from_iter<T: IntoIterator<Item = &'a PlSmallStr>>(iter: T) -> Self {
165        Self(Inner::from_iter(iter.into_iter().map(|x| x.as_str())))
166    }
167}
168
169impl FromIterator<char> for PlSmallStr {
170    #[inline(always)]
171    fn from_iter<I: IntoIterator<Item = char>>(iter: I) -> PlSmallStr {
172        Self(Inner::from_iter(iter))
173    }
174}
175
176impl<'a> FromIterator<&'a char> for PlSmallStr {
177    #[inline(always)]
178    fn from_iter<I: IntoIterator<Item = &'a char>>(iter: I) -> PlSmallStr {
179        Self(Inner::from_iter(iter))
180    }
181}
182
183impl<'a> FromIterator<&'a str> for PlSmallStr {
184    #[inline(always)]
185    fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> PlSmallStr {
186        Self(Inner::from_iter(iter))
187    }
188}
189
190impl FromIterator<String> for PlSmallStr {
191    #[inline(always)]
192    fn from_iter<I: IntoIterator<Item = String>>(iter: I) -> PlSmallStr {
193        Self(Inner::from_iter(iter))
194    }
195}
196
197impl FromIterator<Box<str>> for PlSmallStr {
198    #[inline(always)]
199    fn from_iter<I: IntoIterator<Item = Box<str>>>(iter: I) -> PlSmallStr {
200        Self(Inner::from_iter(iter))
201    }
202}
203
204impl<'a> FromIterator<std::borrow::Cow<'a, str>> for PlSmallStr {
205    #[inline(always)]
206    fn from_iter<I: IntoIterator<Item = std::borrow::Cow<'a, str>>>(iter: I) -> PlSmallStr {
207        Self(Inner::from_iter(iter))
208    }
209}
210
211// PartialEq impls
212
213impl<T> PartialEq<T> for PlSmallStr
214where
215    T: AsRef<str> + ?Sized,
216{
217    #[inline(always)]
218    fn eq(&self, other: &T) -> bool {
219        self.as_str() == other.as_ref()
220    }
221}
222
223impl PartialEq<PlSmallStr> for &str {
224    #[inline(always)]
225    fn eq(&self, other: &PlSmallStr) -> bool {
226        *self == other.as_str()
227    }
228}
229
230impl PartialEq<PlSmallStr> for String {
231    #[inline(always)]
232    fn eq(&self, other: &PlSmallStr) -> bool {
233        self.as_str() == other.as_str()
234    }
235}
236
237// Write
238
239impl core::fmt::Write for PlSmallStr {
240    #[inline(always)]
241    fn write_char(&mut self, c: char) -> std::fmt::Result {
242        self.0.write_char(c)
243    }
244
245    #[inline(always)]
246    fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::fmt::Result {
247        self.0.write_fmt(args)
248    }
249
250    #[inline(always)]
251    fn write_str(&mut self, s: &str) -> std::fmt::Result {
252        self.0.write_str(s)
253    }
254}
255
256// Debug, Display
257
258impl core::fmt::Debug for PlSmallStr {
259    #[inline(always)]
260    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261        self.as_str().fmt(f)
262    }
263}
264
265impl core::fmt::Display for PlSmallStr {
266    #[inline(always)]
267    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
268        self.as_str().fmt(f)
269    }
270}
271
272pub fn unique_column_name() -> PlSmallStr {
273    static COUNTER: AtomicU64 = AtomicU64::new(0);
274    let idx = COUNTER.fetch_add(1, Ordering::Relaxed);
275    format_pl_smallstr!("_POLARS_TMP_{idx}")
276}