refactor(rust): Improve new-streaming groupby performance for high cardinality

crates/polars-compute/src/lib.rs

@@ -1,6 +1,4 @@
 #![cfg_attr(feature = "simd", feature(portable_simd))]
-#![cfg_attr(feature = "simd", feature(core_intrinsics))] // For fadd_algebraic.
-#![cfg_attr(feature = "simd", allow(internal_features))]
 #![cfg_attr(feature = "simd", feature(avx512_target_feature))]
 #![cfg_attr(
     all(feature = "simd", target_arch = "x86_64"),

crates/polars-compute/src/var_cov.rs

@@ -17,25 +17,10 @@
 use arrow::array::{Array, PrimitiveArray};
 use arrow::types::NativeType;
 use num_traits::AsPrimitive;
+use polars_utils::algebraic_ops::*;
 
 const CHUNK_SIZE: usize = 128;
 
-#[inline(always)]
-fn alg_add(a: f64, b: f64) -> f64 {
-    #[cfg(feature = "simd")]
-    {
-        std::intrinsics::fadd_algebraic(a, b)
-    }
-    #[cfg(not(feature = "simd"))]
-    {
-        a + b
-    }
-}
-
-fn alg_sum(it: impl IntoIterator<Item = f64>) -> f64 {
-    it.into_iter().fold(0.0, alg_add)
-}
-
 #[derive(Default, Clone)]
 pub struct VarState {
     weight: f64,
@@ -68,11 +53,11 @@ impl VarState {
         }
 
         let weight = x.len() as f64;
-        let mean = alg_sum(x.iter().copied()) / weight;
+        let mean = alg_sum_f64(x.iter().copied()) / weight;
         Self {
             weight,
             mean,
-            dp: alg_sum(x.iter().map(|&xi| (xi - mean) * (xi - mean))),
+            dp: alg_sum_f64(x.iter().map(|&xi| (xi - mean) * (xi - mean))),
         }
     }
 
@@ -119,13 +104,13 @@ impl CovState {
 
         let weight = x.len() as f64;
         let inv_weight = 1.0 / weight;
-        let mean_x = alg_sum(x.iter().copied()) * inv_weight;
-        let mean_y = alg_sum(y.iter().copied()) * inv_weight;
+        let mean_x = alg_sum_f64(x.iter().copied()) * inv_weight;
+        let mean_y = alg_sum_f64(y.iter().copied()) * inv_weight;
         Self {
             weight,
             mean_x,
             mean_y,
-            dp_xy: alg_sum(
+            dp_xy: alg_sum_f64(
                 x.iter()
                     .zip(y)
                     .map(|(&xi, &yi)| (xi - mean_x) * (yi - mean_y)),
@@ -168,15 +153,15 @@ impl PearsonState {
 
         let weight = x.len() as f64;
         let inv_weight = 1.0 / weight;
-        let mean_x = alg_sum(x.iter().copied()) * inv_weight;
-        let mean_y = alg_sum(y.iter().copied()) * inv_weight;
+        let mean_x = alg_sum_f64(x.iter().copied()) * inv_weight;
+        let mean_y = alg_sum_f64(y.iter().copied()) * inv_weight;
         let mut dp_xx = 0.0;
         let mut dp_xy = 0.0;
         let mut dp_yy = 0.0;
         for (xi, yi) in x.iter().zip(y.iter()) {
-            dp_xx = alg_add(dp_xx, (xi - mean_x) * (xi - mean_x));
-            dp_xy = alg_add(dp_xy, (xi - mean_x) * (yi - mean_y));
-            dp_yy = alg_add(dp_yy, (yi - mean_y) * (yi - mean_y));
+            dp_xx = alg_add_f64(dp_xx, (xi - mean_x) * (xi - mean_x));
+            dp_xy = alg_add_f64(dp_xy, (xi - mean_x) * (yi - mean_y));
+            dp_yy = alg_add_f64(dp_yy, (yi - mean_y) * (yi - mean_y));
         }
         Self {
             weight,

crates/polars-expr/src/groups/mod.rs

@@ -3,15 +3,20 @@ use std::path::Path;
 
 use polars_core::prelude::*;
 use polars_utils::aliases::PlRandomState;
+use polars_utils::cardinality_sketch::CardinalitySketch;
+use polars_utils::hashing::HashPartitioner;
 use polars_utils::IdxSize;
 
 mod row_encoded;
 
 /// A Grouper maps keys to groups, such that duplicate keys map to the same group.
-pub trait Grouper: Any + Send {
+pub trait Grouper: Any + Send + Sync {
     /// Creates a new empty Grouper similar to this one.
     fn new_empty(&self) -> Box<dyn Grouper>;
 
+    /// Reserves space for the given number additional of groups.
+    fn reserve(&mut self, additional: usize);
+
     /// Returns the number of groups in this Grouper.
     fn num_groups(&self) -> IdxSize;
 
@@ -23,29 +28,34 @@ pub trait Grouper: Any + Send {
     /// the ith group of other now has group index group_idxs[i] in self.
     fn combine(&mut self, other: &dyn Grouper, group_idxs: &mut Vec<IdxSize>);
 
-    /// Partitions this Grouper into the given number of partitions.
+    /// Adds the given Grouper into this one, mutating groups_idxs such that
+    /// the group subset[i] of other now has group index group_idxs[i] in self.
     ///
-    /// Updates partition_idxs such that the ith group of self moves to partition
-    /// partition_idxs[i].
+    /// # Safety
+    /// For all i, subset[i] < other.len().
+    unsafe fn gather_combine(
+        &mut self,
+        other: &dyn Grouper,
+        subset: &[IdxSize],
+        group_idxs: &mut Vec<IdxSize>,
+    );
+
+    /// Generate partition indices.
     ///
-    /// It is guaranteed that two equal keys in two independent partition_into
-    /// calls map to the same partition index if the seed and the number of
-    /// partitions is equal.
-    fn partition(
+    /// After this function partitions_idxs[i] will contain the indices for
+    /// partition i, and sketches[i] will contain a cardinality sketch for
+    /// partition i.
+    fn gen_partition_idxs(
         &self,
-        seed: u64,
-        num_partitions: usize,
-        partition_idxs: &mut Vec<IdxSize>,
-    ) -> Vec<Box<dyn Grouper>>;
+        partitioner: &HashPartitioner,
+        partition_idxs: &mut [Vec<IdxSize>],
+        sketches: &mut [CardinalitySketch],
+    );
 
     /// Returns the keys in this Grouper in group order, that is the key for
     /// group i is returned in row i.
     fn get_keys_in_group_order(&self) -> DataFrame;
 
-    /// Returns the keys in this Grouper, mutating group_idxs such that the ith
-    /// key returned corresponds to group group_idxs[i].
-    fn get_keys_groups(&self, group_idxs: &mut Vec<IdxSize>) -> DataFrame;
-
     /// Stores this Grouper at the given path.
     fn store_ooc(&self, _path: &Path) {
         unimplemented!();