polars_plan/dsl/function_expr/
temporal.rs

1use super::*;
2use crate::{map, map_as_slice};
3
4impl From<TemporalFunction> for SpecialEq<Arc<dyn ColumnsUdf>> {
5    fn from(func: TemporalFunction) -> Self {
6        use TemporalFunction::*;
7        match func {
8            Millennium => map!(datetime::millennium),
9            Century => map!(datetime::century),
10            Year => map!(datetime::year),
11            IsLeapYear => map!(datetime::is_leap_year),
12            IsoYear => map!(datetime::iso_year),
13            Month => map!(datetime::month),
14            Quarter => map!(datetime::quarter),
15            Week => map!(datetime::week),
16            WeekDay => map!(datetime::weekday),
17            Duration(tu) => map_as_slice!(impl_duration, tu),
18            Day => map!(datetime::day),
19            OrdinalDay => map!(datetime::ordinal_day),
20            Time => map!(datetime::time),
21            Date => map!(datetime::date),
22            Datetime => map!(datetime::datetime),
23            Hour => map!(datetime::hour),
24            Minute => map!(datetime::minute),
25            Second => map!(datetime::second),
26            Millisecond => map!(datetime::millisecond),
27            Microsecond => map!(datetime::microsecond),
28            Nanosecond => map!(datetime::nanosecond),
29            TotalDays => map!(datetime::total_days),
30            TotalHours => map!(datetime::total_hours),
31            TotalMinutes => map!(datetime::total_minutes),
32            TotalSeconds => map!(datetime::total_seconds),
33            TotalMilliseconds => map!(datetime::total_milliseconds),
34            TotalMicroseconds => map!(datetime::total_microseconds),
35            TotalNanoseconds => map!(datetime::total_nanoseconds),
36            ToString(format) => map!(datetime::to_string, &format),
37            TimeStamp(tu) => map!(datetime::timestamp, tu),
38            #[cfg(feature = "timezones")]
39            ConvertTimeZone(tz) => map!(datetime::convert_time_zone, &tz),
40            WithTimeUnit(tu) => map!(datetime::with_time_unit, tu),
41            CastTimeUnit(tu) => map!(datetime::cast_time_unit, tu),
42            Truncate => {
43                map_as_slice!(datetime::truncate)
44            },
45            #[cfg(feature = "offset_by")]
46            OffsetBy => {
47                map_as_slice!(datetime::offset_by)
48            },
49            #[cfg(feature = "month_start")]
50            MonthStart => map!(datetime::month_start),
51            #[cfg(feature = "month_end")]
52            MonthEnd => map!(datetime::month_end),
53            #[cfg(feature = "timezones")]
54            BaseUtcOffset => map!(datetime::base_utc_offset),
55            #[cfg(feature = "timezones")]
56            DSTOffset => map!(datetime::dst_offset),
57            Round => map_as_slice!(datetime::round),
58            Replace => map_as_slice!(datetime::replace),
59            #[cfg(feature = "timezones")]
60            ReplaceTimeZone(tz, non_existent) => {
61                map_as_slice!(dispatch::replace_time_zone, tz.as_ref(), non_existent)
62            },
63            Combine(tu) => map_as_slice!(temporal::combine, tu),
64            DatetimeFunction {
65                time_unit,
66                time_zone,
67            } => {
68                map_as_slice!(temporal::datetime, &time_unit, time_zone.as_ref())
69            },
70        }
71    }
72}
73
74#[cfg(feature = "dtype-datetime")]
75pub(super) fn datetime(
76    s: &[Column],
77    time_unit: &TimeUnit,
78    time_zone: Option<&TimeZone>,
79) -> PolarsResult<Column> {
80    let col_name = PlSmallStr::from_static("datetime");
81
82    if s.iter().any(|s| s.is_empty()) {
83        return Ok(Column::new_empty(
84            col_name,
85            &DataType::Datetime(
86                time_unit.to_owned(),
87                match time_zone.cloned() {
88                    #[cfg(feature = "timezones")]
89                    Some(v) => Some(v),
90                    _ => {
91                        assert!(
92                            time_zone.is_none(),
93                            "cannot make use of the `time_zone` argument without the 'timezones' feature enabled."
94                        );
95                        None
96                    },
97                },
98            ),
99        ));
100    }
101
102    let year = &s[0];
103    let month = &s[1];
104    let day = &s[2];
105    let hour = &s[3];
106    let minute = &s[4];
107    let second = &s[5];
108    let microsecond = &s[6];
109    let ambiguous = &s[7];
110
111    let max_len = s.iter().map(|s| s.len()).max().unwrap();
112
113    let mut year = year.cast(&DataType::Int32)?;
114    if year.len() < max_len {
115        year = year.new_from_index(0, max_len)
116    }
117    let year = year.i32()?;
118
119    let mut month = month.cast(&DataType::Int8)?;
120    if month.len() < max_len {
121        month = month.new_from_index(0, max_len);
122    }
123    let month = month.i8()?;
124
125    let mut day = day.cast(&DataType::Int8)?;
126    if day.len() < max_len {
127        day = day.new_from_index(0, max_len);
128    }
129    let day = day.i8()?;
130
131    let mut hour = hour.cast(&DataType::Int8)?;
132    if hour.len() < max_len {
133        hour = hour.new_from_index(0, max_len);
134    }
135    let hour = hour.i8()?;
136
137    let mut minute = minute.cast(&DataType::Int8)?;
138    if minute.len() < max_len {
139        minute = minute.new_from_index(0, max_len);
140    }
141    let minute = minute.i8()?;
142
143    let mut second = second.cast(&DataType::Int8)?;
144    if second.len() < max_len {
145        second = second.new_from_index(0, max_len);
146    }
147    let second = second.i8()?;
148
149    let mut nanosecond = microsecond.cast(&DataType::Int32)? * 1_000;
150    if nanosecond.len() < max_len {
151        nanosecond = nanosecond.new_from_index(0, max_len);
152    }
153    let nanosecond = nanosecond.i32()?;
154
155    let mut _ambiguous = ambiguous.cast(&DataType::String)?;
156    if _ambiguous.len() < max_len {
157        _ambiguous = _ambiguous.new_from_index(0, max_len);
158    }
159    let ambiguous = _ambiguous.str()?;
160
161    let ca = DatetimeChunked::new_from_parts(
162        year,
163        month,
164        day,
165        hour,
166        minute,
167        second,
168        nanosecond,
169        ambiguous,
170        time_unit,
171        time_zone.cloned(),
172        col_name,
173    );
174    ca.map(|s| s.into_column())
175}
176
177pub(super) fn combine(s: &[Column], tu: TimeUnit) -> PolarsResult<Column> {
178    let date = &s[0];
179    let time = &s[1];
180
181    let tz = match date.dtype() {
182        DataType::Date => None,
183        DataType::Datetime(_, tz) => tz.as_ref(),
184        _dtype => {
185            polars_bail!(ComputeError: format!("expected Date or Datetime, got {}", _dtype))
186        },
187    };
188
189    let date = date.cast(&DataType::Date)?;
190    let datetime = date.cast(&DataType::Datetime(tu, None)).unwrap();
191
192    let duration = time.cast(&DataType::Duration(tu))?;
193    let result_naive = datetime + duration;
194    match tz {
195        #[cfg(feature = "timezones")]
196        Some(tz) => Ok(polars_ops::prelude::replace_time_zone(
197            result_naive?.datetime().unwrap(),
198            Some(tz),
199            &StringChunked::from_iter(std::iter::once("raise")),
200            NonExistent::Raise,
201        )?
202        .into_column()),
203        _ => result_naive,
204    }
205}
polars_plan/dsl/function_expr/temporal.rs

polars_plan/dsl/function_expr/
temporal.rs
