fix bug in pyarrow list roundtrip

pola-rs · Aug 5, 2021 · 813ce3b · 813ce3b
1 parent d29d7b6
commit 813ce3b
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/polars/polars-core/src/datatypes.rs b/polars/polars-core/src/datatypes.rs
@@ -506,9 +506,11 @@ impl DataType {
             Date32 => ArrowDataType::Date32,
             Date64 => ArrowDataType::Date64,
             Time64(tu) => ArrowDataType::Time64(tu.clone()),
-            List(dt) => {
-                ArrowDataType::List(Box::new(arrow::datatypes::Field::new("", dt.clone(), true)))
-            }
+            List(dt) => ArrowDataType::LargeList(Box::new(arrow::datatypes::Field::new(
+                "",
+                dt.clone(),
+                true,
+            ))),
             Duration(tu) => ArrowDataType::Duration(tu.clone()),
             Null => ArrowDataType::Null,
             #[cfg(feature = "object")]
@@ -635,7 +637,26 @@ impl Schema {
     }
 
     pub fn to_arrow(&self) -> ArrowSchema {
-        let fields = self.fields.iter().map(|f| f.to_arrow()).collect();
+        let fields = self
+            .fields
+            .iter()
+            .map(|f| {
+                match f.data_type() {
+                    // we must call this item, because the arrow crate names this item when creating a
+                    // schema from record batches
+                    DataType::List(dt) => ArrowField::new(
+                        f.name(),
+                        ArrowDataType::LargeList(Box::new(ArrowField::new(
+                            "item",
+                            dt.clone(),
+                            true,
+                        ))),
+                        true,
+                    ),
+                    _ => f.to_arrow(),
+                }
+            })
+            .collect();
         ArrowSchema::new(fields)
     }
 

diff --git a/py-polars/tests/test_pandas.py → py-polars/tests/test_interop.py b/py-polars/tests/test_pandas.py → py-polars/tests/test_interop.py
@@ -1,6 +1,7 @@
 import datetime
 
 import pandas as pd
+import pyarrow as pa
 
 import polars as pl
 
@@ -25,3 +26,8 @@ def test_from_pandas_datetime():
 
     # checks lazy dispatch
     pl.DataFrame([s.rename("foo")])[pl.col("foo").dt.round("hour", 2)]
+
+
+def test_arrow_list_roundtrip():
+    # https://github.com/pola-rs/polars/issues/1064
+    pl.from_arrow(pa.table({"a": [1], "b": [[1, 2]]})).to_arrow()