add DeleteFileIndex, PartitionMap, _added_no_new_delete_files, validation, and tests

Author: gabeiglio

pyiceberg/partitioning.py

@@ -21,7 +21,7 @@
 from dataclasses import dataclass
 from datetime import date, datetime, time
 from functools import cached_property, singledispatch
-from typing import Annotated, Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union
+from typing import Annotated, Any, Callable, Dict, Generic, List, Optional, Tuple, TypeVar, Union
 from urllib.parse import quote_plus
 
 from pydantic import (
@@ -272,6 +272,60 @@ def assign_fresh_partition_spec_ids(spec: PartitionSpec, old_schema: Schema, fre
 T = TypeVar("T")
 
 
+class PartitionMap(Generic[T]):
+    _specs: dict[int, PartitionSpec]
+    _partition_maps: dict[int, dict[Record | None, T]]
+
+    def __init__(self, specs: dict[int, PartitionSpec]):
+        self._specs = specs
+        self._partition_maps = {}
+
+    def __len__(self) -> int:
+        """Return the length of the partition map.
+
+        Returns:
+            length of _partition_maps
+        """
+        return len(self.values())
+
+    def is_empty(self) -> bool:
+        return len(self.values()) == 0
+
+    def contains_key(self, spec_id: int, struct: Record) -> bool:
+        return struct in self._partition_maps.get(spec_id, {})
+
+    def contains_value(self, value: T) -> bool:
+        return value in self.values()
+
+    def get(self, spec_id: int, struct: Record | None) -> Optional[T]:
+        if partition_map := self._partition_maps.get(spec_id):
+            if result := partition_map.get(struct):
+                return result
+        return None
+
+    def put(self, spec_id: int, struct: Record | None, value: T) -> None:
+        if _ := self._specs.get(spec_id):
+            if spec_id not in self._partition_maps:
+                self._partition_maps[spec_id] = {struct: value}
+            else:
+                self._partition_maps[spec_id][struct] = value
+
+    def compute_if_absent(self, spec_id: int, struct: Record, value_factory: Callable[[], T]) -> T:
+        partition_map = self._partition_maps.setdefault(spec_id, {})
+        if struct in partition_map:
+            return partition_map[struct]
+
+        value = value_factory()
+        partition_map[struct] = value
+        return value
+
+    def values(self) -> list[T]:
+        result: list[T] = []
+        for partition_map in self._partition_maps.values():
+            result.extend(partition_map.values())
+        return result
+
+
 class PartitionSpecVisitor(Generic[T], ABC):
     @abstractmethod
     def identity(self, field_id: int, source_name: str, source_id: int) -> T: