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

polars_time/
truncate.rs
