polars_time/
round.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::*;
8use crate::truncate::fast_truncate;
9
10#[inline(always)]
11fn fast_round(t: i64, every: i64) -> i64 {
12 fast_truncate(t + every / 2, every)
13}
14
15pub trait PolarsRound {
16 fn round(&self, every: &StringChunked, tz: Option<&Tz>) -> PolarsResult<Self>
17 where
18 Self: Sized;
19}
20
21impl PolarsRound for DatetimeChunked {
22 fn round(&self, every: &StringChunked, tz: Option<&Tz>) -> PolarsResult<Self> {
23 let time_zone = self.time_zone();
24 let offset = Duration::new(0);
25
26 if every.len() == 1 {
28 if let Some(every) = every.get(0) {
29 let every_parsed = Duration::try_parse(every)?;
30 if every_parsed.negative {
31 polars_bail!(ComputeError: "cannot round a Datetime to a negative duration")
32 }
33 if (time_zone.is_none() || time_zone.as_deref() == Some("UTC"))
34 && (every_parsed.months() == 0 && every_parsed.weeks() == 0)
35 {
36 let every = match self.time_unit() {
39 TimeUnit::Milliseconds => every_parsed.duration_ms(),
40 TimeUnit::Microseconds => every_parsed.duration_us(),
41 TimeUnit::Nanoseconds => every_parsed.duration_ns(),
42 };
43 return Ok(self
44 .apply_values(|t| fast_round(t, every))
45 .into_datetime(self.time_unit(), time_zone.clone()));
46 } else {
47 let w = Window::new(every_parsed, every_parsed, offset);
48 let out = match self.time_unit() {
49 TimeUnit::Milliseconds => {
50 self.try_apply_nonnull_values_generic(|t| w.round_ms(t, tz))
51 },
52 TimeUnit::Microseconds => {
53 self.try_apply_nonnull_values_generic(|t| w.round_us(t, tz))
54 },
55 TimeUnit::Nanoseconds => {
56 self.try_apply_nonnull_values_generic(|t| w.round_ns(t, tz))
57 },
58 };
59 return Ok(out?.into_datetime(self.time_unit(), self.time_zone().clone()));
60 }
61 } else {
62 return Ok(Int64Chunked::full_null(self.name().clone(), self.len())
63 .into_datetime(self.time_unit(), self.time_zone().clone()));
64 }
65 }
66
67 polars_ensure!(
68 self.len() == every.len() || self.len() == 1,
69 length_mismatch = "dt.round",
70 self.len(),
71 every.len()
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::round_ns,
79 TimeUnit::Microseconds => Window::round_us,
80 TimeUnit::Milliseconds => Window::round_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 round 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 PolarsRound for DateChunked {
104 fn round(&self, every: &StringChunked, _tz: Option<&Tz>) -> PolarsResult<Self> {
105 let offset = Duration::new(0);
106 let out = match every.len() {
107 1 => {
108 if let Some(every) = every.get(0) {
109 let every = Duration::try_parse(every)?;
110 if every.negative {
111 polars_bail!(ComputeError: "cannot round a Date to a negative duration")
112 }
113 let w = Window::new(every, every, offset);
114 self.try_apply_nonnull_values_generic(|t| {
115 Ok(
116 (w.round_ms(MILLISECONDS_IN_DAY * t as i64, None)?
117 / MILLISECONDS_IN_DAY) as i32,
118 )
119 })
120 } else {
121 Ok(Int32Chunked::full_null(self.name().clone(), self.len()))
122 }
123 },
124 _ => {
125 polars_ensure!(
126 self.len() == every.len() || self.len() == 1,
127 length_mismatch = "dt.round",
128 self.len(),
129 every.len()
130 );
131 broadcast_try_binary_elementwise(self, every, |opt_t, opt_every| {
132 let mut duration_cache =
134 LruCache::with_capacity((every.len() as f64).sqrt() as usize);
135 match (opt_t, opt_every) {
136 (Some(t), Some(every)) => {
137 let every = *duration_cache.get_or_insert_with(every, Duration::parse);
138
139 if every.negative {
140 polars_bail!(ComputeError: "cannot round a Date to a negative duration")
141 }
142
143 let w = Window::new(every, every, offset);
144 Ok(Some(
145 (w.round_ms(MILLISECONDS_IN_DAY * t as i64, None)?
146 / MILLISECONDS_IN_DAY) as i32,
147 ))
148 },
149 _ => Ok(None),
150 }
151 })
152 },
153 };
154 Ok(out?.into_date())
155 }
156}