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