PyArrow: Pass in null-mask (apache#1264)

* PyArrow: Pass in null-mask * Add missing flag
HonahX · Oct 29, 2024 · 9a6a9a1 · 9a6a9a1
1 parent 58a7be3
commit 9a6a9a1
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
@@ -1807,7 +1807,11 @@ def struct(
             else:
                 raise ResolveError(f"Field is required, and could not be found in the file: {field}")
 
-        return pa.StructArray.from_arrays(arrays=field_arrays, fields=pa.struct(fields))
+        return pa.StructArray.from_arrays(
+            arrays=field_arrays,
+            fields=pa.struct(fields),
+            mask=struct_array.is_null() if isinstance(struct_array, pa.StructArray) else None,
+        )
 
     def field(self, field: NestedField, _: Optional[pa.Array], field_array: Optional[pa.Array]) -> Optional[pa.Array]:
         return field_array

diff --git a/tests/integration/test_writes/test_writes.py b/tests/integration/test_writes/test_writes.py
@@ -1450,6 +1450,47 @@ def test_rewrite_manifest_after_partition_evolution(session_catalog: Catalog) ->
     )
 
 
+@pytest.mark.integration
+def test_writing_null_structs(session_catalog: Catalog) -> None:
+    import pyarrow as pa
+
+    schema = pa.schema([
+        pa.field(
+            "struct_field_1",
+            pa.struct([
+                pa.field("string_nested_1", pa.string()),
+                pa.field("int_item_2", pa.int32()),
+                pa.field("float_item_2", pa.float32()),
+            ]),
+        ),
+    ])
+
+    records = [
+        {
+            "struct_field_1": {
+                "string_nested_1": "nest_1",
+                "int_item_2": 1234,
+                "float_item_2": 1.234,
+            },
+        },
+        {},
+    ]
+
+    try:
+        session_catalog.drop_table(
+            identifier="default.test_writing_null_structs",
+        )
+    except NoSuchTableError:
+        pass
+
+    table = session_catalog.create_table("default.test_writing_null_structs", schema)
+
+    pyarrow_table: pa.Table = pa.Table.from_pylist(records, schema=schema)
+    table.append(pyarrow_table)
+
+    assert pyarrow_table.to_pandas()["struct_field_1"].tolist() == table.scan().to_pandas()["struct_field_1"].tolist()
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("format_version", [1, 2])
 def test_abort_table_transaction_on_exception(