polars_arrow/legacy/kernels/
time.rs

1use std::str::FromStr;
2
3#[cfg(feature = "timezones")]
4use chrono::{LocalResult, NaiveDateTime, TimeZone};
5#[cfg(feature = "timezones")]
6use chrono_tz::Tz;
7#[cfg(feature = "timezones")]
8use polars_error::PolarsResult;
9use polars_error::{PolarsError, polars_bail};
10#[cfg(feature = "serde")]
11use serde::{Deserialize, Serialize};
12use strum_macros::IntoStaticStr;
13
14pub enum Ambiguous {
15    Earliest,
16    Latest,
17    Null,
18    Raise,
19}
20impl FromStr for Ambiguous {
21    type Err = PolarsError;
22
23    fn from_str(s: &str) -> Result<Self, Self::Err> {
24        match s {
25            "earliest" => Ok(Ambiguous::Earliest),
26            "latest" => Ok(Ambiguous::Latest),
27            "raise" => Ok(Ambiguous::Raise),
28            "null" => Ok(Ambiguous::Null),
29            s => polars_bail!(InvalidOperation:
30                "Invalid argument {}, expected one of: \"earliest\", \"latest\", \"null\", \"raise\"", s
31            ),
32        }
33    }
34}
35
36#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, IntoStaticStr)]
37#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
38#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
39#[strum(serialize_all = "snake_case")]
40pub enum NonExistent {
41    Null,
42    Raise,
43}
44
45#[cfg(feature = "timezones")]
46pub fn convert_to_naive_local(
47    from_tz: &Tz,
48    to_tz: &Tz,
49    ndt: NaiveDateTime,
50    ambiguous: Ambiguous,
51    non_existent: NonExistent,
52) -> PolarsResult<Option<NaiveDateTime>> {
53    let ndt = from_tz.from_utc_datetime(&ndt).naive_local();
54    match to_tz.from_local_datetime(&ndt) {
55        LocalResult::Single(dt) => Ok(Some(dt.naive_utc())),
56        LocalResult::Ambiguous(dt_earliest, dt_latest) => match ambiguous {
57            Ambiguous::Earliest => Ok(Some(dt_earliest.naive_utc())),
58            Ambiguous::Latest => Ok(Some(dt_latest.naive_utc())),
59            Ambiguous::Null => Ok(None),
60            Ambiguous::Raise => {
61                polars_bail!(ComputeError: "datetime '{}' is ambiguous in time zone '{}'. Please use `ambiguous` to tell how it should be localized.", ndt, to_tz)
62            },
63        },
64        LocalResult::None => match non_existent {
65            NonExistent::Raise => polars_bail!(ComputeError:
66                "datetime '{}' is non-existent in time zone '{}'. You may be able to use `non_existent='null'` to return `null` in this case.",
67                ndt, to_tz
68            ),
69            NonExistent::Null => Ok(None),
70        },
71    }
72}
73
74/// Same as convert_to_naive_local, but return `None` instead
75/// raising - in some cases this can be used to save a string allocation.
76#[cfg(feature = "timezones")]
77pub fn convert_to_naive_local_opt(
78    from_tz: &Tz,
79    to_tz: &Tz,
80    ndt: NaiveDateTime,
81    ambiguous: Ambiguous,
82) -> Option<Option<NaiveDateTime>> {
83    let ndt = from_tz.from_utc_datetime(&ndt).naive_local();
84    match to_tz.from_local_datetime(&ndt) {
85        LocalResult::Single(dt) => Some(Some(dt.naive_utc())),
86        LocalResult::Ambiguous(dt_earliest, dt_latest) => match ambiguous {
87            Ambiguous::Earliest => Some(Some(dt_earliest.naive_utc())),
88            Ambiguous::Latest => Some(Some(dt_latest.naive_utc())),
89            Ambiguous::Null => Some(None),
90            Ambiguous::Raise => None,
91        },
92        LocalResult::None => None,
93    }
94}