delete_file_index.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

from bisect import bisect_left
from typing import TYPE_CHECKING

from pyiceberg.conversions import from_bytes
from pyiceberg.expressions import EqualTo
from pyiceberg.expressions.visitors import _InclusiveMetricsEvaluator
from pyiceberg.manifest import INITIAL_SEQUENCE_NUMBER, POSITIONAL_DELETE_SCHEMA, DataFile, DataFileContent, ManifestEntry
from pyiceberg.typedef import Record

if TYPE_CHECKING:
    from pyiceberg.schema import Schema

PATH_FIELD_ID = 2147483546


class PositionDeletes:
    """Collects position delete files and indexes them by sequence number."""

    __slots__ = ("_buffer", "_seqs", "_files")

    def __init__(self) -> None:
        self._buffer: list[tuple[DataFile, int]] | None = []
        self._seqs: list[int] = []
        self._files: list[tuple[DataFile, int]] = []

    def add(self, delete_file: DataFile, seq_num: int) -> None:
        if self._buffer is None:
            raise ValueError("Cannot add files after indexing")
        self._buffer.append((delete_file, seq_num))

    def _ensure_indexed(self) -> None:
        if self._buffer is not None:
            self._files = sorted(self._buffer, key=lambda file: file[1])
            self._seqs = [seq for _, seq in self._files]
            self._buffer = None

    def filter_by_seq(self, seq: int) -> list[DataFile]:
        self._ensure_indexed()
        if not self._files:
            return []
        start_idx = bisect_left(self._seqs, seq)
        return [delete_file for delete_file, _ in self._files[start_idx:]]

    def referenced_delete_files(self) -> list[DataFile]:
        self._ensure_indexed()
        return [data_file for data_file, _ in self._files]


class EqualityDeletes(PositionDeletes):
    """Collects equality delete files and indexes them by sequence number."""

    def add(self, delete_file: DataFile, seq_num: int) -> None:
        # Equality deletes are indexed by sequence number - 1 to ensure they only
        # apply to data files added in strictly earlier snapshots.
        super().add(delete_file, seq_num - 1)


def _has_path_bounds(delete_file: DataFile) -> bool:
    lower = delete_file.lower_bounds
    upper = delete_file.upper_bounds
    if not lower or not upper:
        return False

    return PATH_FIELD_ID in lower and PATH_FIELD_ID in upper


def _applies_to_data_file(delete_file: DataFile, data_file: DataFile) -> bool:
    if not _has_path_bounds(delete_file):
        return True

    evaluator = _InclusiveMetricsEvaluator(POSITIONAL_DELETE_SCHEMA, EqualTo("file_path", data_file.file_path))
    return evaluator.eval(delete_file)


def _eq_applies_to_data_file(eq_delete_file: DataFile, data_file: DataFile, schema: Schema) -> bool:
    if not eq_delete_file.equality_ids:
        return True

    for field_id in eq_delete_file.equality_ids:
        if (
            eq_delete_file.lower_bounds
            and eq_delete_file.upper_bounds
            and data_file.lower_bounds
            and data_file.upper_bounds
            and field_id in eq_delete_file.lower_bounds
            and field_id in eq_delete_file.upper_bounds
            and field_id in data_file.lower_bounds
            and field_id in data_file.upper_bounds
        ):
            field_type = schema.find_type(field_id)
            if not field_type.is_primitive:
                continue

            eq_lower = from_bytes(field_type, eq_delete_file.lower_bounds[field_id])
            eq_upper = from_bytes(field_type, eq_delete_file.upper_bounds[field_id])
            data_lower = from_bytes(field_type, data_file.lower_bounds[field_id])
            data_upper = from_bytes(field_type, data_file.upper_bounds[field_id])

            if eq_upper < data_lower or eq_lower > data_upper:
                return False

    return True


def _referenced_data_file_path(delete_file: DataFile) -> str | None:
    """Return the path, if the path bounds evaluate to the same location."""
    lower_bounds = delete_file.lower_bounds
    upper_bounds = delete_file.upper_bounds

    if not lower_bounds or not upper_bounds:
        return None

    lower = lower_bounds.get(PATH_FIELD_ID)
    upper = upper_bounds.get(PATH_FIELD_ID)

    if lower and upper and lower == upper:
        try:
            return lower.decode("utf-8")
        except (UnicodeDecodeError, AttributeError):
            pass

    return None


def _partition_key(spec_id: int, partition: Record | None) -> tuple[int, Record]:
    if partition:
        return spec_id, partition
    return spec_id, Record()  # unpartitioned handling


class DeleteFileIndex:
    """Indexes position and equality delete files by partition and by exact data file path."""

    def __init__(self, schema: Schema | None = None) -> None:
        self._schema = schema
        self._by_partition: dict[tuple[int, Record], PositionDeletes] = {}
        self._by_path: dict[str, PositionDeletes] = {}
        self._eq_by_partition: dict[tuple[int, Record], EqualityDeletes] = {}
        self._global_eq_deletes: EqualityDeletes = EqualityDeletes()

    def is_empty(self) -> bool:
        return (
            not self._by_partition
            and not self._by_path
            and not self._eq_by_partition
            and not self._global_eq_deletes.referenced_delete_files()
        )

    def add_delete_file(self, manifest_entry: ManifestEntry, partition_key: Record | None = None) -> None:
        delete_file = manifest_entry.data_file
        seq = manifest_entry.sequence_number or INITIAL_SEQUENCE_NUMBER

        if delete_file.content == DataFileContent.POSITION_DELETES:
            target_path = _referenced_data_file_path(delete_file)
            if target_path:
                deletes = self._by_path.setdefault(target_path, PositionDeletes())
                deletes.add(delete_file, seq)
            else:
                key = _partition_key(delete_file.spec_id or 0, partition_key)
                deletes = self._by_partition.setdefault(key, PositionDeletes())
                deletes.add(delete_file, seq)
        elif delete_file.content == DataFileContent.EQUALITY_DELETES:
            if partition_key is None or len(partition_key) == 0:
                self._global_eq_deletes.add(delete_file, seq)
            else:
                key = _partition_key(delete_file.spec_id or 0, partition_key)
                deletes = self._eq_by_partition.setdefault(key, EqualityDeletes())
                deletes.add(delete_file, seq)

    def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> set[DataFile]:
        if self.is_empty():
            return set()

        deletes: set[DataFile] = set()
        spec_id = data_file.spec_id or 0
        key = _partition_key(spec_id, partition_key)

        # Add position deletes
        partition_pos_deletes = self._by_partition.get(key)
        if partition_pos_deletes:
            for delete_file in partition_pos_deletes.filter_by_seq(seq_num):
                if _applies_to_data_file(delete_file, data_file):
                    deletes.add(delete_file)

        path_pos_deletes = self._by_path.get(data_file.file_path)
        if path_pos_deletes:
            deletes.update(path_pos_deletes.filter_by_seq(seq_num))

        # Add equality deletes
        candidate_eq_deletes: list[DataFile] = []
        partition_eq_deletes = self._eq_by_partition.get(key)
        if partition_eq_deletes:
            candidate_eq_deletes.extend(partition_eq_deletes.filter_by_seq(seq_num))

        candidate_eq_deletes.extend(self._global_eq_deletes.filter_by_seq(seq_num))

        for eq_delete_file in candidate_eq_deletes:
            if self._schema and not _eq_applies_to_data_file(eq_delete_file, data_file, self._schema):
                continue
            deletes.add(eq_delete_file)

        return deletes

    def referenced_delete_files(self) -> list[DataFile]:
        data_files: list[DataFile] = []

        for deletes in self._by_partition.values():
            data_files.extend(deletes.referenced_delete_files())

        for deletes in self._by_path.values():
            data_files.extend(deletes.referenced_delete_files())

        for deletes in self._eq_by_partition.values():
            data_files.extend(deletes.referenced_delete_files())

        data_files.extend(self._global_eq_deletes.referenced_delete_files())

        return data_files