polars_time/chunkedarray/string/
infer.rs

1use arrow::array::PrimitiveArray;
2use chrono::format::ParseErrorKind;
3use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime};
4use polars_core::prelude::*;
5
6use super::patterns::{self, Pattern};
7#[cfg(feature = "dtype-date")]
8use crate::chunkedarray::date::naive_date_to_date;
9use crate::prelude::string::strptime::StrpTimeState;
10
11polars_utils::regex_cache::cached_regex! {
12    static DATETIME_DMY_RE = r#"(?x)
13        ^
14        ['"]?                        # optional quotes
15        (?:\d{1,2})                  # day
16        [-/\.]                       # separator
17        (?P<month>[01]?\d{1})        # month
18        [-/\.]                       # separator
19        (?:\d{4,})                   # year
20        (?:
21            [T\ ]                    # separator
22            (?:\d{1,2})              # hour
23            :?                       # separator
24            (?:\d{1,2})              # minute
25            (?:
26                :?                   # separator
27                (?:\d{1,2})          # second
28                (?:
29                    \.(?:\d{1,9})    # subsecond
30                )?
31            )?
32        )?
33        ['"]?                        # optional quotes
34        $
35        "#;
36
37    static DATETIME_YMD_RE = r#"(?x)
38            ^
39            ['"]?                      # optional quotes
40            (?:\d{4,})                 # year
41            [-/\.]?                    # separator
42            (?P<month>[01]?\d{1})      # month
43            [-/\.]?                    # separator
44            (?:\d{1,2})                # day
45            (?:
46                [T\ ]                  # separator
47                (?:\d{1,2})            # hour
48                :?                     # separator
49                (?:\d{1,2})            # minute
50                (?:
51                    :?                 # separator
52                    (?:\d{1,2})        # seconds
53                    (?:
54                        \.(?:\d{1,9})  # subsecond
55                    )?
56                )?
57            )?
58            ['"]?                      # optional quotes
59            $
60            "#;
61
62    static DATETIME_YMDZ_RE = r#"(?x)
63            ^
64            ['"]?                  # optional quotes
65            (?:\d{4,})             # year
66            [-/\.]?                # separator
67            (?P<month>[01]?\d{1})  # month
68            [-/\.]?                # separator
69            (?:\d{1,2})            # year
70            [T\ ]                  # separator
71            (?:\d{2})              # hour
72            :?                     # separator
73            (?:\d{2})              # minute
74            (?:
75                :?                 # separator
76                (?:\d{2})          # second
77                (?:
78                    \.(?:\d{1,9})  # subsecond
79                )?
80            )?
81            (?:
82                # offset (e.g. +01:00, +0100, or +01)
83                [+-](?:\d{2})
84                (?::?\d{2})?
85                # or Zulu suffix
86                |Z
87            )
88            ['"]?                  # optional quotes
89            $
90            "#;
91}
92
93impl Pattern {
94    pub fn is_inferable(&self, val: &str) -> bool {
95        match self {
96            Pattern::DateDMY => true, // there are very few Date patterns, so it's cheaper
97            Pattern::DateYMD => true, // to just try them
98            Pattern::Time => true,
99            Pattern::DatetimeDMY => match DATETIME_DMY_RE.captures(val) {
100                Some(search) => (1..=12).contains(
101                    &search
102                        .name("month")
103                        .unwrap()
104                        .as_str()
105                        .parse::<u8>()
106                        .unwrap(),
107                ),
108                None => false,
109            },
110            Pattern::DatetimeYMD => match DATETIME_YMD_RE.captures(val) {
111                Some(search) => (1..=12).contains(
112                    &search
113                        .name("month")
114                        .unwrap()
115                        .as_str()
116                        .parse::<u8>()
117                        .unwrap(),
118                ),
119                None => false,
120            },
121            Pattern::DatetimeYMDZ => match DATETIME_YMDZ_RE.captures(val) {
122                Some(search) => (1..=12).contains(
123                    &search
124                        .name("month")
125                        .unwrap()
126                        .as_str()
127                        .parse::<u8>()
128                        .unwrap(),
129                ),
130                None => false,
131            },
132        }
133    }
134}
135
136pub trait StrpTimeParser<T> {
137    fn parse_bytes(&mut self, val: &[u8], time_unit: Option<TimeUnit>) -> Option<T>;
138}
139
140#[cfg(feature = "dtype-datetime")]
141impl StrpTimeParser<i64> for DatetimeInfer<Int64Type> {
142    fn parse_bytes(&mut self, val: &[u8], time_unit: Option<TimeUnit>) -> Option<i64> {
143        let transform = match time_unit {
144            Some(TimeUnit::Nanoseconds) => datetime_to_timestamp_ns,
145            Some(TimeUnit::Microseconds) => datetime_to_timestamp_us,
146            Some(TimeUnit::Milliseconds) => datetime_to_timestamp_ms,
147            _ => unreachable!(), // time_unit has to be provided for datetime
148        };
149        self.transform_bytes
150            .parse(val, self.latest_fmt.as_bytes())
151            .map(transform)
152            .or_else(|| {
153                // TODO! this will try all patterns.
154                // Somehow we must early escape if value is invalid.
155                for fmt in self.patterns {
156                    if let Some(parsed) = self
157                        .transform_bytes
158                        .parse(val, fmt.as_bytes())
159                        .map(transform)
160                    {
161                        self.latest_fmt = fmt;
162                        return Some(parsed);
163                    }
164                }
165                None
166            })
167    }
168}
169
170#[cfg(feature = "dtype-date")]
171impl StrpTimeParser<i32> for DatetimeInfer<Int32Type> {
172    fn parse_bytes(&mut self, val: &[u8], _time_unit: Option<TimeUnit>) -> Option<i32> {
173        self.transform_bytes
174            .parse(val, self.latest_fmt.as_bytes())
175            .map(|ndt| naive_date_to_date(ndt.date()))
176            .or_else(|| {
177                // TODO! this will try all patterns.
178                // somehow we must early escape if value is invalid
179                for fmt in self.patterns {
180                    if let Some(parsed) = self
181                        .transform_bytes
182                        .parse(val, fmt.as_bytes())
183                        .map(|ndt| naive_date_to_date(ndt.date()))
184                    {
185                        self.latest_fmt = fmt;
186                        return Some(parsed);
187                    }
188                }
189                None
190            })
191    }
192}
193
194#[derive(Clone)]
195pub struct DatetimeInfer<T: PolarsNumericType> {
196    pub pattern: Pattern,
197    patterns: &'static [&'static str],
198    latest_fmt: &'static str,
199    transform: fn(&str, &str) -> Option<T::Native>,
200    transform_bytes: StrpTimeState,
201    pub logical_type: DataType,
202}
203
204pub trait TryFromWithUnit<T>: Sized {
205    type Error;
206    fn try_from_with_unit(pattern: T, time_unit: Option<TimeUnit>) -> PolarsResult<Self>;
207}
208
209#[cfg(feature = "dtype-datetime")]
210impl TryFromWithUnit<Pattern> for DatetimeInfer<Int64Type> {
211    type Error = PolarsError;
212
213    fn try_from_with_unit(value: Pattern, time_unit: Option<TimeUnit>) -> PolarsResult<Self> {
214        let time_unit = time_unit.expect("time_unit must be provided for datetime");
215
216        let transform = match (time_unit, value) {
217            (TimeUnit::Milliseconds, Pattern::DatetimeYMDZ) => transform_tzaware_datetime_ms,
218            (TimeUnit::Milliseconds, _) => transform_datetime_ms,
219            (TimeUnit::Microseconds, Pattern::DatetimeYMDZ) => transform_tzaware_datetime_us,
220            (TimeUnit::Microseconds, _) => transform_datetime_us,
221            (TimeUnit::Nanoseconds, Pattern::DatetimeYMDZ) => transform_tzaware_datetime_ns,
222            (TimeUnit::Nanoseconds, _) => transform_datetime_ns,
223        };
224        let (pattern, patterns) = match value {
225            Pattern::DatetimeDMY | Pattern::DateDMY => {
226                (Pattern::DatetimeDMY, patterns::DATETIME_D_M_Y)
227            },
228            Pattern::DatetimeYMD | Pattern::DateYMD => {
229                (Pattern::DatetimeYMD, patterns::DATETIME_Y_M_D)
230            },
231            Pattern::DatetimeYMDZ => (Pattern::DatetimeYMDZ, patterns::DATETIME_Y_M_D_Z),
232            Pattern::Time => (Pattern::Time, patterns::TIME_H_M_S),
233        };
234
235        Ok(DatetimeInfer {
236            pattern,
237            patterns,
238            latest_fmt: patterns[0],
239            transform,
240            transform_bytes: StrpTimeState::default(),
241            logical_type: DataType::Datetime(time_unit, None),
242        })
243    }
244}
245
246#[cfg(feature = "dtype-date")]
247impl TryFromWithUnit<Pattern> for DatetimeInfer<Int32Type> {
248    type Error = PolarsError;
249
250    fn try_from_with_unit(value: Pattern, _time_unit: Option<TimeUnit>) -> PolarsResult<Self> {
251        match value {
252            Pattern::DateDMY => Ok(DatetimeInfer {
253                pattern: Pattern::DateDMY,
254                patterns: patterns::DATE_D_M_Y,
255                latest_fmt: patterns::DATE_D_M_Y[0],
256                transform: transform_date,
257                transform_bytes: StrpTimeState::default(),
258                logical_type: DataType::Date,
259            }),
260            Pattern::DateYMD => Ok(DatetimeInfer {
261                pattern: Pattern::DateYMD,
262                patterns: patterns::DATE_Y_M_D,
263                latest_fmt: patterns::DATE_Y_M_D[0],
264                transform: transform_date,
265                transform_bytes: StrpTimeState::default(),
266                logical_type: DataType::Date,
267            }),
268            _ => polars_bail!(ComputeError: "could not convert pattern"),
269        }
270    }
271}
272
273impl<T: PolarsNumericType> DatetimeInfer<T> {
274    pub fn parse(&mut self, val: &str) -> Option<T::Native> {
275        match (self.transform)(val, self.latest_fmt) {
276            Some(parsed) => Some(parsed),
277            // try other patterns
278            None => {
279                if !self.pattern.is_inferable(val) {
280                    return None;
281                }
282                for fmt in self.patterns {
283                    if let Some(parsed) = (self.transform)(val, fmt) {
284                        self.latest_fmt = fmt;
285                        return Some(parsed);
286                    }
287                }
288                None
289            },
290        }
291    }
292}
293
294impl<T: PolarsNumericType> DatetimeInfer<T> {
295    pub fn coerce_string(&mut self, ca: &StringChunked) -> Series {
296        let chunks = ca.downcast_iter().map(|array| {
297            let iter = array
298                .into_iter()
299                .map(|opt_val| opt_val.and_then(|val| self.parse(val)));
300            PrimitiveArray::from_trusted_len_iter(iter)
301        });
302        ChunkedArray::<T>::from_chunk_iter(ca.name().clone(), chunks)
303            .into_series()
304            .cast(&self.logical_type)
305            .unwrap()
306            .with_name(ca.name().clone())
307    }
308}
309
310#[cfg(feature = "dtype-date")]
311fn transform_date(val: &str, fmt: &str) -> Option<i32> {
312    NaiveDate::parse_from_str(val, fmt)
313        .ok()
314        .map(naive_date_to_date)
315}
316
317pub(crate) fn parse_datetime(val: &str, fmt: &str) -> Option<NaiveDateTime> {
318    NaiveDateTime::parse_from_str(val, fmt)
319        .or_else(|parse_error| match parse_error.kind() {
320            ParseErrorKind::NotEnough => {
321                NaiveDate::parse_from_str(val, fmt).map(|nd| nd.and_hms_opt(0, 0, 0).unwrap())
322            },
323            _ => Err(parse_error),
324        })
325        .ok()
326}
327
328pub(crate) fn parse_datetime_and_remainder<'a>(
329    val: &'a str,
330    fmt: &str,
331) -> Option<(NaiveDateTime, &'a str)> {
332    NaiveDateTime::parse_and_remainder(val, fmt)
333        .or_else(|parse_error| match parse_error.kind() {
334            ParseErrorKind::NotEnough => NaiveDate::parse_and_remainder(val, fmt)
335                .map(|(nd, r)| (nd.and_hms_opt(0, 0, 0).unwrap(), r)),
336            _ => Err(parse_error),
337        })
338        .ok()
339}
340
341#[cfg(feature = "dtype-datetime")]
342pub(crate) fn transform_datetime_ns(val: &str, fmt: &str) -> Option<i64> {
343    parse_datetime(val, fmt).map(datetime_to_timestamp_ns)
344}
345
346#[cfg(feature = "dtype-datetime")]
347pub(crate) fn transform_datetime_us(val: &str, fmt: &str) -> Option<i64> {
348    parse_datetime(val, fmt).map(datetime_to_timestamp_us)
349}
350
351#[cfg(feature = "dtype-datetime")]
352pub(crate) fn transform_datetime_ms(val: &str, fmt: &str) -> Option<i64> {
353    parse_datetime(val, fmt).map(datetime_to_timestamp_ms)
354}
355
356fn transform_tzaware_datetime_ns(val: &str, fmt: &str) -> Option<i64> {
357    let dt = DateTime::parse_from_str(val, fmt);
358    dt.ok().map(|dt| datetime_to_timestamp_ns(dt.naive_utc()))
359}
360
361fn transform_tzaware_datetime_us(val: &str, fmt: &str) -> Option<i64> {
362    let dt = DateTime::parse_from_str(val, fmt);
363    dt.ok().map(|dt| datetime_to_timestamp_us(dt.naive_utc()))
364}
365
366fn transform_tzaware_datetime_ms(val: &str, fmt: &str) -> Option<i64> {
367    let dt = DateTime::parse_from_str(val, fmt);
368    dt.ok().map(|dt| datetime_to_timestamp_ms(dt.naive_utc()))
369}
370
371pub fn infer_pattern_single(val: &str) -> Option<Pattern> {
372    // Dates come first, because we see datetimes as superset of dates
373    infer_pattern_date_single(val)
374        .or_else(|| infer_pattern_time_single(val))
375        .or_else(|| infer_pattern_datetime_single(val))
376}
377
378pub fn infer_pattern_datetime_single(val: &str) -> Option<Pattern> {
379    if patterns::DATETIME_D_M_Y.iter().any(|fmt| {
380        NaiveDateTime::parse_from_str(val, fmt).is_ok()
381            || NaiveDate::parse_from_str(val, fmt).is_ok()
382    }) {
383        Some(Pattern::DatetimeDMY)
384    } else if patterns::DATETIME_Y_M_D.iter().any(|fmt| {
385        NaiveDateTime::parse_from_str(val, fmt).is_ok()
386            || NaiveDate::parse_from_str(val, fmt).is_ok()
387    }) {
388        Some(Pattern::DatetimeYMD)
389    } else if patterns::DATETIME_Y_M_D_Z
390        .iter()
391        .any(|fmt| NaiveDateTime::parse_from_str(val, fmt).is_ok())
392    {
393        Some(Pattern::DatetimeYMDZ)
394    } else {
395        None
396    }
397}
398
399pub fn infer_pattern_date_single(val: &str) -> Option<Pattern> {
400    if patterns::DATE_D_M_Y
401        .iter()
402        .any(|fmt| NaiveDate::parse_from_str(val, fmt).is_ok())
403    {
404        Some(Pattern::DateDMY)
405    } else if patterns::DATE_Y_M_D
406        .iter()
407        .any(|fmt| NaiveDate::parse_from_str(val, fmt).is_ok())
408    {
409        Some(Pattern::DateYMD)
410    } else {
411        None
412    }
413}
414
415pub fn infer_pattern_time_single(val: &str) -> Option<Pattern> {
416    sniff_time_fmt(val).is_some().then_some(Pattern::Time)
417}
418
419/// Return the first format string from `TIME_H_M_S` that parses `val`, or `None`.
420pub fn sniff_time_fmt(val: &str) -> Option<&'static str> {
421    patterns::TIME_H_M_S
422        .iter()
423        .copied()
424        .find(|fmt| NaiveTime::parse_from_str(val, fmt).is_ok())
425}
426
427/// Scan the non-null values for the first that `infer` accepts.
428pub fn infer_from_values<T>(ca: &StringChunked, infer: impl FnMut(&str) -> Option<T>) -> Option<T> {
429    ca.iter().flatten().find_map(infer)
430}
431
432#[cfg(feature = "dtype-datetime")]
433pub fn to_datetime_with_inferred_tz(
434    ca: &StringChunked,
435    tu: TimeUnit,
436    strict: bool,
437    exact: bool,
438    ambiguous: &StringChunked,
439) -> PolarsResult<DatetimeChunked> {
440    use super::StringMethods;
441
442    let out = if exact {
443        to_datetime(ca, tu, None, ambiguous, false)
444    } else {
445        ca.as_datetime_not_exact(None, tu, false, None, ambiguous, false)
446    }?;
447
448    if strict && ca.null_count() != out.null_count() {
449        polars_core::utils::handle_casting_failures(
450            &ca.clone().into_series(),
451            &out.clone().into_series(),
452        )?;
453    }
454
455    Ok(out)
456}
457
458#[cfg(feature = "dtype-datetime")]
459pub fn to_datetime(
460    ca: &StringChunked,
461    tu: TimeUnit,
462    tz: Option<&TimeZone>,
463    ambiguous: &StringChunked,
464    // Ensure that the inferred time_zone matches the given time_zone.
465    ensure_matching_time_zone: bool,
466) -> PolarsResult<DatetimeChunked> {
467    if ca.null_count() == ca.len() {
468        return Ok(
469            Int64Chunked::full_null(ca.name().clone(), ca.len()).into_datetime(tu, tz.cloned())
470        );
471    }
472
473    let pattern = infer_from_values(ca, infer_pattern_datetime_single)
474        .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?;
475    let mut infer = DatetimeInfer::<Int64Type>::try_from_with_unit(pattern, Some(tu))?;
476    #[cfg(feature = "timezones")]
477    if matches!(pattern, Pattern::DatetimeYMDZ) {
478        polars_ensure!(
479            !ensure_matching_time_zone || tz.is_some(),
480            to_datetime_tz_mismatch
481        );
482    }
483    coerce_string_to_datetime(&mut infer, ca, tz, ambiguous)
484}
485/// Apply a pre-built `DatetimeInfer<Int32Type>` to a `StringChunked`, returning a `DateChunked`.
486#[cfg(feature = "dtype-date")]
487pub fn coerce_string_to_date(
488    infer: &mut DatetimeInfer<Int32Type>,
489    ca: &StringChunked,
490) -> PolarsResult<DateChunked> {
491    infer.coerce_string(ca).date().cloned()
492}
493
494/// Apply a pre-built `DatetimeInfer<Int64Type>` to a `StringChunked`, applying tz handling,
495/// returning a `DatetimeChunked`. Mirrors the post-`coerce_string` logic in `to_datetime`.
496#[cfg(feature = "dtype-datetime")]
497pub fn coerce_string_to_datetime(
498    infer: &mut DatetimeInfer<Int64Type>,
499    ca: &StringChunked,
500    tz: Option<&TimeZone>,
501    ambiguous: &StringChunked,
502) -> PolarsResult<DatetimeChunked> {
503    let DataType::Datetime(tu, _) = &infer.logical_type else {
504        unreachable!()
505    };
506    let tu = *tu;
507    match infer.pattern {
508        #[cfg(feature = "timezones")]
509        Pattern::DatetimeYMDZ => infer.coerce_string(ca).datetime().map(|ca| {
510            let mut ca = ca.clone();
511            ca.set_time_unit_and_time_zone(tu, tz.cloned().unwrap_or(TimeZone::UTC))?;
512            Ok(ca)
513        })?,
514        _ => infer.coerce_string(ca).datetime().map(|ca| {
515            let mut ca = ca.clone();
516            ca.set_time_unit(tu);
517            match tz {
518                #[cfg(feature = "timezones")]
519                Some(tz) => polars_ops::prelude::replace_time_zone(
520                    &ca,
521                    Some(tz),
522                    ambiguous,
523                    NonExistent::Raise,
524                ),
525                _ => Ok(ca),
526            }
527        })?,
528    }
529}
530
531#[cfg(feature = "dtype-date")]
532pub(crate) fn to_date(ca: &StringChunked) -> PolarsResult<DateChunked> {
533    if ca.null_count() == ca.len() {
534        return Ok(Int32Chunked::full_null(ca.name().clone(), ca.len()).into_date());
535    }
536
537    let pattern = infer_from_values(ca, infer_pattern_date_single)
538        .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?;
539    let mut infer = DatetimeInfer::<Int32Type>::try_from_with_unit(pattern, None).unwrap();
540    coerce_string_to_date(&mut infer, ca)
541}
polars_time/chunkedarray/string/infer.rs

polars_time/chunkedarray/string/
infer.rs
