polars_time/
round.rs

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

polars_time/
round.rs
