perf(rust, python): specialized utf8 groupby in streaming

polars/polars-core/src/datatypes/any_value.rs

@@ -609,7 +609,7 @@ impl<'a> From<AnyValue<'a>> for Option<i64> {
 }
 
 impl PartialEq for AnyValue<'_> {
-    // Everything of Any is slow. Don't use.
+    #[inline]
     fn eq(&self, other: &Self) -> bool {
         use AnyValue::*;
         match (self.as_borrowed(), other.as_borrowed()) {

polars/polars-lazy/polars-pipe/Cargo.toml

@@ -18,6 +18,7 @@ polars-io = { version = "0.25.1", path = "../../polars-io", default-features = f
 polars-plan = { version = "0.25.1", path = "../polars-plan", default-features = false, features = ["compile"] }
 polars-utils = { version = "0.25.1", path = "../../polars-utils" }
 rayon.workspace = true
+smartstring = { version = "1" }
 
 [features]
 compile = []

polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs

@@ -23,14 +23,14 @@ use crate::operators::{DataChunk, FinalizedSink, PExecutionContext, Sink, SinkRe
 
 // This is the hash and the Index offset in the linear buffer
 #[derive(Copy, Clone)]
-struct Key {
-    hash: u64,
-    idx: IdxSize,
+pub(super) struct Key {
+    pub(super) hash: u64,
+    pub(super) idx: IdxSize,
 }
 
 impl Key {
     #[inline]
-    fn new(hash: u64, idx: IdxSize) -> Self {
+    pub(super) fn new(hash: u64, idx: IdxSize) -> Self {
         Self { hash, idx }
     }
 }
@@ -237,13 +237,13 @@ impl Sink for GenericGroupbySink {
 
         // a small buffer that holds the current key values
         // if we groupby 2 keys, this holds 2 anyvalues.
-        let mut current_keys_buf = Vec::with_capacity(self.keys.len());
+        let mut current_tuple = Vec::with_capacity(self.keys.len());
 
         for &h in &self.hashes {
             // load the keys in the buffer
-            current_keys_buf.clear();
+            current_tuple.clear();
             for key_iter in key_iters.iter_mut() {
-                unsafe { current_keys_buf.push(key_iter.next().unwrap_unchecked_release()) }
+                unsafe { current_tuple.push(key_iter.next().unwrap_unchecked_release()) }
             }
 
             let partition = hash_to_partition(h, self.pre_agg_partitions.len());
@@ -256,13 +256,13 @@ impl Sink for GenericGroupbySink {
             let entry = current_partition.raw_entry_mut().from_hash(h, |key| {
                 let idx = key.idx as usize;
                 if self.keys_series.len() > 1 {
-                    current_keys_buf.iter().enumerate().all(|(i, key)| unsafe {
+                    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_keys_buf.get_unchecked_release(0)
+                            == current_tuple.get_unchecked_release(0)
                     }
                 }
             });
@@ -282,7 +282,7 @@ impl Sink for GenericGroupbySink {
 
                     unsafe {
                         current_key_values.extend(
-                            current_keys_buf
+                            current_tuple
                                 .iter()
                                 .map(|av| av.clone().into_static().unwrap_unchecked_release()),
                         )

polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/mod.rs

@@ -1,7 +1,9 @@
 pub(crate) mod aggregates;
 mod generic;
 mod primitive;
+mod string;
 mod utils;
 
 pub(crate) use generic::*;
 pub(crate) use primitive::*;
+pub(crate) use string::*;

polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/primitive.rs

@@ -109,7 +109,6 @@ where
     }
 
     fn pre_finalize(&mut self) -> PolarsResult<Vec<DataFrame>> {
-        // TODO! parallel
         let mut aggregators = std::mem::take(&mut self.aggregators);
         let slices = compute_slices(&self.pre_agg_partitions, self.slice);