perf: Limit the cache size for to_datetime

pola-rs · Apr 22, 2024 · 687302e · 687302e
1 parent 0c2783a
commit 687302e
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 36 deletions.
diff --git a/crates/polars-time/src/chunkedarray/string/mod.rs b/crates/polars-time/src/chunkedarray/string/mod.rs
@@ -110,10 +110,13 @@ pub trait StringMethods: AsString {
         };
         let use_cache = use_cache && string_ca.len() > 50;
 
-        let mut convert = CachedFunc::new(|s| {
-            let naive_time = NaiveTime::parse_from_str(s, fmt).ok()?;
-            Some(time_to_time64ns(&naive_time))
-        });
+        let mut convert = CachedFunc::new(
+            |s| {
+                let naive_time = NaiveTime::parse_from_str(s, fmt).ok()?;
+                Some(time_to_time64ns(&naive_time))
+            },
+            (string_ca.len() as f64).sqrt() as usize,
+        );
         let ca = string_ca.apply_generic(|opt_s| convert.eval(opt_s?, use_cache));
         Ok(ca.with_name(string_ca.name()).into())
     }
@@ -237,21 +240,27 @@ pub trait StringMethods: AsString {
         // We can use the fast parser.
         let ca = if let Some(fmt_len) = strptime::fmt_len(fmt.as_bytes()) {
             let mut strptime_cache = StrpTimeState::default();
-            let mut convert = CachedFunc::new(|s: &str| {
-                // SAFETY: fmt_len is correct, it was computed with this `fmt` str.
-                match unsafe { strptime_cache.parse(s.as_bytes(), fmt.as_bytes(), fmt_len) } {
-                    // Fallback to chrono.
-                    None => NaiveDate::parse_from_str(s, &fmt).ok(),
-                    Some(ndt) => Some(ndt.date()),
-                }
-                .map(naive_date_to_date)
-            });
+            let mut convert = CachedFunc::new(
+                |s: &str| {
+                    // SAFETY: fmt_len is correct, it was computed with this `fmt` str.
+                    match unsafe { strptime_cache.parse(s.as_bytes(), fmt.as_bytes(), fmt_len) } {
+                        // Fallback to chrono.
+                        None => NaiveDate::parse_from_str(s, &fmt).ok(),
+                        Some(ndt) => Some(ndt.date()),
+                    }
+                    .map(naive_date_to_date)
+                },
+                (string_ca.len() as f64).sqrt() as usize,
+            );
             string_ca.apply_generic(|val| convert.eval(val?, use_cache))
         } else {
-            let mut convert = CachedFunc::new(|s| {
-                let naive_date = NaiveDate::parse_from_str(s, &fmt).ok()?;
-                Some(naive_date_to_date(naive_date))
-            });
+            let mut convert = CachedFunc::new(
+                |s| {
+                    let naive_date = NaiveDate::parse_from_str(s, &fmt).ok()?;
+                    Some(naive_date_to_date(naive_date))
+                },
+                (string_ca.len() as f64).sqrt() as usize,
+            );
             string_ca.apply_generic(|val| convert.eval(val?, use_cache))
         };
 
@@ -286,10 +295,13 @@ pub trait StringMethods: AsString {
         if tz_aware {
             #[cfg(feature = "timezones")]
             {
-                let mut convert = CachedFunc::new(|s: &str| {
-                    let dt = DateTime::parse_from_str(s, &fmt).ok()?;
-                    Some(func(dt.naive_utc()))
-                });
+                let mut convert = CachedFunc::new(
+                    |s: &str| {
+                        let dt = DateTime::parse_from_str(s, &fmt).ok()?;
+                        Some(func(dt.naive_utc()))
+                    },
+                    (string_ca.len() as f64).sqrt() as usize,
+                );
                 Ok(string_ca
                     .apply_generic(|opt_s| convert.eval(opt_s?, use_cache))
                     .with_name(string_ca.name())
@@ -308,16 +320,23 @@ pub trait StringMethods: AsString {
             // We can use the fast parser.
             let ca = if let Some(fmt_len) = self::strptime::fmt_len(fmt.as_bytes()) {
                 let mut strptime_cache = StrpTimeState::default();
-                let mut convert = CachedFunc::new(|s: &str| {
-                    // SAFETY: fmt_len is correct, it was computed with this `fmt` str.
-                    match unsafe { strptime_cache.parse(s.as_bytes(), fmt.as_bytes(), fmt_len) } {
-                        None => transform(s, &fmt),
-                        Some(ndt) => Some(func(ndt)),
-                    }
-                });
+                let mut convert = CachedFunc::new(
+                    |s: &str| {
+                        // SAFETY: fmt_len is correct, it was computed with this `fmt` str.
+                        match unsafe { strptime_cache.parse(s.as_bytes(), fmt.as_bytes(), fmt_len) }
+                        {
+                            None => transform(s, &fmt),
+                            Some(ndt) => Some(func(ndt)),
+                        }
+                    },
+                    (string_ca.len() as f64).sqrt() as usize,
+                );
                 string_ca.apply_generic(|opt_s| convert.eval(opt_s?, use_cache))
             } else {
-                let mut convert = CachedFunc::new(|s| transform(s, &fmt));
+                let mut convert = CachedFunc::new(
+                    |s| transform(s, &fmt),
+                    (string_ca.len() as f64).sqrt() as usize,
+                );
                 string_ca.apply_generic(|opt_s| convert.eval(opt_s?, use_cache))
             };
             let dt = ca.with_name(string_ca.name()).into_datetime(tu, None);

diff --git a/crates/polars-utils/src/cache.rs b/crates/polars-utils/src/cache.rs
@@ -7,11 +7,9 @@ use ahash::RandomState;
 use bytemuck::allocation::zeroed_vec;
 use bytemuck::Zeroable;
 
-use crate::aliases::PlHashMap;
-
 pub struct CachedFunc<T, R, F> {
     func: F,
-    cache: PlHashMap<T, R>,
+    cache: FastFixedCache<T, R>,
 }
 
 impl<T, R, F> CachedFunc<T, R, F>
@@ -20,19 +18,18 @@ where
     T: std::hash::Hash + Eq + Clone,
     R: Copy,
 {
-    pub fn new(func: F) -> Self {
+    pub fn new(func: F, size: usize) -> Self {
         Self {
             func,
-            cache: PlHashMap::with_capacity_and_hasher(0, Default::default()),
+            cache: FastFixedCache::new(size),
         }
     }
 
     pub fn eval(&mut self, x: T, use_cache: bool) -> R {
         if use_cache {
             *self
                 .cache
-                .entry(x)
-                .or_insert_with_key(|xr| (self.func)(xr.clone()))
+                .get_or_insert_with(&x, |xr| (self.func)(xr.clone()))
         } else {
             (self.func)(x)
         }