small fixes and tests for PartitionMap

Author: gabeiglio

pyiceberg/partitioning.py

@@ -274,7 +274,7 @@ def assign_fresh_partition_spec_ids(spec: PartitionSpec, old_schema: Schema, fre
 
 class PartitionMap(Generic[T]):
     _specs: dict[int, PartitionSpec]
-    _partition_maps: dict[int, dict[Record, T]]
+    _partition_maps: dict[int, dict[Record | None, T]]
 
     def __init__(self, specs: dict[int, PartitionSpec]):
         self._specs = specs
@@ -286,34 +286,38 @@ def __len__(self) -> int:
         Returns:
             length of _partition_maps
         """
-        return len(self._partition_maps.values())
+        return len(self.values())
 
     def is_empty(self) -> bool:
         return len(self._partition_maps.values()) == 0
 
     def contains_key(self, spec_id: int, struct: Record) -> bool:
-        try:
-            return struct in self._partition_maps[spec_id]
-        except KeyError as _:
-            return False
+        return self._partition_maps.get(spec_id) is not None
 
     def contains_value(self, value: T) -> bool:
         return value in self._partition_maps.values()
 
-    def get(self, spec_id: int, struct: Record) -> Optional[T]:
+    def get(self, spec_id: int, struct: Record | None) -> Optional[T]:
         if partition_map := self._partition_maps.get(spec_id):
-            return partition_map.get(struct)
+            if result := partition_map.get(struct):
+                return result
         return None
 
-    def put(self, spec_id: int, struct: Record, value: T) -> None:
+    def put(self, spec_id: int, struct: Record | None, value: T) -> None:
         if _ := self._specs.get(spec_id):
-            self._partition_maps[spec_id] = {struct: value}
+            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: T, value_factory: Callable[[], T]) -> T:
-        if partition_map := self._partition_maps.get(spec_id):
-            if val := partition_map.get(struct):
-                return val
-        return value_factory()
+        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] = []

pyiceberg/table/update/validate.py

@@ -56,7 +56,7 @@ def _index_if_needed(self) -> None:
             self._indexed = True
 
     def add_entry(self, entry: ManifestEntry) -> None:
-        if self._buffer is None:
+        if self._indexed:
             raise Exception("Can't add files upon indexing.")
         self._buffer.append(entry)
 
@@ -99,7 +99,7 @@ def _index_if_needed(self) -> None:
 
     def add_entry(self, spec: PartitionSpec, entry: ManifestEntry) -> None:
         # TODO: Equality deletes should consider the spec to get the equality fields
-        if self._buffer is None:
+        if self._indexed:
             raise Exception("Can't add files upon indexing.")
         self._buffer.append(entry)
 

tests/table/test_partitioning.py

@@ -21,7 +21,7 @@
 
 import pytest
 
-from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
+from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionMap, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.transforms import (
     BucketTransform,
@@ -225,5 +225,63 @@ def test_deserialize_partition_field_v3() -> None:
     assert field == PartitionField(source_id=1, field_id=1000, transform=TruncateTransform(width=19), name="str_truncate")
 
 
-def test_partition_map() -> None:
-    pass
+@pytest.fixture
+def specs_set() -> dict[int, PartitionSpec]:
+    return {
+        0: UNPARTITIONED_PARTITION_SPEC,
+        1: PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=DayTransform(), name="dayPartition"), spec_id=1),
+        2: PartitionSpec(
+            PartitionField(source_id=1, field_id=1000, transform=DayTransform(), name="dayPartition"),
+            PartitionField(source_id=2, field_id=1001, transform=IdentityTransform(), name="identityPartition"),
+            spec_id=2,
+        ),
+    }
+
+
+def test_empty_partition_map() -> None:
+    specs: dict[int, PartitionSpec] = {UNPARTITIONED_PARTITION_SPEC.spec_id: UNPARTITIONED_PARTITION_SPEC}
+    partition_map: PartitionMap[str] = PartitionMap(specs)
+    assert partition_map.is_empty()
+    assert len(partition_map) == 0
+    assert not partition_map.contains_key(1, Record(1))
+    assert len(partition_map.values()) == 0
+
+
+def test_size_partition_map(specs_set: dict[int, PartitionSpec]) -> None:
+    partition_map: PartitionMap[str] = PartitionMap(specs_set)
+    partition_map.put(UNPARTITIONED_PARTITION_SPEC.spec_id, None, "v1")
+    partition_map.put(specs_set[1].spec_id, Record("aaa"), "v2")
+    partition_map.put(specs_set[1].spec_id, Record("bbb"), "v3")
+    partition_map.put(specs_set[2].spec_id, Record("ccc", 2), "v4")
+    assert not partition_map.is_empty()
+    assert len(partition_map) == 4
+    # assert partition_map.get(UNPARTITIONED_PARTITION_SPEC.spec_id, None) == "v1"
+
+
+def test_put_and_get_partition_map(specs_set: dict[int, PartitionSpec]) -> None:
+    partition_map: PartitionMap[str] = PartitionMap(specs_set)
+    partition_map.put(UNPARTITIONED_PARTITION_SPEC.spec_id, None, "v1")
+    partition_map.put(specs_set[1].spec_id, Record("aaa", 1), "v2")
+    assert partition_map.get(UNPARTITIONED_PARTITION_SPEC.spec_id, None) == "v1"
+    assert partition_map.get(specs_set[1].spec_id, Record("aaa", 1)) == "v2"
+
+
+def test_values_partition_map(specs_set: dict[int, PartitionSpec]) -> None:
+    partition_map: PartitionMap[str] = PartitionMap(specs_set)
+    partition_map.put(UNPARTITIONED_PARTITION_SPEC.spec_id, None, "v1")
+    partition_map.put(specs_set[1].spec_id, Record("aaa"), "v2")
+    partition_map.put(specs_set[1].spec_id, Record("bbb"), "v3")
+    partition_map.put(specs_set[2].spec_id, Record("ccc", 2), "v4")
+    assert partition_map.values() == ["v1", "v2", "v3", "v4"]
+
+
+def test_compute_if_absent_partition_map(specs_set: dict[int, PartitionSpec]) -> None:
+    partition_map: PartitionMap[str] = PartitionMap(specs_set)
+
+    result1 = partition_map.compute_if_absent(specs_set[1].spec_id, Record("a"), "v1", lambda: "v1")
+    assert result1 == "v1"
+    assert partition_map.get(specs_set[1].spec_id, Record("a")) == "v1"
+
+    result2 = partition_map.compute_if_absent(specs_set[1].spec_id, Record("a"), "v2", lambda: "v2")
+    assert result2 == "v1"
+    assert partition_map.get(specs_set[1].spec_id, Record("a")) == "v1"