Add DeleteFileIndex support for EqualityDeletes

Author: rambleraptor

tests/table/test_delete_file_index.py

@@ -24,14 +24,21 @@
 from pyiceberg.types import IntegerType, NestedField
 
 
-def _create_data_file(file_path: str = "s3://bucket/data.parquet", spec_id: int = 0) -> DataFile:
+def _create_data_file(
+    file_path: str = "s3://bucket/data.parquet",
+    spec_id: int = 0,
+    lower_bounds: dict[int, bytes] | None = None,
+    upper_bounds: dict[int, bytes] | None = None,
+) -> DataFile:
     data_file = DataFile.from_args(
         content=DataFileContent.DATA,
         file_path=file_path,
         file_format=FileFormat.PARQUET,
         partition=Record(),
         record_count=100,
         file_size_in_bytes=1000,
+        lower_bounds=lower_bounds,
+        upper_bounds=upper_bounds,
     )
     data_file._spec_id = spec_id
     return data_file
@@ -84,6 +91,27 @@ def _create_deletion_vector(
     return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file)
 
 
+def _create_equality_delete(
+    sequence_number: int = 1,
+    spec_id: int = 0,
+    lower_bounds: dict[int, bytes] | None = None,
+    upper_bounds: dict[int, bytes] | None = None,
+) -> ManifestEntry:
+    delete_file = DataFile.from_args(
+        content=DataFileContent.EQUALITY_DELETES,
+        file_path=f"s3://bucket/eq-delete-{sequence_number}.parquet",
+        file_format=FileFormat.PARQUET,
+        partition=Record(),
+        record_count=10,
+        file_size_in_bytes=100,
+        equality_ids=[1],
+        lower_bounds=lower_bounds,
+        upper_bounds=upper_bounds,
+    )
+    delete_file._spec_id = spec_id
+    return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file)
+
+
 def test_empty_index() -> None:
     index = DeleteFileIndex()
     data_file = _create_data_file()
@@ -236,46 +264,35 @@ def test_equality_delete_metrics_filtering() -> None:
     schema = Schema(NestedField(1, "id", IntegerType(), required=True))
     index = DeleteFileIndex(schema=schema)
 
-    def _create_data_file_with_metrics(file_path: str, lower: int, upper: int) -> DataFile:
-        data_file = DataFile.from_args(
-            content=DataFileContent.DATA,
-            file_path=file_path,
-            file_format=FileFormat.PARQUET,
-            partition=Record(),
-            record_count=100,
-            file_size_in_bytes=1000,
-            lower_bounds={1: to_bytes(IntegerType(), lower)},
-            upper_bounds={1: to_bytes(IntegerType(), upper)},
-        )
-        data_file._spec_id = 0
-        return data_file
-
-    def _create_equality_delete_with_metrics(sequence_number: int, lower: int, upper: int) -> ManifestEntry:
-        delete_file = DataFile.from_args(
-            content=DataFileContent.EQUALITY_DELETES,
-            file_path=f"s3://bucket/eq-delete-{sequence_number}.parquet",
-            file_format=FileFormat.PARQUET,
-            partition=Record(),
-            record_count=10,
-            file_size_in_bytes=100,
-            equality_ids=[1],
-            lower_bounds={1: to_bytes(IntegerType(), lower)},
-            upper_bounds={1: to_bytes(IntegerType(), upper)},
-        )
-        delete_file._spec_id = 0
-        return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file)
-
     # Equality delete for rows where id is between 10 and 20
-    index.add_delete_file(_create_equality_delete_with_metrics(sequence_number=100, lower=10, upper=20))
+    index.add_delete_file(
+        _create_equality_delete(
+            sequence_number=100,
+            lower_bounds={1: to_bytes(IntegerType(), 10)},
+            upper_bounds={1: to_bytes(IntegerType(), 20)},
+        )
+    )
 
     # Data file with id between 0 and 5 (no overlap)
-    file_no_overlap = _create_data_file_with_metrics("s3://bucket/no_overlap.parquet", 0, 5)
+    file_no_overlap = _create_data_file(
+        "s3://bucket/no_overlap.parquet",
+        lower_bounds={1: to_bytes(IntegerType(), 0)},
+        upper_bounds={1: to_bytes(IntegerType(), 5)},
+    )
     assert len(index.for_data_file(1, file_no_overlap)) == 0
 
     # Data file with id between 15 and 25 (overlap)
-    file_overlap = _create_data_file_with_metrics("s3://bucket/overlap.parquet", 15, 25)
+    file_overlap = _create_data_file(
+        "s3://bucket/overlap.parquet",
+        lower_bounds={1: to_bytes(IntegerType(), 15)},
+        upper_bounds={1: to_bytes(IntegerType(), 25)},
+    )
     assert len(index.for_data_file(1, file_overlap)) == 1
 
     # Data file with id between 25 and 30 (no overlap)
-    file_no_overlap_2 = _create_data_file_with_metrics("s3://bucket/no_overlap_2.parquet", 25, 30)
+    file_no_overlap_2 = _create_data_file(
+        "s3://bucket/no_overlap_2.parquet",
+        lower_bounds={1: to_bytes(IntegerType(), 25)},
+        upper_bounds={1: to_bytes(IntegerType(), 30)},
+    )
     assert len(index.for_data_file(1, file_no_overlap_2)) == 0