struct: handle validity (#3217)

pola-rs · Apr 23, 2022 · c47e275 · c47e275
1 parent 6677344
commit c47e275
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git a/polars/polars-core/src/series/from.rs b/polars/polars-core/src/series/from.rs
@@ -379,11 +379,29 @@ impl Series {
                     chunks[0].clone()
                 };
                 let arr = convert_inner_types(&arr);
-                let struct_arr = arr.as_any().downcast_ref::<StructArray>().unwrap();
-                assert!(
-                    struct_arr.validity().is_none(),
-                    "polars struct does not support validity"
-                );
+                let mut struct_arr =
+                    std::borrow::Cow::Borrowed(arr.as_any().downcast_ref::<StructArray>().unwrap());
+
+                if let Some(validity) = struct_arr.validity() {
+                    let new_values = struct_arr
+                        .values()
+                        .iter()
+                        .map(|arr| {
+                            Arc::from(match arr.validity() {
+                                None => arr.with_validity(Some(validity.clone())),
+                                Some(arr_validity) => {
+                                    arr.with_validity(Some(arr_validity & validity))
+                                }
+                            })
+                        })
+                        .collect();
+
+                    struct_arr = std::borrow::Cow::Owned(StructArray::new(
+                        struct_arr.data_type().clone(),
+                        new_values,
+                        None,
+                    ));
+                }
                 let fields = struct_arr
                     .values()
                     .iter()

diff --git a/py-polars/tests/test_struct.py b/py-polars/tests/test_struct.py
@@ -219,3 +219,11 @@ def build_struct_df(data: list) -> DataFrame:
     assert df["struct_list_struct_col"].struct.field("list_struct_col")[0].struct.field(
         "inner"
     ).to_list() == [1]
+
+
+def test_struct_with_validity() -> None:
+    data = [{"a": {"b": 1}}, {"a": None}]
+    tbl = pa.Table.from_pylist(data)
+    df = pl.from_arrow(tbl)
+    assert isinstance(df, pl.DataFrame)
+    assert df["a"].to_list() == [{"b": 1}, {"b": None}]