feat: add mor delete file index support

Author: Babar

pyiceberg/io/pyarrow.py

@@ -978,18 +978,30 @@ def _get_file_format(file_format: FileFormat, **kwargs: Dict[str, Any]) -> ds.Fi
         raise ValueError(f"Unsupported file format: {file_format}")
 
 
-def _read_deletes(io: FileIO, data_file: DataFile) -> Dict[str, pa.ChunkedArray]:
+def _construct_fragment(io: FileIO, data_file: DataFile, file_format_kwargs: Dict[str, Any] = EMPTY_DICT) -> ds.Fragment:
+    with io.new_input(data_file.file_path).open() as fi:
+        return _get_file_format(
+            data_file.file_format, **file_format_kwargs
+        ).make_fragment(fi)
+
+def _read_deletes(io: FileIO, data_file: DataFile) -> Union[Dict[str, pa.ChunkedArray], pa.Table]:
     if data_file.file_format == FileFormat.PARQUET:
-        with io.new_input(data_file.file_path).open() as fi:
-            delete_fragment = _get_file_format(
-                data_file.file_format, dictionary_columns=("file_path",), pre_buffer=True, buffer_size=ONE_MEGABYTE
-            ).make_fragment(fi)
-            table = ds.Scanner.from_fragment(fragment=delete_fragment).to_table()
-        table = table.unify_dictionaries()
-        return {
-            file.as_py(): table.filter(pc.field("file_path") == file).column("pos")
-            for file in table.column("file_path").chunks[0].dictionary
-        }
+        delete_fragment = _construct_fragment(
+            io,
+            data_file,
+            file_format_kwargs={"dictionary_columns": ("file_path",), "pre_buffer": True, "buffer_size": ONE_MEGABYTE},
+        )
+        table = ds.Scanner.from_fragment(fragment=delete_fragment).to_table()
+        if data_file.content == DataFileContent.POSITION_DELETES:
+            table = table.unify_dictionaries()
+            return {
+                file.as_py(): table.filter(pc.field("file_path") == file).column("pos")
+                for file in table.column("file_path").chunks[0].dictionary
+            }
+        elif data_file.content == DataFileContent.EQUALITY_DELETES:
+            return table
+        else:
+            raise ValueError(f"Unsupported delete file content: {data_file.content}")
     elif data_file.file_format == FileFormat.PUFFIN:
         with io.new_input(data_file.file_path).open() as fi:
             payload = fi.read()
@@ -1445,7 +1457,7 @@ def _task_to_record_batches(
     bound_row_filter: BooleanExpression,
     projected_schema: Schema,
     projected_field_ids: Set[int],
-    positional_deletes: Optional[List[ChunkedArray]],
+    deletes: Optional[List[Union[pa.ChunkedArray, pa.Table]]],
     case_sensitive: bool,
     name_mapping: Optional[NameMapping] = None,
     partition_spec: Optional[PartitionSpec] = None,
@@ -1479,9 +1491,18 @@ def _task_to_record_batches(
             schema=physical_schema,
             # This will push down the query to Arrow.
             # But in case there are positional deletes, we have to apply them first
-            filter=pyarrow_filter if not positional_deletes else None,
+            filter=pyarrow_filter if not deletes else None,
             columns=[col.name for col in file_project_schema.columns],
         )
+        positional_deletes = []
+        combined_eq_deletes = []
+        if deletes:
+            positional_deletes = [d for d in deletes if isinstance(d, pa.ChunkedArray)]
+            equality_deletes = [d for d in deletes if isinstance(d, pa.Table)]
+            if equality_deletes:
+                task_eq_files = [df for df in task.delete_files if df.content == DataFileContent.EQUALITY_DELETES]
+                # Group and combine equality deletes
+                combined_eq_deletes = group_equality_deletes(task_eq_files, equality_deletes)
 
         next_index = 0
         batches = fragment_scanner.to_batches()
@@ -1499,6 +1520,17 @@ def _task_to_record_batches(
             if current_batch.num_rows == 0:
                 continue
 
+            if combined_eq_deletes:
+                table = pa.Table.from_batches([current_batch])
+                for equality_ids, combined_table in combined_eq_deletes:
+                    table = _apply_equality_deletes(table, combined_table, equality_ids, file_schema)
+                    if table.num_rows == 0:
+                        break
+                if table.num_rows > 0:
+                    current_batch = table.combine_chunks().to_batches()[0]
+                else:
+                    continue
+
             # Apply the user filter
             if pyarrow_filter is not None:
                 # Temporary fix until PyArrow 21 is released ( https://github.com/apache/arrow/pull/46057 )
@@ -1529,9 +1561,16 @@ def _task_to_record_batches(
 
 
 def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> Dict[str, List[ChunkedArray]]:
-    deletes_per_file: Dict[str, List[ChunkedArray]] = {}
-    unique_deletes = set(itertools.chain.from_iterable([task.delete_files for task in tasks]))
-    if len(unique_deletes) > 0:
+    deletes_per_file: Dict[str, List[Union[pa.ChunkedArray, pa.Table]]] = {}
+
+    # Position Deletes
+    unique_deletes = {
+        df
+        for task in tasks
+        for df in task.delete_files
+        if df.content == DataFileContent.POSITION_DELETES and df.file_format != FileFormat.PUFFIN
+    }
+    if unique_deletes:
         executor = ExecutorFactory.get_or_create()
         deletes_per_files: Iterator[Dict[str, ChunkedArray]] = executor.map(
             lambda args: _read_deletes(*args),
@@ -1543,7 +1582,44 @@ def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> Dict[st
                     deletes_per_file[file].append(arr)
                 else:
                     deletes_per_file[file] = [arr]
+    # Deletion Vectors
+    deletion_vectors = {
+        df
+        for task in tasks
+        for df in task.delete_files
+        if df.content == DataFileContent.POSITION_DELETES and df.file_format == FileFormat.PUFFIN
+    }
+    if deletion_vectors:
+        executor = ExecutorFactory.get_or_create()
+        dv_results = executor.map(
+            lambda args: _read_deletes(*args),
+            [(_fs_from_file_path(io, delete_file.file_path), delete_file) for delete_file in deletion_vectors],
+        )
+        for delete in dv_results:
+            for file, arr in delete.items():
+                # Deletion vectors replace all position deletes for a file
+                deletes_per_file[file] = [arr]
+
+    # Equality Deletes
+    equality_delete_tasks = []
+    for task in tasks:
+        equality_deletes = [df for df in task.delete_files if df.content == DataFileContent.EQUALITY_DELETES]
+        if equality_deletes:
+            for delete_file in equality_deletes:
+                equality_delete_tasks.append((task.file.file_path, delete_file))
+
+    if equality_delete_tasks:
+        executor = ExecutorFactory.get_or_create()
 
+        # Processing equality delete tasks in parallel like position deletes
+        equality_delete_results = executor.map(
+            lambda args: (args[0], _read_deletes(_fs_from_file_path(io, args[1].file_path), args[1])),
+            equality_delete_tasks,
+        )
+        for file_path, equality_delete_table in equality_delete_results:
+            if file_path not in deletes_per_file:
+                deletes_per_file[file_path] = []
+            deletes_per_file[file_path].append(equality_delete_table)
     return deletes_per_file
 
 
@@ -2799,3 +2875,56 @@ def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Ar
     field_array = arrow_table[path_parts[0]]
     # Navigate into the struct using the remaining path parts
     return pc.struct_field(field_array, path_parts[1:])
+
+
+def group_equality_deletes(
+    task_eq_files: List[DataFile], equality_delete_tables: List[pa.Table]
+) -> List[Tuple[List[int], pa.Table]]:
+    """Group equality delete tables by their equality IDs."""
+    equality_delete_groups: Dict[frozenset[int], List[Tuple[List[int], pa.Table]]] = {}
+
+    for delete_file, delete_table in zip(task_eq_files, equality_delete_tables):
+        if delete_file.equality_ids:
+            key = frozenset(delete_file.equality_ids)
+
+            # Add to the appropriate group
+            if key not in equality_delete_groups:
+                equality_delete_groups[key] = []
+            equality_delete_groups[key].append((delete_file.equality_ids, delete_table))
+
+    # Combine tables with the same equality IDs
+    combined_deletes = []
+    for items in equality_delete_groups.values():
+        # Use the original equality IDs from the first item
+        original_ids = items[0][0]
+        tables = [item[1] for item in items]
+
+        if tables:
+            combined_table = pa.concat_tables(tables)
+            combined_deletes.append((original_ids, combined_table))
+
+    return combined_deletes
+
+
+def _apply_equality_deletes(
+    data_table: pa.Table, delete_table: pa.Table, equality_ids: List[int], data_schema: Optional[Schema]
+) -> pa.Table:
+    """Apply equality deletes to a data table.
+
+    Filter out rows from the table that match the equality delete table the conditions in it.
+    Args:
+        data_table: A PyArrow table which has data to filter
+        delete_table: A PyArrow table containing the equality deletes
+        equality_ids: A List of field IDs to use for equality comparison
+        data_schema: The schema of the PyArrow table
+    Returns:
+        A filtered PyArrow table with matching rows removed
+    """
+    if len(delete_table) == 0:
+        return data_table
+    if data_schema is None:
+        raise ValueError("Schema is required for applying equality deletes")
+    equality_columns = [data_schema.find_field(fid).name for fid in equality_ids]
+    # Use PyArrow's join function with left anti join type
+    result = data_table.join(delete_table.select(equality_columns), keys=equality_columns, join_type="left anti")
+    return result

pyiceberg/table/__init__.py

@@ -41,7 +41,6 @@
 )
 
 from pydantic import Field
-from sortedcontainers import SortedList
 
 import pyiceberg.expressions.parser as parser
 from pyiceberg.expressions import (
@@ -64,7 +63,6 @@
 )
 from pyiceberg.io import FileIO, load_file_io
 from pyiceberg.manifest import (
-    POSITIONAL_DELETE_SCHEMA,
     DataFile,
     DataFileContent,
     ManifestContent,
@@ -78,6 +76,7 @@
     PartitionSpec,
 )
 from pyiceberg.schema import Schema
+from pyiceberg.table.delete_file_index import DeleteFileIndex
 from pyiceberg.table.inspect import InspectTable
 from pyiceberg.table.locations import LocationProvider, load_location_provider
 from pyiceberg.table.metadata import (
@@ -1793,29 +1792,20 @@ def _min_sequence_number(manifests: List[ManifestFile]) -> int:
         return INITIAL_SEQUENCE_NUMBER
 
 
-def _match_deletes_to_data_file(data_entry: ManifestEntry, positional_delete_entries: SortedList[ManifestEntry]) -> Set[DataFile]:
-    """Check if the delete file is relevant for the data file.
-
-    Using the column metrics to see if the filename is in the lower and upper bound.
+def _match_deletes_to_data_file(data_entry: ManifestEntry, delete_file_index: DeleteFileIndex) -> Set[DataFile]:
+    """Check if delete files are relevant for the data file.
 
     Args:
-        data_entry (ManifestEntry): The manifest entry path of the datafile.
-        positional_delete_entries (List[ManifestEntry]): All the candidate positional deletes manifest entries.
+        data_entry (ManifestEntry): The manifest entry of the data file.
+        delete_file_index (DeleteFileIndex): Index containing all delete files.
 
     Returns:
-        A set of files that are relevant for the data file.
+        A set of delete files that are relevant for the data file.
     """
-    relevant_entries = positional_delete_entries[positional_delete_entries.bisect_right(data_entry) :]
-
-    if len(relevant_entries) > 0:
-        evaluator = _InclusiveMetricsEvaluator(POSITIONAL_DELETE_SCHEMA, EqualTo("file_path", data_entry.data_file.file_path))
-        return {
-            positional_delete_entry.data_file
-            for positional_delete_entry in relevant_entries
-            if evaluator.eval(positional_delete_entry.data_file)
-        }
-    else:
-        return set()
+    candidate_deletes = delete_file_index.for_data_file(
+        data_entry.sequence_number or 0, data_entry.data_file, partition_key=data_entry.data_file.partition
+    )
+    return set(candidate_deletes)
 
 
 class DataScan(TableScan):
@@ -1921,7 +1911,7 @@ def plan_files(self) -> Iterable[FileScanTask]:
         min_sequence_number = _min_sequence_number(manifests)
 
         data_entries: List[ManifestEntry] = []
-        positional_delete_entries = SortedList(key=lambda entry: entry.sequence_number or INITIAL_SEQUENCE_NUMBER)
+        delete_file_index = DeleteFileIndex(self.table_metadata.schema(), self.table_metadata.specs())
 
         executor = ExecutorFactory.get_or_create()
         for manifest_entry in chain(
@@ -1942,19 +1932,16 @@ def plan_files(self) -> Iterable[FileScanTask]:
             data_file = manifest_entry.data_file
             if data_file.content == DataFileContent.DATA:
                 data_entries.append(manifest_entry)
-            elif data_file.content == DataFileContent.POSITION_DELETES:
-                positional_delete_entries.add(manifest_entry)
-            elif data_file.content == DataFileContent.EQUALITY_DELETES:
-                raise ValueError("PyIceberg does not yet support equality deletes: https://github.com/apache/iceberg/issues/6568")
+            elif data_file.content in (DataFileContent.POSITION_DELETES, DataFileContent.EQUALITY_DELETES):
+                delete_file_index.add_delete_file(manifest_entry, partition_key=data_file.partition)
             else:
                 raise ValueError(f"Unknown DataFileContent ({data_file.content}): {manifest_entry}")
-
         return [
             FileScanTask(
                 data_entry.data_file,
                 delete_files=_match_deletes_to_data_file(
                     data_entry,
-                    positional_delete_entries,
+                    delete_file_index,
                 ),
                 residual=residual_evaluators[data_entry.data_file.spec_id](data_entry.data_file).residual_for(
                     data_entry.data_file.partition