From 19f093954615766a55991e4a6d8ee439ab18d2f6 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Mon, 15 Apr 2024 16:36:54 +0200
Subject: [PATCH] perf: Fix regression that led to using only a single thread
 (#15667)

---
 crates/polars-core/src/frame/group_by/hashing.rs     | 2 +-
 crates/polars-core/src/frame/group_by/into_groups.rs | 8 +++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/crates/polars-core/src/frame/group_by/hashing.rs b/crates/polars-core/src/frame/group_by/hashing.rs
index 0114478cfd5a..418471abc388 100644
--- a/crates/polars-core/src/frame/group_by/hashing.rs
+++ b/crates/polars-core/src/frame/group_by/hashing.rs
@@ -124,6 +124,7 @@ where
 // have the code duplication
 pub(crate) fn group_by_threaded_slice<T, IntoSlice>(
     keys: Vec<IntoSlice>,
+    n_partitions: usize,
     sorted: bool,
 ) -> GroupsProxy
 where
@@ -131,7 +132,6 @@ where
     <T as ToTotalOrd>::TotalOrdItem: Send + Hash + Eq + Sync + Copy + DirtyHash,
     IntoSlice: AsRef<[T]> + Send + Sync,
 {
-    let n_partitions = keys.len();
     let init_size = get_init_size();
 
     // We will create a hashtable in every thread.
diff --git a/crates/polars-core/src/frame/group_by/into_groups.rs b/crates/polars-core/src/frame/group_by/into_groups.rs
index fb19d9578c8f..f7066c7b6bcd 100644
--- a/crates/polars-core/src/frame/group_by/into_groups.rs
+++ b/crates/polars-core/src/frame/group_by/into_groups.rs
@@ -36,7 +36,7 @@ where
                 .downcast_iter()
                 .map(|arr| arr.values().as_slice())
                 .collect::<Vec<_>>();
-            group_by_threaded_slice(keys, sorted)
+            group_by_threaded_slice(keys, n_partitions, sorted)
         } else {
             let keys = ca
                 .downcast_iter()
@@ -261,9 +261,10 @@ impl IntoGroupsProxy for BinaryChunked {
         let bh = self.to_bytes_hashes(multithreaded, Default::default());
 
         let out = if multithreaded {
+            let n_partitions = bh.len();
             // Take slices so that the vecs are not cloned.
             let bh = bh.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
-            group_by_threaded_slice(bh, sorted)
+            group_by_threaded_slice(bh, n_partitions, sorted)
         } else {
             group_by(bh[0].iter(), sorted)
         };
@@ -277,9 +278,10 @@ impl IntoGroupsProxy for BinaryOffsetChunked {
         let bh = self.to_bytes_hashes(multithreaded, Default::default());
 
         let out = if multithreaded {
+            let n_partitions = bh.len();
             // Take slices so that the vecs are not cloned.
             let bh = bh.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
-            group_by_threaded_slice(bh, sorted)
+            group_by_threaded_slice(bh, n_partitions, sorted)
         } else {
             group_by(bh[0].iter(), sorted)
         };