fix: Fix df.chunked for struct (#16504)

pola-rs · May 26, 2024 · 805acd8 · 805acd8
1 parent af3a42f
commit 805acd8
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/crates/polars-core/src/frame/chunks.rs b/crates/polars-core/src/frame/chunks.rs
@@ -29,12 +29,17 @@ impl DataFrame {
             let columns = self
                 .get_columns()
                 .iter()
-                .map(|s| {
-                    if s.n_chunks() == 1 {
-                        s.clone()
-                    } else {
-                        s.replace_with_chunk(s.chunks()[i].clone())
-                    }
+                .map(|s| match s.dtype() {
+                    #[cfg(feature = "dtype-struct")]
+                    DataType::Struct(_) => {
+                        let mut ca = s.struct_().unwrap().clone();
+                        for field in ca.fields_mut().iter_mut() {
+                            *field = field.replace_with_chunk(field.chunks()[i].clone())
+                        }
+                        ca.update_chunks(0);
+                        ca.into_series()
+                    },
+                    _ => s.replace_with_chunk(s.chunks()[i].clone()),
                 })
                 .collect::<Vec<_>>();
 

diff --git a/py-polars/tests/unit/datatypes/test_struct.py b/py-polars/tests/unit/datatypes/test_struct.py
@@ -910,3 +910,16 @@ def test_struct_field_recognized_as_renaming_expr_16480() -> None:
 
     q = q.select("x")
     assert q.collect().to_dict(as_series=False) == {"x": [1]}
+
+
+def test_struct_filter_chunked_16498() -> None:
+    with pl.StringCache():
+        N = 5
+        df_orig1 = pl.DataFrame({"cat_a": ["remove"] * N, "cat_b": ["b"] * N})
+
+        df_orig2 = pl.DataFrame({"cat_a": ["a"] * N, "cat_b": ["b"] * N})
+
+        df = pl.concat([df_orig1, df_orig2], rechunk=False).cast(pl.Categorical)
+        df = df.select(pl.struct(pl.all()).alias("s"))
+        df = df.filter(pl.col("s").struct.field("cat_a") != pl.lit("remove"))
+        assert df.shape == (5, 1)