diff --git a/crates/polars-time/src/truncate.rs b/crates/polars-time/src/truncate.rs index a8ee41daefdd7..7940f4e5a7069 100644 --- a/crates/polars-time/src/truncate.rs +++ b/crates/polars-time/src/truncate.rs @@ -16,7 +16,6 @@ impl PolarsTruncate for DatetimeChunked { fn truncate(&self, tz: Option<&Tz>, every: &StringChunked, offset: &str) -> PolarsResult { let offset: Duration = Duration::parse(offset); let time_zone = self.time_zone(); - let mut duration_cache_opt: Option> = None; // Let's check if we can use a fastpath... if every.len() == 1 { @@ -42,19 +41,28 @@ impl PolarsTruncate for DatetimeChunked { }) .into_datetime(self.time_unit(), time_zone.clone())); } else { - // A sqrt(n) cache is not too small, not too large. - duration_cache_opt = - Some(FastFixedCache::new((every.len() as f64).sqrt() as usize)); - duration_cache_opt - .as_mut() - .map(|cache| *cache.insert(every.to_string(), every_parsed)); + let w = Window::new(every_parsed, every_parsed, offset); + let out = match self.time_unit() { + TimeUnit::Milliseconds => { + self.try_apply_nonnull_values_generic(|t| w.truncate_ms(t, tz)) + }, + TimeUnit::Microseconds => { + self.try_apply_nonnull_values_generic(|t| w.truncate_us(t, tz)) + }, + TimeUnit::Nanoseconds => { + self.try_apply_nonnull_values_generic(|t| w.truncate_ns(t, tz)) + }, + }; + return Ok(out?.into_datetime(self.time_unit(), self.time_zone().clone())); } + } else { + return Ok(Int64Chunked::full_null(self.name(), self.len()) + .into_datetime(self.time_unit(), self.time_zone().clone())); } } - let mut duration_cache = match duration_cache_opt { - Some(cache) => cache, - None => FastFixedCache::new((every.len() as f64).sqrt() as usize), - }; + + // A sqrt(n) cache is not too small, not too large. + let mut duration_cache = FastFixedCache::new((every.len() as f64).sqrt() as usize); let func = match self.time_unit() { TimeUnit::Nanoseconds => Window::truncate_ns, @@ -62,44 +70,23 @@ impl PolarsTruncate for DatetimeChunked { TimeUnit::Milliseconds => Window::truncate_ms, }; - // TODO: optimize the code below, so it does the following: - // - convert to naive - // - truncate all naively - // - localize, preserving the fold of the original datetime. - // The last step is the non-trivial one. But it should be worth it, - // and faster than the current approach of truncating everything - // as tz-aware. + let out = broadcast_try_binary_elementwise(self, every, |opt_timestamp, opt_every| match ( + opt_timestamp, + opt_every, + ) { + (Some(timestamp), Some(every)) => { + let every = + *duration_cache.get_or_insert_with(every, |every| Duration::parse(every)); - let out = match every.len() { - 1 => { - if let Some(every) = every.get(0) { - let every = Duration::parse(every); - if every.negative { - polars_bail!(ComputeError: "cannot truncate a Datetime to a negative duration") - } - let w = Window::new(every, every, offset); - self.try_apply_nonnull_values_generic(|t| func(&w, t, tz)) - } else { - Ok(Int64Chunked::full_null(self.name(), self.len())) + if every.negative { + polars_bail!(ComputeError: "cannot truncate a Datetime to a negative duration") } - }, - _ => broadcast_try_binary_elementwise(self, every, |opt_timestamp, opt_every| { - match (opt_timestamp, opt_every) { - (Some(timestamp), Some(every)) => { - let every = *duration_cache - .get_or_insert_with(every, |every| Duration::parse(every)); - if every.negative { - polars_bail!(ComputeError: "cannot truncate a Datetime to a negative duration") - } - - let w = Window::new(every, every, offset); - func(&w, timestamp, tz).map(Some) - }, - _ => Ok(None), - } - }), - }; + let w = Window::new(every, every, offset); + func(&w, timestamp, tz).map(Some) + }, + _ => Ok(None), + }); Ok(out?.into_datetime(self.time_unit(), self.time_zone().clone())) } } @@ -112,8 +99,6 @@ impl PolarsTruncate for DateChunked { offset: &str, ) -> PolarsResult { let offset = Duration::parse(offset); - // A sqrt(n) cache is not too small, not too large. - let mut duration_cache = FastFixedCache::new((every.len() as f64).sqrt() as usize); let out = match every.len() { 1 => { if let Some(every) = every.get(0) { @@ -131,6 +116,8 @@ impl PolarsTruncate for DateChunked { } }, _ => broadcast_try_binary_elementwise(self, every, |opt_t, opt_every| { + // A sqrt(n) cache is not too small, not too large. + let mut duration_cache = FastFixedCache::new((every.len() as f64).sqrt() as usize); match (opt_t, opt_every) { (Some(t), Some(every)) => { let every = *duration_cache