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::FastFixedCache;
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        let time_zone = self.time_zone();
24        let offset = Duration::new(0);
25
26        // Let's check if we can use a fastpath...
27        if every.len() == 1 {
28            if let Some(every) = every.get(0) {
29                let every_parsed = Duration::parse(every);
30                if every_parsed.negative {
31                    polars_bail!(ComputeError: "cannot truncate 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                    // ... yes we can! Weeks, months, and time zones require extra logic.
37                    // But in this simple case, it's just simple integer arithmetic.
38                    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_truncate(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.truncate_ms(t, tz))
51                        },
52                        TimeUnit::Microseconds => {
53                            self.try_apply_nonnull_values_generic(|t| w.truncate_us(t, tz))
54                        },
55                        TimeUnit::Nanoseconds => {
56                            self.try_apply_nonnull_values_generic(|t| w.truncate_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        // A sqrt(n) cache is not too small, not too large.
68        let mut duration_cache = FastFixedCache::new((every.len() as f64).sqrt() as usize);
69
70        let func = match self.time_unit() {
71            TimeUnit::Nanoseconds => Window::truncate_ns,
72            TimeUnit::Microseconds => Window::truncate_us,
73            TimeUnit::Milliseconds => Window::truncate_ms,
74        };
75
76        let out = broadcast_try_binary_elementwise(self, every, |opt_timestamp, opt_every| match (
77            opt_timestamp,
78            opt_every,
79        ) {
80            (Some(timestamp), Some(every)) => {
81                let every =
82                    *duration_cache.get_or_insert_with(every, |every| Duration::parse(every));
83
84                if every.negative {
85                    polars_bail!(ComputeError: "cannot truncate a Datetime to a negative duration")
86                }
87
88                let w = Window::new(every, every, offset);
89                func(&w, timestamp, tz).map(Some)
90            },
91            _ => Ok(None),
92        });
93        Ok(out?.into_datetime(self.time_unit(), self.time_zone().clone()))
94    }
95}
96
97impl PolarsTruncate for DateChunked {
98    fn truncate(&self, _tz: Option<&Tz>, every: &StringChunked) -> PolarsResult<Self> {
99        let offset = Duration::new(0);
100        let out = match every.len() {
101            1 => {
102                if let Some(every) = every.get(0) {
103                    let every = Duration::parse(every);
104                    if every.negative {
105                        polars_bail!(ComputeError: "cannot truncate a Date to a negative duration")
106                    }
107                    let w = Window::new(every, every, offset);
108                    self.try_apply_nonnull_values_generic(|t| {
109                        Ok((w.truncate_ms(MILLISECONDS_IN_DAY * t as i64, None)?
110                            / MILLISECONDS_IN_DAY) as i32)
111                    })
112                } else {
113                    Ok(Int32Chunked::full_null(self.name().clone(), self.len()))
114                }
115            },
116            _ => broadcast_try_binary_elementwise(self, every, |opt_t, opt_every| {
117                // A sqrt(n) cache is not too small, not too large.
118                let mut duration_cache = FastFixedCache::new((every.len() as f64).sqrt() as usize);
119                match (opt_t, opt_every) {
120                    (Some(t), Some(every)) => {
121                        let every = *duration_cache
122                            .get_or_insert_with(every, |every| Duration::parse(every));
123
124                        if every.negative {
125                            polars_bail!(ComputeError: "cannot truncate a Date to a negative duration")
126                        }
127
128                        let w = Window::new(every, every, offset);
129                        Ok(Some(
130                            (w.truncate_ms(MILLISECONDS_IN_DAY * t as i64, None)?
131                                / MILLISECONDS_IN_DAY) as i32,
132                        ))
133                    },
134                    _ => Ok(None),
135                }
136            }),
137        };
138        Ok(out?.into_date())
139    }
140}
polars_time/truncate.rs

polars_time/
truncate.rs
