polars_time/
truncate.rs
1use arrow::legacy::time_zone::Tz;
2use arrow::temporal_conversions::MILLISECONDS_IN_DAY;
3use polars_core::prelude::arity::broadcast_try_binary_elementwise;
4use polars_core::prelude::*;
5use polars_utils::cache::LruCache;
6
7use crate::prelude::*;
8
9pub trait PolarsTruncate {
10 fn truncate(&self, tz: Option<&Tz>, every: &StringChunked) -> PolarsResult<Self>
11 where
12 Self: Sized;
13}
14
15#[inline(always)]
16pub(crate) fn fast_truncate(t: i64, every: i64) -> i64 {
17 let remainder = t % every;
18 t - (remainder + every * (remainder < 0) as i64)
19}
20
21impl PolarsTruncate for DatetimeChunked {
22 fn truncate(&self, tz: Option<&Tz>, every: &StringChunked) -> PolarsResult<Self> {
23 polars_ensure!(
24 self.len() == every.len() || self.len() == 1 || every.len() == 1,
25 length_mismatch = "dt.truncate",
26 self.len(),
27 every.len()
28 );
29
30 let time_zone = self.time_zone();
31 let offset = Duration::new(0);
32
33 if every.len() == 1 {
35 if let Some(every) = every.get(0) {
36 let every_parsed = Duration::parse(every);
37 if every_parsed.negative {
38 polars_bail!(ComputeError: "cannot truncate a Datetime to a negative duration")
39 }
40 if (time_zone.is_none() || time_zone.as_deref() == Some("UTC"))
41 && (every_parsed.months() == 0 && every_parsed.weeks() == 0)
42 {
43 let every = match self.time_unit() {
46 TimeUnit::Milliseconds => every_parsed.duration_ms(),
47 TimeUnit::Microseconds => every_parsed.duration_us(),
48 TimeUnit::Nanoseconds => every_parsed.duration_ns(),
49 };
50 return Ok(self
51 .apply_values(|t| fast_truncate(t, every))
52 .into_datetime(self.time_unit(), time_zone.clone()));
53 } else {
54 let w = Window::new(every_parsed, every_parsed, offset);
55 let out = match self.time_unit() {
56 TimeUnit::Milliseconds => {
57 self.try_apply_nonnull_values_generic(|t| w.truncate_ms(t, tz))
58 },
59 TimeUnit::Microseconds => {
60 self.try_apply_nonnull_values_generic(|t| w.truncate_us(t, tz))
61 },
62 TimeUnit::Nanoseconds => {
63 self.try_apply_nonnull_values_generic(|t| w.truncate_ns(t, tz))
64 },
65 };
66 return Ok(out?.into_datetime(self.time_unit(), self.time_zone().clone()));
67 }
68 } else {
69 return Ok(Int64Chunked::full_null(self.name().clone(), self.len())
70 .into_datetime(self.time_unit(), self.time_zone().clone()));
71 }
72 }
73
74 let mut duration_cache = LruCache::with_capacity((every.len() as f64).sqrt() as usize);
76
77 let func = match self.time_unit() {
78 TimeUnit::Nanoseconds => Window::truncate_ns,
79 TimeUnit::Microseconds => Window::truncate_us,
80 TimeUnit::Milliseconds => Window::truncate_ms,
81 };
82
83 let out = broadcast_try_binary_elementwise(self, every, |opt_timestamp, opt_every| match (
84 opt_timestamp,
85 opt_every,
86 ) {
87 (Some(timestamp), Some(every)) => {
88 let every = *duration_cache.get_or_insert_with(every, Duration::parse);
89
90 if every.negative {
91 polars_bail!(ComputeError: "cannot truncate a Datetime to a negative duration")
92 }
93
94 let w = Window::new(every, every, offset);
95 func(&w, timestamp, tz).map(Some)
96 },
97 _ => Ok(None),
98 });
99 Ok(out?.into_datetime(self.time_unit(), self.time_zone().clone()))
100 }
101}
102
103impl PolarsTruncate for DateChunked {
104 fn truncate(&self, _tz: Option<&Tz>, every: &StringChunked) -> PolarsResult<Self> {
105 polars_ensure!(
106 self.len() == every.len() || self.len() == 1 || every.len() == 1,
107 length_mismatch = "dt.truncate",
108 self.len(),
109 every.len()
110 );
111
112 let offset = Duration::new(0);
113 let out = match every.len() {
114 1 => {
115 if let Some(every) = every.get(0) {
116 let every = Duration::parse(every);
117 if every.negative {
118 polars_bail!(ComputeError: "cannot truncate a Date to a negative duration")
119 }
120 let w = Window::new(every, every, offset);
121 self.try_apply_nonnull_values_generic(|t| {
122 Ok((w.truncate_ms(MILLISECONDS_IN_DAY * t as i64, None)?
123 / MILLISECONDS_IN_DAY) as i32)
124 })
125 } else {
126 Ok(Int32Chunked::full_null(self.name().clone(), self.len()))
127 }
128 },
129 _ => broadcast_try_binary_elementwise(self, every, |opt_t, opt_every| {
130 let mut duration_cache =
132 LruCache::with_capacity((every.len() as f64).sqrt() as usize);
133 match (opt_t, opt_every) {
134 (Some(t), Some(every)) => {
135 let every = *duration_cache.get_or_insert_with(every, Duration::parse);
136
137 if every.negative {
138 polars_bail!(ComputeError: "cannot truncate a Date to a negative duration")
139 }
140
141 let w = Window::new(every, every, offset);
142 Ok(Some(
143 (w.truncate_ms(MILLISECONDS_IN_DAY * t as i64, None)?
144 / MILLISECONDS_IN_DAY) as i32,
145 ))
146 },
147 _ => Ok(None),
148 }
149 }),
150 };
151 Ok(out?.into_date())
152 }
153}