Adding support for struct columns in DataFrame.to_pandas (#3178)

pola-rs · Apr 19, 2022 · 6220dcd · 6220dcd
1 parent d7c9b1f
commit 6220dcd
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 2 deletions.
diff --git a/polars/polars-core/src/frame/mod.rs b/polars/polars-core/src/frame/mod.rs
@@ -556,8 +556,7 @@ impl DataFrame {
             .ok_or_else(|| {
                 PolarsError::NoData("Can not determine number of chunks if there is no data".into())
             })?
-            .chunks()
-            .len())
+            .n_chunks())
     }
 
     /// Get a reference to the schema fields of the `DataFrame`.

diff --git a/polars/polars-core/src/series/implementations/struct_.rs b/polars/polars-core/src/series/implementations/struct_.rs
@@ -87,6 +87,11 @@ impl SeriesTrait for SeriesWrap<StructChunked> {
         self.0.name()
     }
 
+    fn chunk_lengths(&self) -> ChunkIdIter {
+        let s = self.0.fields().first().unwrap();
+        s.chunk_lengths()
+    }
+
     /// Number of chunks in this Series
     fn n_chunks(&self) -> usize {
         let s = self.0.fields().first().unwrap();

diff --git a/py-polars/tests/test_struct.py b/py-polars/tests/test_struct.py
@@ -1,3 +1,5 @@
+import pandas as pd
+
 import polars as pl
 
 
@@ -141,3 +143,12 @@ def test_nested_struct() -> None:
 def test_eager_struct() -> None:
     s = pl.struct([pl.Series([1, 2, 3]), pl.Series(["a", "b", "c"])], eager=True)
     assert s.dtype == pl.Struct
+
+
+def test_struct_to_pandas() -> None:
+    df = pd.DataFrame([{"a": {"b": {"c": 2}}}])
+    pl_df = pl.from_pandas(df)
+
+    assert isinstance(pl_df.dtypes[0], pl.datatypes.Struct)
+
+    assert pl_df.to_pandas().equals(df)