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

polars_time/
truncate.rs
