diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 3dffc2270c..7f9165fe4b 100644 --- a/pyiceberg/table/__init__.py +++ b/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.""" diff --git a/tests/table/test_init.py b/tests/table/test_init.py index 3f1e97768c..76544eb9b2 100644 --- a/tests/table/test_init.py +++ b/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}" diff --git a/tests/table/test_snapshots.py b/tests/table/test_snapshots.py index 5f1680ed59..4f4f8099ae 100644 --- a/tests/table/test_snapshots.py +++ b/tests/table/test_snapshots.py @@ -616,7 +616,8 @@ 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: @@ -624,28 +625,37 @@ def make_file() -> DataFile: 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)" )