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