fix: empty arrays in ak.to_parquet with extensionarray=True (#3234)

scikit-hep · Sep 9, 2024 · 6b39e43 · 6b39e43
1 parent dc79aa1
commit 6b39e43
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/src/awkward/_connect/pyarrow/table_conv.py b/src/awkward/_connect/pyarrow/table_conv.py
@@ -221,7 +221,7 @@ def replace_schema(table: pyarrow.Table, new_schema: pyarrow.Schema) -> pyarrow.
         new_batches.append(
             pyarrow.RecordBatch.from_arrays(arrays=columns, schema=new_schema)
         )
-    return pyarrow.Table.from_batches(new_batches)
+    return pyarrow.Table.from_batches(new_batches, schema=new_schema)
 
 
 def array_with_replacement_type(

diff --git a/tests/test_2772_parquet_extn_array_metadata.py b/tests/test_2772_parquet_extn_array_metadata.py
@@ -206,3 +206,12 @@ def test_selective_parquet(tmp_path):
     ak.to_parquet(ak_tbl, filename)
     tbl_tr = ak.from_parquet(filename, columns=["struct_array", "indexed"])
     assert to_list(tbl_tr["struct_array"]) == to_list(ak_tbl["struct_array"])
+
+
+@pytest.mark.parametrize("doit", [False, True])
+def test_empty(tmp_path, doit):
+    filename = os.path.join(tmp_path, "whatever.parquet")
+
+    ak.to_parquet(ak.Array([{"x": 1, "y": 1.1}])[0:0], filename, extensionarray=doit)
+
+    assert str(ak.from_parquet(filename).type) == "0 * {x: int64, y: float64}"