feat: Implements agg_list for NullChunked (pola-rs#15439)

cmdlineluser · Apr 2, 2024 · 802d6c8 · 802d6c8
1 parent c8640cb
commit 802d6c8
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 1 deletion.
diff --git a/crates/polars-core/src/chunked_array/builder/list/mod.rs b/crates/polars-core/src/chunked_array/builder/list/mod.rs
@@ -16,7 +16,7 @@ pub use boolean::*;
 #[cfg(feature = "dtype-categorical")]
 use categorical::*;
 use dtypes::*;
-use null::*;
+pub use null::*;
 pub use primitive::*;
 
 use super::*;

diff --git a/crates/polars-core/src/chunked_array/builder/list/null.rs b/crates/polars-core/src/chunked_array/builder/list/null.rs
@@ -18,6 +18,12 @@ impl ListNullChunkedBuilder {
         value_builder.extend_nulls(s.len());
         self.builder.try_push_valid().unwrap();
     }
+
+    pub(crate) fn append_with_len(&mut self, len: usize) {
+        let value_builder = self.builder.mut_values();
+        value_builder.extend_nulls(len);
+        self.builder.try_push_valid().unwrap();
+    }
 }
 
 impl ListBuilderTrait for ListNullChunkedBuilder {

diff --git a/crates/polars-core/src/frame/group_by/aggregations/agg_list.rs b/crates/polars-core/src/frame/group_by/aggregations/agg_list.rs
@@ -5,6 +5,8 @@ use polars_utils::unwrap::UnwrapUncheckedRelease;
 use super::*;
 #[cfg(feature = "dtype-struct")]
 use crate::chunked_array::builder::AnonymousOwnedListBuilder;
+use crate::chunked_array::builder::ListNullChunkedBuilder;
+use crate::series::implementations::null::NullChunked;
 
 pub trait AggList {
     /// # Safety
@@ -156,6 +158,27 @@ where
     }
 }
 
+impl AggList for NullChunked {
+    unsafe fn agg_list(&self, groups: &GroupsProxy) -> Series {
+        match groups {
+            GroupsProxy::Idx(groups) => {
+                let mut builder = ListNullChunkedBuilder::new(self.name(), groups.len());
+                for idx in groups.all().iter() {
+                    builder.append_with_len(idx.len());
+                }
+                builder.finish().into_series()
+            },
+            GroupsProxy::Slice { groups, .. } => {
+                let mut builder = ListNullChunkedBuilder::new(self.name(), groups.len());
+                for [_, len] in groups {
+                    builder.append_with_len(*len as usize);
+                }
+                builder.finish().into_series()
+            },
+        }
+    }
+}
+
 impl AggList for BooleanChunked {
     unsafe fn agg_list(&self, groups: &GroupsProxy) -> Series {
         match groups {

diff --git a/crates/polars-core/src/series/implementations/null.rs b/crates/polars-core/src/series/implementations/null.rs
@@ -106,6 +106,11 @@ impl PrivateSeries for NullChunked {
         })
     }
 
+    #[cfg(feature = "algorithm_group_by")]
+    unsafe fn agg_list(&self, groups: &GroupsProxy) -> Series {
+        AggList::agg_list(self, groups)
+    }
+
     fn _get_flags(&self) -> Settings {
         Settings::empty()
     }
@@ -183,6 +188,10 @@ impl SeriesTrait for NullChunked {
         NullChunked::new(self.name.clone(), self.len()).into_series()
     }
 
+    fn drop_nulls(&self) -> Series {
+        NullChunked::new(self.name.clone(), 0).into_series()
+    }
+
     fn cast(&self, data_type: &DataType) -> PolarsResult<Series> {
         Ok(Series::full_null(self.name.as_ref(), self.len(), data_type))
     }

diff --git a/py-polars/tests/unit/operations/test_window.py b/py-polars/tests/unit/operations/test_window.py
@@ -442,3 +442,10 @@ def test_window_13173() -> None:
         "val": ["2", "3"],
         "min_val_per_color": ["2", "3"],
     }
+
+
+def test_window_agg_list_null_15437() -> None:
+    df = pl.DataFrame({"a": [None]})
+    output = df.select(pl.concat_list("a").over(1))
+    expected = pl.DataFrame({"a": [[None]]})
+    assert_frame_equal(output, expected)