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

polars_time/
truncate.rs
