Last few bits

Author: Fokko

pyiceberg/io/pyarrow.py

@@ -1772,12 +1772,9 @@ def data_file_statistics_from_parquet_metadata(
     )
 
 
-def write_file(io: FileIO, table_metadata: TableMetadata, tasks: Iterable["WriteTask"]) -> Iterator[DataFile]:
+def write_file(io: FileIO, table_metadata: TableMetadata, tasks: Iterator[WriteTask]) -> Iterator[DataFile]:
     from pyiceberg.table import PropertyUtil, TableProperties
 
-    schema = table_metadata.schema()
-    arrow_file_schema = schema.as_arrow()
-
     parquet_writer_kwargs = _get_parquet_writer_kwargs(table_metadata.properties)
     row_group_size = PropertyUtil.property_as_int(
         properties=table_metadata.properties,

pyiceberg/manifest.py

@@ -338,7 +338,7 @@ class DataFile(Record):
     split_offsets: Optional[List[int]]
     equality_ids: Optional[List[int]]
     sort_order_id: Optional[int]
-    spec_id: Optional[int]
+    spec_id: int
 
     def __setattr__(self, name: str, value: Any) -> None:
         """Assign a key/value to a DataFile."""

pyiceberg/table/__init__.py

@@ -465,10 +465,10 @@ def delete(self, delete_filter: BooleanExpression, snapshot_properties: Dict[str
             warnings.warn("PyIceberg only supports copy on write")
 
         with self.update_snapshot(snapshot_properties=snapshot_properties).delete() as delete_snapshot:
-            delete_snapshot.delete_by_predicate(delete_filter)  # type: ignore
+            delete_snapshot.delete_by_predicate(delete_filter)
 
         # Check if there are any files that require an actual rewrite of a data file
-        if delete_snapshot.rewrites_needed is True:  # type: ignore
+        if delete_snapshot.rewrites_needed is True:
             # When we want to filter out certain rows, we want to invert the expression
             # delete id = 22 means that we want to look for that value, and then remove
             # if from the Parquet file
@@ -2767,7 +2767,7 @@ def _parquet_files_to_data_files(table_metadata: TableMetadata, file_paths: List
     yield from parquet_files_to_data_files(io=io, table_metadata=table_metadata, file_paths=iter(file_paths))
 
 
-class _MergingSnapshotProducer(UpdateTableMetadata["_MergingSnapshotProducer"]):
+class _MergingSnapshotProducer(UpdateTableMetadata[U], Generic[U]):
     commit_uuid: uuid.UUID
     _operation: Operation
     _snapshot_id: int
@@ -2798,11 +2798,11 @@ def __init__(
         self.snapshot_properties = snapshot_properties
         self._manifest_counter = itertools.count(0)
 
-    def append_data_file(self, data_file: DataFile) -> _MergingSnapshotProducer:
+    def append_data_file(self, data_file: DataFile) -> _MergingSnapshotProducer[U]:
         self._added_data_files.append(data_file)
         return self
 
-    def delete_data_file(self, data_file: DataFile) -> _MergingSnapshotProducer:
+    def delete_data_file(self, data_file: DataFile) -> _MergingSnapshotProducer[U]:
         self._deleted_data_files.add(data_file)
         return self
 
@@ -2893,7 +2893,7 @@ def _summary(self, snapshot_properties: Dict[str, str] = EMPTY_DICT) -> Summary:
             for data_file in self._deleted_data_files:
                 ssc.remove_file(
                     data_file=data_file,
-                    partition_spec=specs.get(data_file.spec_id),
+                    partition_spec=specs[data_file.spec_id],
                     schema=self._transaction.table_metadata.schema(),
                 )
 
@@ -2953,7 +2953,7 @@ def _commit(self) -> UpdatesAndRequirements:
         )
 
 
-class DeleteFiles(_MergingSnapshotProducer):
+class DeleteFiles(_MergingSnapshotProducer["DeleteFiles"]):
     """Will delete manifest entries from the current snapshot based on the predicate.
 
     This will produce a DELETE snapshot:
@@ -3102,7 +3102,7 @@ def files_affected(self) -> bool:
         return len(self._deleted_entries()) > 0
 
 
-class FastAppendFiles(_MergingSnapshotProducer):
+class FastAppendFiles(_MergingSnapshotProducer["FastAppendFiles"]):
     def _existing_manifests(self) -> List[ManifestFile]:
         """To determine if there are any existing manifest files.
 
@@ -3131,7 +3131,7 @@ def _deleted_entries(self) -> List[ManifestEntry]:
         return []
 
 
-class OverwriteFiles(_MergingSnapshotProducer):
+class OverwriteFiles(_MergingSnapshotProducer["OverwriteFiles"]):
     """Overwrites data from the table. This will produce an OVERWRITE snapshot.
 
     Data and delete files were added and removed in a logical overwrite operation.

tests/integration/test_writes/test_writes.py

@@ -185,10 +185,11 @@ def test_summaries(spark: SparkSession, session_catalog: Catalog, arrow_table_wi
     ).collect()
 
     operations = [row.operation for row in rows]
-    assert operations == ['append', 'append', 'overwrite']
+    assert operations == ['append', 'append', 'delete', 'overwrite']
 
     summaries = [row.summary for row in rows]
 
+    # Append
     assert summaries[0] == {
         'added-data-files': '1',
         'added-files-size': '5459',
@@ -201,6 +202,7 @@ def test_summaries(spark: SparkSession, session_catalog: Catalog, arrow_table_wi
         'total-records': '3',
     }
 
+    # Append
     assert summaries[1] == {
         'added-data-files': '1',
         'added-files-size': '5459',
@@ -213,13 +215,24 @@ def test_summaries(spark: SparkSession, session_catalog: Catalog, arrow_table_wi
         'total-records': '6',
     }
 
+    # Delete
     assert summaries[2] == {
-        'added-data-files': '1',
-        'added-files-size': '5459',
-        'added-records': '3',
         'deleted-data-files': '2',
         'deleted-records': '6',
         'removed-files-size': '10918',
+        'total-data-files': '0',
+        'total-delete-files': '0',
+        'total-equality-deletes': '0',
+        'total-files-size': '0',
+        'total-position-deletes': '0',
+        'total-records': '0',
+    }
+
+    # Overwrite
+    assert summaries[3] == {
+        'added-data-files': '1',
+        'added-files-size': '5459',
+        'added-records': '3',
         'total-data-files': '1',
         'total-delete-files': '0',
         'total-equality-deletes': '0',
@@ -247,9 +260,9 @@ def test_data_files(spark: SparkSession, session_catalog: Catalog, arrow_table_w
     """
     ).collect()
 
-    assert [row.added_data_files_count for row in rows] == [1, 1, 0, 1, 1]
-    assert [row.existing_data_files_count for row in rows] == [0, 0, 0, 0, 0]
-    assert [row.deleted_data_files_count for row in rows] == [0, 0, 1, 0, 0]
+    assert [row.added_data_files_count for row in rows] == [1, 0, 1, 0, 1, 1]
+    assert [row.existing_data_files_count for row in rows] == [0, 0, 0, 0, 0, 0]
+    assert [row.deleted_data_files_count for row in rows] == [0, 1, 0, 1, 0, 0]
 
 
 @pytest.mark.integration
@@ -476,7 +489,7 @@ def test_summaries_with_only_nulls(
     ).collect()
 
     operations = [row.operation for row in rows]
-    assert operations == ['append', 'append', 'overwrite']
+    assert operations == ['append', 'append', 'delete', 'overwrite']
 
     summaries = [row.summary for row in rows]
 
@@ -502,14 +515,23 @@ def test_summaries_with_only_nulls(
     }
 
     assert summaries[2] == {
+        'deleted-data-files': '1',
+        'deleted-records': '2',
         'removed-files-size': '4239',
+        'total-data-files': '0',
+        'total-delete-files': '0',
         'total-equality-deletes': '0',
+        'total-files-size': '0',
         'total-position-deletes': '0',
-        'deleted-data-files': '1',
+        'total-records': '0',
+    }
+
+    assert summaries[3] == {
+        'total-data-files': '0',
         'total-delete-files': '0',
+        'total-equality-deletes': '0',
         'total-files-size': '0',
-        'deleted-records': '2',
-        'total-data-files': '0',
+        'total-position-deletes': '0',
         'total-records': '0',
     }
 
@@ -731,13 +753,14 @@ def test_inspect_snapshots(
         assert isinstance(snapshot_id.as_py(), int)
 
     assert df['parent_id'][0].as_py() is None
-    assert df['parent_id'][1:] == df['snapshot_id'][:2]
+    assert df['parent_id'][1:].to_pylist() == df['snapshot_id'][:-1].to_pylist()
 
-    assert [operation.as_py() for operation in df['operation']] == ['append', 'overwrite', 'append']
+    assert [operation.as_py() for operation in df['operation']] == ['append', 'delete', 'overwrite', 'append']
 
     for manifest_list in df['manifest_list']:
         assert manifest_list.as_py().startswith("s3://")
 
+    # Append
     assert df['summary'][0].as_py() == [
         ('added-files-size', '5459'),
         ('added-data-files', '1'),
@@ -750,6 +773,19 @@ def test_inspect_snapshots(
         ('total-equality-deletes', '0'),
     ]
 
+    # Delete
+    assert df['summary'][1].as_py() == [
+        ('removed-files-size', '5459'),
+        ('deleted-data-files', '1'),
+        ('deleted-records', '3'),
+        ('total-data-files', '0'),
+        ('total-delete-files', '0'),
+        ('total-records', '0'),
+        ('total-files-size', '0'),
+        ('total-position-deletes', '0'),
+        ('total-equality-deletes', '0'),
+    ]
+
     lhs = spark.table(f"{identifier}.snapshots").toPandas()
     rhs = df.to_pandas()
     for column in df.column_names: