polars_time/chunkedarray/string/
mod.rs

1pub mod infer;
2use chrono::DateTime;
3mod patterns;
4mod strptime;
5pub use patterns::Pattern;
6#[cfg(feature = "dtype-time")]
7use polars_core::chunked_array::temporal::time_to_time64ns;
8use polars_core::prelude::arity::unary_elementwise;
9use polars_utils::cache::LruCachedFunc;
10
11use super::*;
12#[cfg(feature = "dtype-date")]
13use crate::chunkedarray::date::naive_date_to_date;
14use crate::prelude::string::strptime::StrpTimeState;
15
16#[cfg(feature = "dtype-time")]
17fn time_pattern<F, K>(val: &str, convert: F) -> Option<&'static str>
18// (string, fmt) -> PolarsResult
19where
20    F: Fn(&str, &str) -> chrono::ParseResult<K>,
21{
22    patterns::TIME_H_M_S
23        .iter()
24        .chain(patterns::TIME_H_M_S)
25        .find(|fmt| convert(val, fmt).is_ok())
26        .copied()
27}
28
29fn datetime_pattern<F, K>(val: &str, convert: F) -> Option<&'static str>
30// (string, fmt) -> PolarsResult
31where
32    F: Fn(&str, &str) -> chrono::ParseResult<K>,
33{
34    patterns::DATETIME_Y_M_D
35        .iter()
36        .chain(patterns::DATETIME_D_M_Y)
37        .find(|fmt| convert(val, fmt).is_ok())
38        .copied()
39}
40
41fn date_pattern<F, K>(val: &str, convert: F) -> Option<&'static str>
42// (string, fmt) -> PolarsResult
43where
44    F: Fn(&str, &str) -> chrono::ParseResult<K>,
45{
46    patterns::DATE_Y_M_D
47        .iter()
48        .chain(patterns::DATE_D_M_Y)
49        .find(|fmt| convert(val, fmt).is_ok())
50        .copied()
51}
52
53pub trait StringMethods: AsString {
54    #[cfg(feature = "dtype-time")]
55    /// Parsing string values and return a [`TimeChunked`]
56    fn as_time(&self, fmt: Option<&str>, use_cache: bool) -> PolarsResult<TimeChunked> {
57        let string_ca = self.as_string();
58        let fmt = match fmt {
59            Some(fmt) => fmt,
60            None => {
61                if string_ca.null_count() == string_ca.len() {
62                    return Ok(
63                        Int64Chunked::full_null(string_ca.name().clone(), string_ca.len())
64                            .into_time(),
65                    );
66                }
67                infer::infer_from_values(string_ca, |val| {
68                    time_pattern(val, NaiveTime::parse_from_str)
69                })
70                .ok_or_else(|| polars_err!(parse_fmt_idk = "time"))?
71            },
72        };
73        let use_cache = use_cache && string_ca.len() > 50;
74
75        let mut convert = LruCachedFunc::new(
76            |s| {
77                let naive_time = NaiveTime::parse_from_str(s, fmt).ok()?;
78                Some(time_to_time64ns(&naive_time))
79            },
80            (string_ca.len() as f64).sqrt() as usize,
81        );
82        let ca = unary_elementwise(string_ca, |opt_s| convert.eval(opt_s?, use_cache));
83        Ok(ca.with_name(string_ca.name().clone()).into_time())
84    }
85
86    #[cfg(feature = "dtype-date")]
87    /// Parsing string values and return a [`DateChunked`]
88    /// Different from `as_date` this function allows matches that not contain the whole string
89    /// e.g. "foo-2021-01-01-bar" could match "2021-01-01"
90    fn as_date_not_exact(&self, fmt: Option<&str>) -> PolarsResult<DateChunked> {
91        let string_ca = self.as_string();
92        let fmt = match fmt {
93            Some(fmt) => fmt,
94            None => {
95                if string_ca.null_count() == string_ca.len() {
96                    return Ok(
97                        Int32Chunked::full_null(string_ca.name().clone(), string_ca.len())
98                            .into_date(),
99                    );
100                }
101                infer::infer_from_values(string_ca, |val| {
102                    date_pattern(val, NaiveDate::parse_from_str)
103                })
104                .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?
105            },
106        };
107        let ca = unary_elementwise(string_ca, |opt_s| {
108            let mut s = opt_s?;
109            while !s.is_empty() {
110                match NaiveDate::parse_and_remainder(s, fmt) {
111                    Ok((nd, _)) => return Some(naive_date_to_date(nd)),
112                    Err(_) => {
113                        let mut it = s.chars();
114                        it.next();
115                        s = it.as_str();
116                    },
117                }
118            }
119
120            None
121        });
122        Ok(ca.with_name(string_ca.name().clone()).into_date())
123    }
124
125    #[cfg(feature = "dtype-datetime")]
126    /// Parsing string values and return a [`DatetimeChunked`]
127    /// Different from `as_datetime` this function allows matches that not contain the whole string
128    /// e.g. "foo-2021-01-01-bar" could match "2021-01-01"
129    fn as_datetime_not_exact(
130        &self,
131        fmt: Option<&str>,
132        tu: TimeUnit,
133        tz_aware: bool,
134        tz: Option<&TimeZone>,
135        _ambiguous: &StringChunked,
136        // Ensure that the inferred time_zone matches the given time_zone.
137        ensure_matching_tz: bool,
138    ) -> PolarsResult<DatetimeChunked> {
139        let string_ca = self.as_string();
140        let had_format = fmt.is_some();
141        let fmt = match fmt {
142            Some(fmt) => fmt,
143            None => {
144                if string_ca.null_count() == string_ca.len() {
145                    return Ok(
146                        Int64Chunked::full_null(string_ca.name().clone(), string_ca.len())
147                            .into_datetime(tu, tz.cloned()),
148                    );
149                }
150                infer::infer_from_values(string_ca, |val| {
151                    datetime_pattern(val, NaiveDateTime::parse_from_str)
152                        .or_else(|| datetime_pattern(val, NaiveDate::parse_from_str))
153                })
154                .ok_or_else(|| polars_err!(parse_fmt_idk = "datetime"))?
155            },
156        };
157
158        let func = match tu {
159            TimeUnit::Nanoseconds => datetime_to_timestamp_ns,
160            TimeUnit::Microseconds => datetime_to_timestamp_us,
161            TimeUnit::Milliseconds => datetime_to_timestamp_ms,
162        };
163
164        let ca = unary_elementwise(string_ca, |opt_s| {
165            let mut s = opt_s?;
166            while !s.is_empty() {
167                let timestamp = if tz_aware {
168                    DateTime::parse_and_remainder(s, fmt)
169                        .ok()
170                        .map(|(dt, _r)| func(dt.naive_utc()))
171                } else {
172                    infer::parse_datetime_and_remainder(s, fmt).map(|(nd, _r)| func(nd))
173                };
174                match timestamp {
175                    Some(ts) => return Some(ts),
176                    None => {
177                        let mut it = s.chars();
178                        it.next();
179                        s = it.as_str();
180                    },
181                }
182            }
183            None
184        })
185        .with_name(string_ca.name().clone());
186
187        polars_ensure!(
188            !ensure_matching_tz || had_format || !(tz_aware && tz.is_none()),
189            to_datetime_tz_mismatch
190        );
191
192        match (tz_aware, tz) {
193            #[cfg(feature = "timezones")]
194            (false, Some(tz)) => polars_ops::prelude::replace_time_zone(
195                &ca.into_datetime(tu, None),
196                Some(tz),
197                _ambiguous,
198                NonExistent::Raise,
199            ),
200            #[cfg(feature = "timezones")]
201            (true, tz) => Ok(ca.into_datetime(tu, Some(tz.cloned().unwrap_or(TimeZone::UTC)))),
202            _ => Ok(ca.into_datetime(tu, None)),
203        }
204    }
205
206    #[cfg(feature = "dtype-date")]
207    /// Parsing string values and return a [`DateChunked`]
208    fn as_date(&self, fmt: Option<&str>, use_cache: bool) -> PolarsResult<DateChunked> {
209        let string_ca = self.as_string();
210        let fmt = match fmt {
211            Some(fmt) => fmt,
212            None => return infer::to_date(string_ca),
213        };
214        let use_cache = use_cache && string_ca.len() > 50;
215        let fmt = strptime::compile_fmt(fmt)?;
216
217        // We can use the fast parser.
218        let ca = if strptime::fast_parser_supported(fmt.as_bytes()) {
219            let mut strptime_cache = StrpTimeState::default();
220            let mut convert = LruCachedFunc::new(
221                |s: &str| {
222                    match strptime_cache.parse(s.as_bytes(), fmt.as_bytes()) {
223                        // Fallback to chrono.
224                        None => NaiveDate::parse_from_str(s, &fmt).ok(),
225                        Some(ndt) => Some(ndt.date()),
226                    }
227                    .map(naive_date_to_date)
228                },
229                (string_ca.len() as f64).sqrt() as usize,
230            );
231            unary_elementwise(string_ca, |val| convert.eval(val?, use_cache))
232        } else {
233            let mut convert = LruCachedFunc::new(
234                |s| {
235                    let naive_date = NaiveDate::parse_from_str(s, &fmt).ok()?;
236                    Some(naive_date_to_date(naive_date))
237                },
238                (string_ca.len() as f64).sqrt() as usize,
239            );
240            unary_elementwise(string_ca, |val| convert.eval(val?, use_cache))
241        };
242
243        Ok(ca.with_name(string_ca.name().clone()).into_date())
244    }
245
246    #[cfg(feature = "dtype-datetime")]
247    /// Parsing string values and return a [`DatetimeChunked`].
248    fn as_datetime(
249        &self,
250        fmt: Option<&str>,
251        tu: TimeUnit,
252        use_cache: bool,
253        tz_aware: bool,
254        tz: Option<&TimeZone>,
255        ambiguous: &StringChunked,
256    ) -> PolarsResult<DatetimeChunked> {
257        let string_ca = self.as_string();
258        let fmt = match fmt {
259            Some(fmt) => fmt,
260            None => return infer::to_datetime(string_ca, tu, tz, ambiguous, true),
261        };
262        let fmt = strptime::compile_fmt(fmt)?;
263        let use_cache = use_cache && string_ca.len() > 50;
264
265        let func = match tu {
266            TimeUnit::Nanoseconds => datetime_to_timestamp_ns,
267            TimeUnit::Microseconds => datetime_to_timestamp_us,
268            TimeUnit::Milliseconds => datetime_to_timestamp_ms,
269        };
270
271        if tz_aware {
272            #[cfg(feature = "timezones")]
273            {
274                let mut convert = LruCachedFunc::new(
275                    |s: &str| {
276                        let dt = DateTime::parse_from_str(s, &fmt).ok()?;
277                        Some(func(dt.naive_utc()))
278                    },
279                    (string_ca.len() as f64).sqrt() as usize,
280                );
281                Ok(
282                    unary_elementwise(string_ca, |opt_s| convert.eval(opt_s?, use_cache))
283                        .with_name(string_ca.name().clone())
284                        .into_datetime(tu, Some(tz.cloned().unwrap_or(TimeZone::UTC))),
285                )
286            }
287            #[cfg(not(feature = "timezones"))]
288            {
289                panic!("activate 'timezones' feature")
290            }
291        } else {
292            let transform = match tu {
293                TimeUnit::Nanoseconds => infer::transform_datetime_ns,
294                TimeUnit::Microseconds => infer::transform_datetime_us,
295                TimeUnit::Milliseconds => infer::transform_datetime_ms,
296            };
297            let ca = if strptime::fast_parser_supported(fmt.as_bytes()) {
298                let mut strptime_cache = StrpTimeState::default();
299                let mut convert = LruCachedFunc::new(
300                    |s: &str| match strptime_cache.parse(s.as_bytes(), fmt.as_bytes()) {
301                        None => transform(s, &fmt),
302                        Some(ndt) => Some(func(ndt)),
303                    },
304                    (string_ca.len() as f64).sqrt() as usize,
305                );
306                unary_elementwise(string_ca, |opt_s| convert.eval(opt_s?, use_cache))
307            } else {
308                let mut convert = LruCachedFunc::new(
309                    |s| transform(s, &fmt),
310                    (string_ca.len() as f64).sqrt() as usize,
311                );
312                unary_elementwise(string_ca, |opt_s| convert.eval(opt_s?, use_cache))
313            };
314            let dt = ca
315                .with_name(string_ca.name().clone())
316                .into_datetime(tu, None);
317            match tz {
318                #[cfg(feature = "timezones")]
319                Some(tz) => polars_ops::prelude::replace_time_zone(
320                    &dt,
321                    Some(tz),
322                    ambiguous,
323                    NonExistent::Raise,
324                ),
325                _ => Ok(dt),
326            }
327        }
328    }
329}
330
331pub trait AsString {
332    fn as_string(&self) -> &StringChunked;
333}
334
335impl AsString for StringChunked {
336    fn as_string(&self) -> &StringChunked {
337        self
338    }
339}
340
341impl StringMethods for StringChunked {}
polars_time/chunkedarray/string/mod.rs

polars_time/chunkedarray/string/
mod.rs
