polars_utils/
pl_str.rs

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