Add equality-delete reads and atomic upserts

pyiceberg/manifest.py

@@ -295,7 +295,7 @@ def __repr__(self) -> str:
         NestedField(
             field_id=135,
             name="equality_ids",
-            field_type=ListType(element_id=136, element_type=LongType(), element_required=True),
+            field_type=ListType(element_id=136, element_type=IntegerType(), element_required=True),
             required=False,
             doc="Field ids used to determine row equality in equality delete files.",
         ),
@@ -390,7 +390,7 @@ def __repr__(self) -> str:
         NestedField(
             field_id=135,
             name="equality_ids",
-            field_type=ListType(element_id=136, element_type=LongType(), element_required=True),
+            field_type=ListType(element_id=136, element_type=IntegerType(), element_required=True),
             required=False,
             doc="Field ids used to determine row equality in equality delete files.",
         ),
@@ -1260,11 +1260,13 @@ def __init__(
         output_file: OutputFile,
         snapshot_id: int,
         avro_compression: AvroCompressionCodec,
+        content: ManifestContent = ManifestContent.DATA,
     ):
         super().__init__(spec, schema, output_file, snapshot_id, avro_compression)
+        self._content = content
 
     def content(self) -> ManifestContent:
-        return ManifestContent.DATA
+        return self._content
 
     @property
     def version(self) -> TableVersion:
@@ -1274,10 +1276,15 @@ def version(self) -> TableVersion:
     def _meta(self) -> dict[str, str]:
         return {
             **super()._meta,
-            "content": "data",
+            "content": "data" if self._content == ManifestContent.DATA else "deletes",
         }
 
     def prepare_entry(self, entry: ManifestEntry) -> ManifestEntry:
+        entry_content = ManifestContent.DATA if entry.data_file.content == DataFileContent.DATA else ManifestContent.DELETES
+        if entry_content != self._content:
+            raise ValueError(
+                f"Cannot write {entry.data_file.content.name.lower()} file to a {self._content.name.lower()} manifest"
+            )
         if entry.sequence_number is None:
             if entry.snapshot_id is not None and entry.snapshot_id != self._snapshot_id:
                 raise ValueError(f"Found unassigned sequence number for an entry from snapshot: {entry.snapshot_id}")
@@ -1293,11 +1300,14 @@ def write_manifest(
     output_file: OutputFile,
     snapshot_id: int,
     avro_compression: AvroCompressionCodec,
+    content: ManifestContent = ManifestContent.DATA,
 ) -> ManifestWriter:
     if format_version == 1:
+        if content != ManifestContent.DATA:
+            raise ValueError("Cannot write delete manifests in a v1 table")
         return ManifestWriterV1(spec, schema, output_file, snapshot_id, avro_compression)
     elif format_version == 2:
-        return ManifestWriterV2(spec, schema, output_file, snapshot_id, avro_compression)
+        return ManifestWriterV2(spec, schema, output_file, snapshot_id, avro_compression, content)
     else:
         raise ValueError(f"Cannot write manifest for table version: {format_version}")
 

pyiceberg/table/__init__.py

@@ -989,6 +989,120 @@ def upsert(
 
         return UpsertResult(rows_updated=update_row_cnt, rows_inserted=insert_row_cnt)
 
+    def upsert_by_equality_delete(
+        self,
+        df: pa.Table,
+        join_cols: list[str] | None = None,
+        case_sensitive: bool = True,
+        branch: str | None = MAIN_BRANCH,
+        snapshot_properties: dict[str, str] = EMPTY_DICT,
+        delete_df: pa.Table | None = None,
+    ) -> None:
+        """Atomically commit source-key deletes and replacement rows in one snapshot.
+
+        ``df`` contains replacement rows: their keys are deleted and the rows are
+        appended. ``delete_df`` optionally contains delete-only rows; only its key
+        columns are written to the equality-delete file.
+
+        Snapshot visibility is atomic: readers see either the previous snapshot or
+        both the equality deletes and replacement data. Physical files are written
+        before the metadata commit, so a failed write or catalog commit can leave
+        unreferenced files for normal orphan-file maintenance.
+
+        Concurrent commits are ordered by their Iceberg snapshot sequence. A later
+        equality-delete commit can remove matching rows from an earlier concurrent
+        commit; rows committed later survive. Partitioned tables require every
+        partition source column in ``join_cols`` and currently reject evolved
+        partition specs because one partition-scoped delete cannot safely cover
+        data written under a different spec.
+        """
+        try:
+            import pyarrow as pa
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError("For writes PyArrow needs to be installed") from e
+
+        from pyiceberg.io.pyarrow import (
+            _check_pyarrow_schema_compatible,
+            _dataframe_to_data_files,
+            _dataframe_to_equality_delete_files,
+        )
+        from pyiceberg.table import upsert_util
+
+        if not isinstance(df, pa.Table):
+            raise ValueError(f"Expected pa.Table, got: {df}")
+        if delete_df is not None and not isinstance(delete_df, pa.Table):
+            raise ValueError(f"Expected delete_df to be pa.Table, got: {delete_df}")
+
+        if join_cols is None:
+            join_cols = []
+            for field_id in self.table_metadata.schema().identifier_field_ids:
+                column_name = self.table_metadata.schema().find_column_name(field_id)
+                if column_name is None:
+                    raise ValueError(f"Field ID could not be found: {field_id}")
+                join_cols.append(column_name)
+        if not join_cols:
+            raise ValueError("Join columns could not be found, please set identifier-field-ids or pass in explicitly.")
+        downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
+        _check_pyarrow_schema_compatible(
+            self.table_metadata.schema(),
+            provided_schema=df.schema,
+            downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us,
+            format_version=self.table_metadata.format_version,
+        )
+        delete_key_tables = [df.select(join_cols)]
+        if delete_df is not None:
+            _check_pyarrow_schema_compatible(
+                self.table_metadata.schema().select(*join_cols, case_sensitive=case_sensitive),
+                provided_schema=delete_df.select(join_cols).schema,
+                downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us,
+                format_version=self.table_metadata.format_version,
+            )
+            delete_key_tables.append(delete_df.select(join_cols))
+        delete_rows = pa.concat_tables(delete_key_tables)
+        if upsert_util.has_duplicate_rows(delete_rows, join_cols):
+            raise ValueError("Duplicate rows found in source datasets based on the key columns. No upsert executed")
+        if delete_rows.num_rows == 0:
+            return
+        if not self.table_metadata.spec().is_unpartitioned() and len(self.table_metadata.specs()) > 1:
+            raise NotImplementedError("Equality-delete upserts do not yet support evolved partition specs")
+
+        equality_ids = [
+            self.table_metadata.schema().find_field(column_name, case_sensitive=case_sensitive).field_id
+            for column_name in join_cols
+        ]
+        commit_uuid = uuid.uuid4()
+        counter = itertools.count(0)
+        equality_delete_files = list(
+            _dataframe_to_equality_delete_files(
+                table_metadata=self.table_metadata,
+                df=delete_rows,
+                equality_ids=equality_ids,
+                io=self._table.io,
+                write_uuid=commit_uuid,
+                counter=counter,
+            )
+        )
+        data_files = (
+            list(
+                _dataframe_to_data_files(
+                    table_metadata=self.table_metadata,
+                    df=df,
+                    io=self._table.io,
+                    write_uuid=commit_uuid,
+                    counter=counter,
+                )
+            )
+            if df.num_rows > 0
+            else []
+        )
+
+        with self.update_snapshot(snapshot_properties=snapshot_properties, branch=branch).row_delta() as row_delta:
+            row_delta.commit_uuid = commit_uuid
+            for equality_delete_file in equality_delete_files:
+                row_delta.append_data_file(equality_delete_file)
+            for data_file in data_files:
+                row_delta.append_data_file(data_file)
+
     def _find_referenced_data_files(self, file_paths: list[str]) -> list[str]:
         """Return file_paths already referenced by data files in the current snapshot."""
         snapshot = self.table_metadata.current_snapshot()
@@ -1702,6 +1816,26 @@ def upsert(
                 snapshot_properties=snapshot_properties,
             )
 
+    def upsert_by_equality_delete(
+        self,
+        df: pa.Table,
+        join_cols: list[str] | None = None,
+        case_sensitive: bool = True,
+        branch: str | None = MAIN_BRANCH,
+        snapshot_properties: dict[str, str] = EMPTY_DICT,
+        delete_df: pa.Table | None = None,
+    ) -> None:
+        """Atomically commit equality deletes and replacement rows; see the transaction API for semantics."""
+        with self.transaction() as tx:
+            tx.upsert_by_equality_delete(
+                df=df,
+                join_cols=join_cols,
+                delete_df=delete_df,
+                case_sensitive=case_sensitive,
+                branch=branch,
+                snapshot_properties=snapshot_properties,
+            )
+
     def append(
         self,
         df: pa.Table | pa.RecordBatchReader,
@@ -2255,19 +2389,13 @@ def from_rest_response(
         Returns:
             A FileScanTask with the converted data and delete files.
 
-        Raises:
-            NotImplementedError: If equality delete files are encountered.
         """
-        from pyiceberg.catalog.rest.scan_planning import RESTEqualityDeleteFile
-
         data_file = _rest_file_to_data_file(rest_task.data_file)
 
         resolved_deletes: set[DataFile] = set()
         if rest_task.delete_file_references:
             for idx in rest_task.delete_file_references:
                 delete_file = delete_files[idx]
-                if isinstance(delete_file, RESTEqualityDeleteFile):
-                    raise NotImplementedError(f"PyIceberg does not yet support equality deletes: {delete_file.file_path}")
                 resolved_deletes.add(_rest_file_to_data_file(delete_file))
 
         return FileScanTask(
@@ -2279,7 +2407,7 @@ def from_rest_response(
 
 def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile:
     """Convert a REST content file to a manifest DataFile."""
-    from pyiceberg.catalog.rest.scan_planning import RESTDataFile
+    from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTEqualityDeleteFile
 
     if isinstance(rest_file, RESTDataFile):
         column_sizes = rest_file.column_sizes.to_dict() if rest_file.column_sizes else None
@@ -2292,6 +2420,8 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile:
         null_value_counts = None
         nan_value_counts = None
 
+    equality_ids = rest_file.equality_ids if isinstance(rest_file, RESTEqualityDeleteFile) else None
+
     data_file = DataFile.from_args(
         content=DataFileContent.from_rest_type(rest_file.content),
         file_path=rest_file.file_path,
@@ -2305,6 +2435,7 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile:
         nan_value_counts=nan_value_counts,
         split_offsets=rest_file.split_offsets,
         sort_order_id=rest_file.sort_order_id,
+        equality_ids=equality_ids,
     )
     data_file.spec_id = rest_file.spec_id
     return data_file
@@ -2804,7 +2935,7 @@ def plan_files(
             List of FileScanTasks that contain both data and delete files.
         """
         data_entries: list[ManifestEntry] = []
-        delete_index = DeleteFileIndex()
+        delete_index = DeleteFileIndex(self.table_metadata.schema())
 
         residual_evaluators: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict(self._build_residual_evaluator)
 
@@ -2815,10 +2946,8 @@ def plan_files(
             data_file = manifest_entry.data_file
             if data_file.content == DataFileContent.DATA:
                 data_entries.append(manifest_entry)
-            elif data_file.content == DataFileContent.POSITION_DELETES:
+            elif data_file.content in (DataFileContent.POSITION_DELETES, DataFileContent.EQUALITY_DELETES):
                 delete_index.add_delete_file(manifest_entry, partition_key=data_file.partition)
-            elif data_file.content == DataFileContent.EQUALITY_DELETES:
-                raise ValueError("PyIceberg does not yet support equality deletes: https://github.com/apache/iceberg/issues/6568")
             else:
                 raise ValueError(f"Unknown DataFileContent ({data_file.content}): {manifest_entry}")
 
@@ -2912,6 +3041,8 @@ class WriteTask:
     record_batches: list[pa.RecordBatch]
     sort_order_id: int | None = None
     partition_key: PartitionKey | None = None
+    content: DataFileContent = DataFileContent.DATA
+    equality_ids: tuple[int, ...] | None = None
 
     def generate_data_file_filename(self, extension: str) -> str:
         # Mimics the behavior in the Java API: