Make method generic

pola-rs · Nov 23, 2023 · 406277d · 406277d
1 parent 9df8b61
commit 406277d
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/crates/polars-core/src/chunked_array/logical/categorical/builder.rs b/crates/polars-core/src/chunked_array/logical/categorical/builder.rs
@@ -336,7 +336,10 @@ impl<'a> CategoricalChunkedBuilder<'a> {
     }
 
     /// Build a global string cached [`CategoricalChunked`] from a local [`Dictionary`].
-    pub(super) fn global_map_from_local(&mut self, keys: &UInt32Array, values: Utf8Array<i64>) {
+    pub(super) fn global_map_from_local<I>(&mut self, keys: I, values: Utf8Array<i64>)
+    where
+        I: IntoIterator<Item = Option<u32>> + Send,
+    {
         // locally we don't need a hashmap because we all categories are 1 integer apart
         // so the index is local, and the values is global
         let mut local_to_global: Vec<u32> = Vec::with_capacity(values.len());
@@ -370,8 +373,8 @@ impl<'a> CategoricalChunkedBuilder<'a> {
             keys.into_iter()
                 .map(|opt_k| {
                     opt_k.map(|cat| {
-                        debug_assert!((*cat as usize) < local_to_global.len());
-                        *unsafe { local_to_global.get_unchecked(*cat as usize) }
+                        debug_assert!((cat as usize) < local_to_global.len());
+                        *unsafe { local_to_global.get_unchecked(cat as usize) }
                     })
                 })
                 .collect::<UInt32Vec>()
@@ -469,8 +472,8 @@ impl<'a> CategoricalChunkedBuilder<'a> {
         if using_string_cache() {
             if let RevMappingBuilder::Local(ref mut mut_arr) = self.reverse_mapping {
                 let arr: Utf8Array<_> = std::mem::take(mut_arr).into();
-                let keys = std::mem::take(&mut self.cat_builder).into();
-                self.global_map_from_local(&keys, arr);
+                let keys: UInt32Array = std::mem::take(&mut self.cat_builder).into();
+                self.global_map_from_local(keys, arr);
             }
         }
 

diff --git a/crates/polars-core/src/chunked_array/logical/categorical/from.rs b/crates/polars-core/src/chunked_array/logical/categorical/from.rs
@@ -93,7 +93,7 @@ impl CategoricalChunked {
     ) -> Self {
         if using_string_cache() {
             let mut builder = CategoricalChunkedBuilder::new(name, keys.len());
-            builder.global_map_from_local(keys, values.clone());
+            builder.global_map_from_local(keys.clone(), values.clone());
             builder.finish()
         } else {
             CategoricalChunked::from_chunks_original(

diff --git a/crates/polars-core/src/chunked_array/logical/categorical/mod.rs b/crates/polars-core/src/chunked_array/logical/categorical/mod.rs
@@ -105,11 +105,10 @@ impl CategoricalChunked {
             RevMapping::Local(categories, _) => categories,
             RevMapping::Enum(categories, _) => categories,
         };
-        let physical = self.physical().rechunk();
-        let arr = physical.downcast_get(0).unwrap();
+        let physical = self.physical();
 
         let mut builder = CategoricalChunkedBuilder::new(self.name(), physical.len());
-        builder.global_map_from_local(arr, categories.clone());
+        builder.global_map_from_local(physical, categories.clone());
         Ok(builder.finish())
     }