reduce memory usage of multiple key hash

pola-rs · Jun 24, 2021 · 9fb25c7 · 9fb25c7
1 parent 0cb5f44
commit 9fb25c7
Show file tree

Hide file tree

Showing 8 changed files with 210 additions and 110 deletions.
diff --git a/polars/polars-arrow/src/vec.rs b/polars/polars-arrow/src/vec.rs
@@ -169,6 +169,10 @@ impl<T> AlignedVec<T> {
         self.inner.as_mut_slice()
     }
 
+    pub fn as_slice(&self) -> &[T] {
+        self.inner.as_slice()
+    }
+
     pub fn capacity(&self) -> usize {
         self.inner.capacity()
     }

diff --git a/polars/polars-core/src/chunked_array/ops/apply.rs b/polars/polars-core/src/chunked_array/ops/apply.rs
@@ -125,6 +125,23 @@ where
             .map(f)
             .collect()
     }
+    fn apply_to_slice<F, V>(&'a self, f: F, slice: &mut [V])
+    where
+        F: Fn(Option<T::Native>, &V) -> V,
+    {
+        assert!(slice.len() >= self.len());
+
+        let mut idx = 0;
+        self.downcast_iter().for_each(|arr| {
+            arr.into_iter().for_each(|opt_val| {
+                // Safety:
+                // length asserted above
+                let item = unsafe { slice.get_unchecked_mut(idx) };
+                *item = f(opt_val, item);
+                idx += 1;
+            })
+        });
+    }
 }
 
 impl<'a> ChunkApply<'a, bool, bool> for BooleanChunked {
@@ -180,6 +197,24 @@ impl<'a> ChunkApply<'a, bool, bool> for BooleanChunked {
     {
         self.into_iter().enumerate().map(f).collect()
     }
+
+    fn apply_to_slice<F, T>(&'a self, f: F, slice: &mut [T])
+    where
+        F: Fn(Option<bool>, &T) -> T,
+    {
+        assert!(slice.len() >= self.len());
+
+        let mut idx = 0;
+        self.downcast_iter().for_each(|arr| {
+            arr.into_iter().for_each(|opt_val| {
+                // Safety:
+                // length asserted above
+                let item = unsafe { slice.get_unchecked_mut(idx) };
+                *item = f(opt_val, item);
+                idx += 1;
+            })
+        });
+    }
 }
 
 impl<'a> ChunkApply<'a, &'a str, Cow<'a, str>> for Utf8Chunked {
@@ -246,6 +281,24 @@ impl<'a> ChunkApply<'a, &'a str, Cow<'a, str>> for Utf8Chunked {
     {
         self.into_iter().enumerate().map(f).collect()
     }
+
+    fn apply_to_slice<F, T>(&'a self, f: F, slice: &mut [T])
+    where
+        F: Fn(Option<&'a str>, &T) -> T,
+    {
+        assert!(slice.len() >= self.len());
+
+        let mut idx = 0;
+        self.downcast_iter().for_each(|arr| {
+            arr.into_iter().for_each(|opt_val| {
+                // Safety:
+                // length asserted above
+                let item = unsafe { slice.get_unchecked_mut(idx) };
+                *item = f(opt_val, item);
+                idx += 1;
+            })
+        });
+    }
 }
 
 impl ChunkApplyKernel<BooleanArray> for BooleanChunked {
@@ -396,4 +449,24 @@ impl<'a> ChunkApply<'a, Series, Series> for ListChunked {
     {
         self.into_iter().enumerate().map(f).collect()
     }
+
+    fn apply_to_slice<F, T>(&'a self, f: F, slice: &mut [T])
+    where
+        F: Fn(Option<Series>, &T) -> T,
+    {
+        assert!(slice.len() >= self.len());
+
+        let mut idx = 0;
+        self.downcast_iter().for_each(|arr| {
+            arr.iter().for_each(|opt_val| {
+                let opt_val = opt_val.map(|arrayref| Series::try_from(("", arrayref)).unwrap());
+
+                // Safety:
+                // length asserted above
+                let item = unsafe { slice.get_unchecked_mut(idx) };
+                *item = f(opt_val, item);
+                idx += 1;
+            })
+        });
+    }
 }
diff --git a/polars/polars-core/src/chunked_array/ops/mod.rs b/polars/polars-core/src/chunked_array/ops/mod.rs
@@ -474,6 +474,12 @@ pub trait ChunkApply<'a, A, B> {
     fn apply_with_idx_on_opt<F>(&'a self, f: F) -> Self
     where
         F: Fn((usize, Option<A>)) -> Option<B> + Copy;
+
+    /// Apply a closure elementwise and write results to a mutable slice.
+    fn apply_to_slice<F, T>(&'a self, f: F, slice: &mut [T])
+    // (value of chunkedarray, value of slice) -> value of slice
+    where
+        F: Fn(Option<A>, &T) -> T;
 }
 
 /// Aggregation operations

diff --git a/polars/polars-core/src/frame/groupby/mod.rs b/polars/polars-core/src/frame/groupby/mod.rs
@@ -320,7 +320,18 @@ impl DataFrame {
             by.iter()
                 .map(|s| match s.dtype() {
                     DataType::Categorical => s.cast::<UInt32Type>().unwrap(),
-                    _ => s.clone(),
+                    DataType::Float32 => s.bit_repr_small().into_series(),
+                    // otherwise we use the vec hash for float
+                    #[cfg(feature = "dtype-u64")]
+                    DataType::Float64 => s.bit_repr_large().into_series(),
+                    _ => {
+                        // is date like
+                        if !s.is_numeric() && s.is_numeric_physical() {
+                            s.to_physical_repr()
+                        } else {
+                            s.clone()
+                        }
+                    }
                 })
                 .collect(),
         )?;

diff --git a/polars/polars-core/src/series/implementations/dates.rs b/polars/polars-core/src/series/implementations/dates.rs
@@ -119,7 +119,7 @@ macro_rules! impl_dyn_series {
                 try_physical_dispatch!(self, zip_with_same_type, mask, other)
             }
 
-            fn vec_hash(&self, random_state: RandomState) -> UInt64Chunked {
+            fn vec_hash(&self, random_state: RandomState) -> AlignedVec<u64> {
                 cast_and_apply!(self, vec_hash, random_state)
             }
 

diff --git a/polars/polars-core/src/series/implementations/mod.rs b/polars/polars-core/src/series/implementations/mod.rs
@@ -78,10 +78,14 @@ macro_rules! impl_dyn_series {
                 (&self.0).into_partial_eq_inner()
             }
 
-            fn vec_hash(&self, random_state: RandomState) -> UInt64Chunked {
+            fn vec_hash(&self, random_state: RandomState) -> AlignedVec<u64> {
                 self.0.vec_hash(random_state)
             }
 
+            fn vec_hash_combine(&self, build_hasher: RandomState, hashes: &mut [u64]) {
+                self.0.vec_hash_combine(build_hasher, hashes)
+            }
+
             fn agg_mean(&self, groups: &[(u32, Vec<u32>)]) -> Option<Series> {
                 self.0.agg_mean(groups)
             }

diff --git a/polars/polars-core/src/series/mod.rs b/polars/polars-core/src/series/mod.rs
@@ -61,7 +61,10 @@ pub(crate) mod private {
         fn into_partial_eq_inner<'a>(&'a self) -> Box<dyn PartialEqInner + 'a> {
             unimplemented!()
         }
-        fn vec_hash(&self, _build_hasher: RandomState) -> UInt64Chunked {
+        fn vec_hash(&self, _build_hasher: RandomState) -> AlignedVec<u64> {
+            unimplemented!()
+        }
+        fn vec_hash_combine(&self, _build_hasher: RandomState, _hashes: &mut [u64]) {
             unimplemented!()
         }
         fn agg_mean(&self, _groups: &[(u32, Vec<u32>)]) -> Option<Series> {