Fix read_parquet with list having nested struct (#2991)

pola-rs · Mar 31, 2022 · 12044c0 · 12044c0
1 parent 6b2e7b5
commit 12044c0
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/polars/polars-core/src/datatypes.rs b/polars/polars-core/src/datatypes.rs
@@ -953,6 +953,11 @@ impl From<&ArrowDataType> for DataType {
             ArrowDataType::Time64(_) | ArrowDataType::Time32(_) => DataType::Time,
             #[cfg(feature = "dtype-categorical")]
             ArrowDataType::Dictionary(_, _, _) => DataType::Categorical(None),
+            #[cfg(feature = "dtype-struct")]
+            ArrowDataType::Struct(fields) => {
+                let fields: Vec<Field> = fields.iter().map(|fld| fld.into()).collect();
+                DataType::Struct(fields)
+            }
             ArrowDataType::Extension(name, _, _) if name == "POLARS_EXTENSION_TYPE" => {
                 #[cfg(feature = "object")]
                 {

diff --git a/py-polars/tests/io/test_parquet.py b/py-polars/tests/io/test_parquet.py
@@ -130,3 +130,18 @@ def test_parquet_datetime() -> None:
     f.seek(0)
     read = pl.read_parquet(f)
     assert read.frame_equal(df)
+
+
+def test_nested_parquet() -> None:
+    f = io.BytesIO()
+    data = [
+        {"a": [{"b": 0}]},
+        {"a": [{"b": 1}, {"b": 2}]},
+    ]
+    df = pd.DataFrame(data)
+    df.to_parquet(f)
+
+    read = pl.read_parquet(f, use_pyarrow=True)
+    assert read.columns == ["a"]
+    assert isinstance(read.dtypes[0], pl.datatypes.List)
+    assert isinstance(read.dtypes[0].inner, pl.datatypes.Struct)