perf: cache Transaction.table_metadata between reads

pyiceberg/table/__init__.py

@@ -241,6 +241,7 @@ class Transaction:
     _autocommit: bool
     _updates: tuple[TableUpdate, ...]
     _requirements: tuple[TableRequirement, ...]
+    _table_metadata_cache: tuple[TableMetadata, tuple[TableUpdate, ...], TableMetadata] | None
 
     def __init__(self, table: Table, autocommit: bool = False):
         """Open a transaction to stage and commit changes to a table.
@@ -255,10 +256,21 @@ def __init__(self, table: Table, autocommit: bool = False):
         self._requirements = ()
         self._snapshot_producers: list[_SnapshotProducer[Any]] = []
         self._failed = False
+        self._table_metadata_cache = None
 
     @property
     def table_metadata(self) -> TableMetadata:
-        return update_table_metadata(self._table.metadata, self._updates)
+        base, updates = self._table.metadata, self._updates
+        # update_table_metadata replays every staged update via model_copy(deep=True);
+        # the cache is keyed on the identity of its inputs so it self-invalidates
+        # whenever _updates is reassigned (tuple += creates a new object) or the
+        # underlying table metadata is refreshed.
+        cached = self._table_metadata_cache
+        if cached is not None and cached[0] is base and cached[1] is updates:
+            return cached[2]
+        result = update_table_metadata(base, updates)
+        self._table_metadata_cache = (base, updates, result)
+        return result
 
     def __enter__(self) -> Transaction:
         """Start a transaction to update the table."""

tests/table/test_init.py

@@ -2010,3 +2010,54 @@ def _spy(*args: Any, **kwargs: Any) -> FileIO:
 
     assert seen_locations, "expected at least one load_file_io call"
     assert all(loc is not None for loc in seen_locations), f"load_file_io called without a location: {seen_locations}"
+
+
+def test_transaction_table_metadata_cached(table_v2: Table) -> None:
+    """Repeated reads of an unchanged transaction state recompute at most once.
+
+    `Transaction.table_metadata` replays every staged update through
+    `model_copy(deep=True)`, so the cost of a read scales with the size of the
+    metadata (the snapshot list in particular), not with the work being done.
+    """
+    from unittest import mock
+
+    from pyiceberg.table.update import SetPropertiesUpdate, update_table_metadata
+
+    with mock.patch("pyiceberg.table.update_table_metadata", wraps=update_table_metadata) as spy:
+        txn = table_v2.transaction()
+
+        first = txn.table_metadata
+        for _ in range(10):
+            assert txn.table_metadata is first
+        assert spy.call_count == 1, f"expected 1 recompute for repeated reads, got {spy.call_count}"
+
+        txn._stage((SetPropertiesUpdate(updates={"k": "v"}),))
+        second = txn.table_metadata
+        assert second is not first
+        assert second.properties["k"] == "v"
+        for _ in range(10):
+            assert txn.table_metadata is second
+        assert spy.call_count == 2, f"expected 2 recomputes after one staged update, got {spy.call_count}"
+
+
+def test_transaction_table_metadata_cached_with_updates_already_staged(table_v2: Table) -> None:
+    """The cache must still hold once `_updates` is non-empty.
+
+    An empty-`_updates` short circuit (`return self._table.metadata` when nothing
+    is staged) would leave this case uncovered, and it is the expensive one:
+    `CreateTableTransaction` seeds `_updates` with ~10 entries before any write,
+    so a create-then-append transaction replays all of them on every read.
+    """
+    from unittest import mock
+
+    from pyiceberg.table.update import SetPropertiesUpdate, update_table_metadata
+
+    txn = table_v2.transaction()
+    txn._stage((SetPropertiesUpdate(updates={"staged": "before"}),))
+
+    with mock.patch("pyiceberg.table.update_table_metadata", wraps=update_table_metadata) as spy:
+        first = txn.table_metadata
+        for _ in range(10):
+            assert txn.table_metadata is first
+        assert first.properties["staged"] == "before"
+        assert spy.call_count == 1, f"expected 1 recompute with updates already staged, got {spy.call_count}"

tests/table/test_snapshots.py

@@ -616,36 +616,46 @@ def test_snapshot_producer_bounded_metadata_access(table_v2: Table) -> None:
     """
     from unittest import mock
 
-    from pyiceberg.table.update import update_table_metadata
+    from pyiceberg.table import Transaction
+    from pyiceberg.table.metadata import TableMetadata
     from pyiceberg.table.update.snapshot import _FastAppendFiles, _MergeAppendFiles
 
     def make_file() -> DataFile:
         return DataFile.from_args(content=DataFileContent.DATA, record_count=1, file_size_in_bytes=1, partition=Record())
 
     txn = table_v2.transaction()
 
-    with mock.patch("pyiceberg.table.update_table_metadata", wraps=update_table_metadata) as spy:
+    # Counts property reads rather than update_table_metadata calls: hoisting removes reads,
+    # and a read that the cache serves is still a read the hoisting was meant to remove.
+    reads = [0]
+    fget = Transaction.__dict__["table_metadata"].fget
+
+    def counting(self: Transaction) -> TableMetadata:
+        reads[0] += 1
+        return fget(self)
+
+    with mock.patch.object(Transaction, "table_metadata", property(counting)):
         # _summary() cost must not scale with the number of data files
-        def summary_calls(n_files: int) -> int:
+        def summary_reads(n_files: int) -> int:
             append = _FastAppendFiles(operation=Operation.APPEND, transaction=txn, io=table_v2.io)
             for _ in range(n_files):
                 append.append_data_file(make_file())
-            spy.reset_mock()
+            reads[0] = 0
             append._summary()
-            return spy.call_count
+            return reads[0]
 
-        few, many = summary_calls(10), summary_calls(100)
-        assert few == many, f"_summary() update_table_metadata calls scale with file count ({few} vs {many})"
-        assert many <= 2, f"_summary() triggered {many} update_table_metadata calls; expected O(1)"
+        few, many = summary_reads(10), summary_reads(100)
+        assert few == many, f"_summary() table_metadata reads scale with file count ({few} vs {many})"
+        assert many <= 2, f"_summary() made {many} table_metadata reads; expected O(1)"
 
-        # _MergeAppendFiles.__init__ should add exactly one call over _FastAppendFiles.__init__
-        spy.reset_mock()
+        # _MergeAppendFiles.__init__ should add exactly one read over _FastAppendFiles.__init__
+        reads[0] = 0
         _FastAppendFiles(operation=Operation.APPEND, transaction=txn, io=table_v2.io)
-        fast_init = spy.call_count
-        spy.reset_mock()
+        fast_init = reads[0]
+        reads[0] = 0
         _MergeAppendFiles(operation=Operation.APPEND, transaction=txn, io=table_v2.io)
-        merge_init = spy.call_count
+        merge_init = reads[0]
         assert merge_init - fast_init == 1, (
-            f"_MergeAppendFiles.__init__ made {merge_init - fast_init} extra update_table_metadata "
-            "calls over its superclass; expected 1 (hoisted)"
+            f"_MergeAppendFiles.__init__ made {merge_init - fast_init} extra table_metadata "
+            "reads over its superclass; expected 1 (hoisted)"
         )