fix(rust, python): fix generic streaming groupby on logical types (po…

…la-rs#5752)
zundertj · Jan 7, 2023 · 300ad43 · 300ad43
1 parent ab7fc17
commit 300ad43
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/polars/polars-core/src/frame/row.rs b/polars/polars-core/src/frame/row.rs
@@ -500,7 +500,7 @@ impl<'a> AnyValueBuffer<'a> {
         }
     }
 
-    pub fn new(dtype: &DataType, capacity: usize) -> AnyValueBuffer<'_> {
+    pub fn new(dtype: &DataType, capacity: usize) -> AnyValueBuffer<'a> {
         (dtype, capacity).into()
     }
 }

diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs
@@ -147,7 +147,7 @@ impl GenericGroupbySink {
                             .output_schema
                             .iter_dtypes()
                             .take(n_keys)
-                            .map(|dtype| AnyValueBuffer::new(dtype, agg_map.len()))
+                            .map(|dtype| AnyValueBuffer::new(&dtype.to_physical(), agg_map.len()))
                             .collect::<Vec<_>>();
                         let dtypes = agg_fns
                             .iter()
@@ -251,15 +251,17 @@ impl Sink for GenericGroupbySink {
             let current_key_values = unsafe { self.keys.get_unchecked_release_mut(partition) };
 
             let entry = current_partition.raw_entry_mut().from_hash(h, |key| {
-                let idx = key.idx as usize;
-                if self.keys_series.len() > 1 {
-                    current_tuple.iter().enumerate().all(|(i, key)| unsafe {
-                        current_key_values.get_unchecked_release(i + idx) == key
-                    })
-                } else {
-                    unsafe {
-                        current_key_values.get_unchecked_release(idx)
-                            == current_tuple.get_unchecked_release(0)
+                key.hash == h && {
+                    let idx = key.idx as usize;
+                    if self.keys_series.len() > 1 {
+                        current_tuple.iter().enumerate().all(|(i, key)| unsafe {
+                            current_key_values.get_unchecked_release(i + idx) == key
+                        })
+                    } else {
+                        unsafe {
+                            current_key_values.get_unchecked_release(idx)
+                                == current_tuple.get_unchecked_release(0)
+                        }
                     }
                 }
             });