per(rust, python): use global string cache for rev cat (#6067)

pola-rs · Jan 5, 2023 · a8e7e4a · a8e7e4a
1 parent 41f0aa6
commit a8e7e4a
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/polars/polars-core/src/chunked_array/logical/categorical/builder.rs b/polars/polars-core/src/chunked_array/logical/categorical/builder.rs
@@ -122,8 +122,16 @@ impl RevMapping {
     /// str to Categorical
     pub fn find(&self, value: &str) -> Option<u32> {
         match self {
-            Self::Global(map, a, _) => {
-                map.iter()
+            Self::Global(rev_map, a, id) => {
+                // fast path is check
+                if using_string_cache() {
+                    let map = crate::STRING_CACHE.read_map();
+                    if map.uuid == *id {
+                        return map.get_cat(value);
+                    }
+                }
+                rev_map
+                    .iter()
                     // Safety:
                     // value is always within bounds
                     .find(|(_k, &v)| (unsafe { a.value_unchecked(v as usize) } == value))

diff --git a/polars/polars-core/src/chunked_array/logical/categorical/stringcache.rs b/polars/polars-core/src/chunked_array/logical/categorical/stringcache.rs
@@ -145,6 +145,22 @@ impl SCacheInner {
         global_idx
     }
 
+    #[inline]
+    pub(crate) fn get_cat(&self, s: &str) -> Option<u32> {
+        let h = StringCache::get_hash_builder().hash_single(s);
+        // as StrHashGlobal may allocate a string
+        self.map
+            .raw_entry()
+            .from_hash(h, |key| {
+                (key.hash == h) && {
+                    let pos = key.idx as usize;
+                    let value = unsafe { self.payloads.get_unchecked(pos) };
+                    s == value.as_str()
+                }
+            })
+            .map(|(k, _)| k.idx)
+    }
+
     #[inline]
     pub(crate) fn insert(&mut self, s: &str) -> u32 {
         let h = StringCache::get_hash_builder().hash_single(s);