add list groupby feature

polars/Cargo.toml

@@ -83,6 +83,7 @@ decompress = ["polars-io/decompress"]
 mode = ["polars-core/mode", "polars-lazy/mode"]
 take_opt_iter = ["polars-core/take_opt_iter"]
 extract_jsonpath = ["polars-core/extract_jsonpath", "polars-core/strings"]
+groupby_list = ["polars-core/groupby_list"]
 
 # don't use this
 private = []

polars/polars-core/Cargo.toml

@@ -63,6 +63,8 @@ reinterpret = []
 take_opt_iter = []
 mode = []
 extract_jsonpath = ["serde_json", "jsonpath_lib"]
+# allow groupby operation on list type
+groupby_list = []
 
 
 # opt-in datatypes for Series

polars/polars-core/src/frame/groupby/mod.rs

@@ -2,6 +2,8 @@ use self::hashing::*;
 use crate::chunked_array::builder::PrimitiveChunkedBuilder;
 use crate::frame::select::Selection;
 use crate::prelude::*;
+#[cfg(feature = "groupby_list")]
+use crate::utils::Wrap;
 use crate::utils::{accumulate_dataframes_vertical, set_partition_size, split_ca, NoNull};
 use crate::vector_hasher::{AsU64, StrHash};
 use crate::POOL;
@@ -169,7 +171,13 @@ impl IntoGroupTuples for CategoricalChunked {
     }
 }
 
-impl IntoGroupTuples for ListChunked {}
+impl IntoGroupTuples for ListChunked {
+    #[cfg(feature = "groupby_list")]
+    fn group_tuples(&self, _multithreaded: bool) -> GroupTuples {
+        groupby(self.into_iter().map(|opt_s| opt_s.map(Wrap)))
+    }
+}
+
 #[cfg(feature = "object")]
 impl<T> IntoGroupTuples for ObjectChunked<T>
 where

polars/polars-core/src/series/mod.rs

@@ -8,14 +8,20 @@ pub mod implementations;
 pub(crate) mod iterator;
 
 use crate::chunked_array::{builder::get_list_builder, ChunkIdIter};
+#[cfg(feature = "groupby_list")]
+use crate::utils::Wrap;
 use crate::utils::{split_ca, split_series};
 use crate::{series::arithmetic::coerce_lhs_rhs, POOL};
+#[cfg(feature = "groupby_list")]
+use ahash::RandomState;
 use arrow::compute::cast;
 use itertools::Itertools;
 use num::NumCast;
 use rayon::prelude::*;
 use std::any::Any;
 use std::convert::TryFrom;
+#[cfg(feature = "groupby_list")]
+use std::hash::{Hash, Hasher};
 use std::ops::Deref;
 use std::sync::Arc;
 
@@ -1163,6 +1169,25 @@ impl<'a> (dyn SeriesTrait + 'a) {
 #[derive(Clone)]
 pub struct Series(pub Arc<dyn SeriesTrait>);
 
+#[cfg(feature = "groupby_list")]
+impl PartialEq for Wrap<Series> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.series_equal_missing(other)
+    }
+}
+
+#[cfg(feature = "groupby_list")]
+impl Eq for Wrap<Series> {}
+
+#[cfg(feature = "groupby_list")]
+impl Hash for Wrap<Series> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        let rs = RandomState::with_seeds(0, 0, 0, 0);
+        let h = UInt64Chunked::new_from_aligned_vec("", self.0.vec_hash(rs)).sum();
+        h.hash(state)
+    }
+}
+
 impl Series {
     pub(crate) fn get_inner_mut(&mut self) -> &mut dyn SeriesTrait {
         if Arc::weak_count(&self.0) + Arc::strong_count(&self.0) != 1 {

polars/polars-core/src/utils.rs

@@ -11,6 +11,7 @@ use rayon::prelude::*;
 use std::borrow::Cow;
 use std::ops::{Deref, DerefMut};
 
+#[repr(transparent)]
 pub struct Wrap<T>(pub T);
 
 impl<T> Deref for Wrap<T> {

polars/src/lib.rs

@@ -118,6 +118,7 @@
 //!     - `downsample` - [downsample operation](crate::frame::DataFrame::downsample) on `DataFrame`s
 //!     - `asof_join` - Join as of, to join on nearest keys instead of exact equality match.
 //!     - `cross_join` - Create the cartesian product of two DataFrames.
+//!     - `groupby_list` - Allow groupby operation on keys of type List.
 //! * `Series` operations:
 //!     - `is_in` - [Check for membership in `Series`](crate::chunked_array::ops::IsIn)
 //!     - `zip_with` - [Zip two Series/ ChunkedArrays](crate::chunked_array::ops::ChunkZip)