delay initialising cache

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<Self> {
         let offset: Duration = Duration::parse(offset);
         let time_zone = self.time_zone();
-        let mut duration_cache_opt: Option<FastFixedCache<String, Duration>> = None;
 
         // Let's check if we can use a fastpath...
         if every.len() == 1 {
@@ -42,64 +41,52 @@ 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,
             TimeUnit::Microseconds => Window::truncate_us,
             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<Self> {
         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