| # 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. |
| |
| """Regression test for caching of OVERWRITE changes across commit retries. |
| |
| On every OVERWRITE retry, pypaimon re-scanned the full target partitions to |
| recompute the files to delete. Under concurrent writers this made each retry as |
| expensive as the first attempt. OverwriteChangesProvider caches the existing |
| files of the target partitions and, on retry, reuses them when the snapshots in |
| between can be applied from target-partition DELTA manifests, instead of a full |
| re-scan. |
| |
| The test deterministically forces ``K`` conflicts, each advancing the latest |
| snapshot with an append to an unrelated partition, and asserts the full scan |
| runs once (not ``K + 1``) and the cache is advanced by a delta probe per retry. |
| """ |
| |
| import os |
| import shutil |
| import tempfile |
| import unittest |
| |
| import pandas as pd |
| import pyarrow as pa |
| |
| from pypaimon import CatalogFactory, Schema |
| from pypaimon.write.commit.overwrite_changes_provider import OverwriteChangesProvider |
| |
| |
| class OverwriteChangesCacheTest(unittest.TestCase): |
| |
| def setUp(self): |
| self.temp_dir = tempfile.mkdtemp(prefix="ow_cache_") |
| self.warehouse = os.path.join(self.temp_dir, 'wh') |
| self.catalog = CatalogFactory.create({"warehouse": self.warehouse}) |
| self.catalog.create_database("test_db", True) |
| |
| pa_schema = pa.schema([('f0', pa.int32()), ('f1', pa.string())]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| partition_keys=['f0'], |
| options={'dynamic-partition-overwrite': 'false'}, |
| ) |
| self.catalog.create_table('test_db.t', schema, False) |
| self.table = self.catalog.get_table('test_db.t') |
| |
| # Seed: f0=1 is the overwrite target, f0=2 untouched. |
| self._append(pd.DataFrame({'f0': [1, 1, 2], 'f1': ['a', 'b', 'c']})) |
| |
| def tearDown(self): |
| shutil.rmtree(self.temp_dir, ignore_errors=True) |
| |
| def _append(self, df): |
| wb = self.table.new_batch_write_builder() |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(df) |
| c.commit(w.prepare_commit()) |
| w.close() |
| c.close() |
| |
| def test_overwrite_scan_runs_once_not_once_per_retry(self): |
| K = 3 # force 3 conflicts; the 4th attempt wins |
| |
| wb = self.table.new_batch_write_builder().overwrite({'f0': 1}) |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']})) |
| messages = w.prepare_commit() |
| |
| fsc = c.file_store_commit |
| |
| # --- count provider full scans and delta probes (class-level spies) --- |
| counts = {'full_scan': 0, 'probe': 0} |
| orig_full = OverwriteChangesProvider._full_scan_manifest_entries |
| orig_probe = OverwriteChangesProvider._read_delta_manifest_entries |
| |
| def spy_full(self, *a, **k): |
| counts['full_scan'] += 1 |
| return orig_full(self, *a, **k) |
| |
| def spy_probe(self, *a, **k): |
| counts['probe'] += 1 |
| return orig_probe(self, *a, **k) |
| |
| # --- inject K CAS conflicts; each advances latest via an APPEND to an |
| # unrelated partition (f0=99) so the delta probe finds the target |
| # (f0=1) untouched and the cache is reused. |
| orig_cas = fsc.snapshot_commit.commit |
| cas = {'fails': 0} |
| |
| def patched_cas(base_snapshot_uuid, snapshot, statistics): |
| if snapshot.commit_kind == "OVERWRITE" and cas['fails'] < K: |
| cas['fails'] += 1 |
| self._append(pd.DataFrame({'f0': [99], 'f1': [f'x{cas["fails"]}']})) |
| return False |
| return orig_cas(base_snapshot_uuid, snapshot, statistics) |
| |
| fsc.snapshot_commit.commit = patched_cas |
| OverwriteChangesProvider._full_scan_manifest_entries = spy_full |
| OverwriteChangesProvider._read_delta_manifest_entries = spy_probe |
| try: |
| c.commit(messages) |
| c.close() |
| finally: |
| OverwriteChangesProvider._full_scan_manifest_entries = orig_full |
| OverwriteChangesProvider._read_delta_manifest_entries = orig_probe |
| |
| # Harness sanity: we really did force K conflicts and then converged. |
| self.assertEqual(cas['fails'], K, "expected exactly K forced conflicts") |
| |
| print(f"\n[overwrite-cache] K={K} conflicts -> " |
| f"full_scan={counts['full_scan']}, probe={counts['probe']} " |
| f"(target: full_scan=1, probe=K)") |
| |
| # #7894: the full overwrite scan runs once; each retry reuses the cache |
| # after a cheap delta probe of the (unrelated) intervening append. |
| self.assertEqual( |
| counts['full_scan'], 1, |
| f"full overwrite scan ran {counts['full_scan']}x; should run once and " |
| f"reuse the cache on retries") |
| self.assertEqual( |
| counts['probe'], K, |
| f"delta probe ran {counts['probe']}x; should probe once per retry " |
| f"(= K = {K})") |
| |
| # Sanity: f0=1 overwritten, f0=2 preserved, f0=99 appended K times. |
| read_builder = self.table.new_read_builder() |
| actual = read_builder.new_read().to_pandas( |
| read_builder.new_scan().plan().splits()) |
| self.assertEqual(sorted(actual[actual['f0'] == 1]['f1'].tolist()), ['new']) |
| self.assertEqual(sorted(actual[actual['f0'] == 2]['f1'].tolist()), ['c']) |
| self.assertEqual(len(actual[actual['f0'] == 99]), K) |
| |
| def test_cache_applies_delta_when_concurrent_append_hits_target_partition(self): |
| # Concurrent appends hit the overwrite target; retry advances cached |
| # target-partition state from APPEND deltas instead of rebuilding. |
| K = 3 |
| |
| wb = self.table.new_batch_write_builder().overwrite({'f0': 1}) |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']})) |
| messages = w.prepare_commit() |
| |
| fsc = c.file_store_commit |
| |
| counts = {'full_scan': 0, 'probe': 0} |
| orig_full = OverwriteChangesProvider._full_scan_manifest_entries |
| orig_probe = OverwriteChangesProvider._read_delta_manifest_entries |
| |
| def spy_full(self, *a, **k): |
| counts['full_scan'] += 1 |
| return orig_full(self, *a, **k) |
| |
| def spy_probe(self, *a, **k): |
| counts['probe'] += 1 |
| return orig_probe(self, *a, **k) |
| |
| orig_cas = fsc.snapshot_commit.commit |
| cas = {'fails': 0} |
| |
| def patched_cas(base_snapshot_uuid, snapshot, statistics): |
| if snapshot.commit_kind == "OVERWRITE" and cas['fails'] < K: |
| cas['fails'] += 1 |
| self._append(pd.DataFrame({'f0': [1], 'f1': [f'y{cas["fails"]}']})) |
| return False |
| return orig_cas(base_snapshot_uuid, snapshot, statistics) |
| |
| fsc.snapshot_commit.commit = patched_cas |
| OverwriteChangesProvider._full_scan_manifest_entries = spy_full |
| OverwriteChangesProvider._read_delta_manifest_entries = spy_probe |
| try: |
| c.commit(messages) |
| c.close() |
| finally: |
| OverwriteChangesProvider._full_scan_manifest_entries = orig_full |
| OverwriteChangesProvider._read_delta_manifest_entries = orig_probe |
| |
| self.assertEqual(cas['fails'], K, "expected exactly K forced conflicts") |
| |
| # Target APPEND deltas can be applied to cached state; no full rebuild. |
| self.assertEqual(counts['full_scan'], 1, |
| f"full scan ran {counts['full_scan']}x; APPEND deltas " |
| f"should advance cached target state") |
| self.assertEqual(counts['probe'], K, |
| f"delta probe ran {counts['probe']}x; once per retry (= K)") |
| |
| # Overwrite wins: f0=1 is just 'new', f0=2 untouched. |
| read_builder = self.table.new_read_builder() |
| actual = read_builder.new_read().to_pandas( |
| read_builder.new_scan().plan().splits()) |
| self.assertEqual(sorted(actual[actual['f0'] == 1]['f1'].tolist()), ['new']) |
| self.assertEqual(sorted(actual[actual['f0'] == 2]['f1'].tolist()), ['c']) |
| |
| def _overwrite_partition(self, part_val, f1_val): |
| wb = self.table.new_batch_write_builder().overwrite({'f0': part_val}) |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(pd.DataFrame({'f0': [part_val], 'f1': [f1_val]})) |
| c.commit(w.prepare_commit()) |
| w.close() |
| c.close() |
| |
| def _run_with_conflicts(self, c, messages, K, concurrent_fn): |
| # Run an overwrite commit, forcing K CAS conflicts (each calls |
| # concurrent_fn(i) to advance the latest snapshot). Returns this commit's |
| # own OverwriteChangesProvider so the caller can read its counters |
| # (only this provider, not any concurrent writer's). |
| fsc = c.file_store_commit |
| captured = {} |
| orig_factory = fsc._overwrite_changes_provider |
| |
| def capturing_factory(*a, **k): |
| captured['provider'] = orig_factory(*a, **k) |
| return captured['provider'] |
| |
| fsc._overwrite_changes_provider = capturing_factory |
| |
| orig_cas = fsc.snapshot_commit.commit |
| cas = {'fails': 0} |
| |
| def patched_cas(base_snapshot_uuid, snapshot, statistics): |
| if snapshot.commit_kind == "OVERWRITE" and cas['fails'] < K: |
| cas['fails'] += 1 |
| concurrent_fn(cas['fails']) |
| return False |
| return orig_cas(base_snapshot_uuid, snapshot, statistics) |
| |
| fsc.snapshot_commit.commit = patched_cas |
| |
| c.commit(messages) |
| c.close() |
| |
| self.assertEqual(cas['fails'], K, "expected exactly K forced conflicts") |
| return captured['provider'] |
| |
| def test_cache_applies_delta_when_concurrent_overwrite_hits_target_partition(self): |
| # Concurrent overwrites produce DELETE+ADD deltas. They can be applied |
| # to the cached target state just like APPEND deltas. |
| K = 2 |
| wb = self.table.new_batch_write_builder().overwrite({'f0': 1}) |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']})) |
| provider = self._run_with_conflicts( |
| c, w.prepare_commit(), K, |
| lambda i: self._overwrite_partition(1, f'z{i}')) |
| |
| self.assertEqual(provider.full_scan_count, 1) |
| self.assertEqual(provider.delta_probe_count, K) |
| self.assertEqual(provider.delta_apply_count, K) |
| |
| read_builder = self.table.new_read_builder() |
| actual = read_builder.new_read().to_pandas( |
| read_builder.new_scan().plan().splits()) |
| self.assertEqual(sorted(actual[actual['f0'] == 1]['f1'].tolist()), ['new']) |
| self.assertEqual(sorted(actual[actual['f0'] == 2]['f1'].tolist()), ['c']) |
| |
| def test_whole_table_overwrite_advances_by_append_delta(self): |
| # Whole-table overwrite (no partition filter) can still advance through |
| # APPEND deltas because every appended file belongs to the target state. |
| K = 2 |
| wb = self.table.new_batch_write_builder().overwrite() |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']})) |
| provider = self._run_with_conflicts( |
| c, w.prepare_commit(), K, |
| lambda i: self._append(pd.DataFrame({'f0': [99], 'f1': [f'x{i}']}))) |
| |
| self.assertEqual(provider.full_scan_count, 1) |
| self.assertEqual(provider.delta_probe_count, K) |
| self.assertEqual(provider.delta_apply_count, K) |
| |
| read_builder = self.table.new_read_builder() |
| actual = read_builder.new_read().to_pandas( |
| read_builder.new_scan().plan().splits()) |
| self.assertEqual(sorted(actual['f1'].tolist()), ['new']) |
| |
| def test_dynamic_partition_overwrite_reuses_cache(self): |
| # Dynamic-partition overwrite is scoped to the data's partitions, so an |
| # unrelated concurrent append lets the cache be reused. |
| pa_schema = pa.schema([('f0', pa.int32()), ('f1', pa.string())]) |
| schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['f0']) |
| self.catalog.create_table('test_db.t_dyn', schema, False) |
| table = self.catalog.get_table('test_db.t_dyn') |
| |
| def append(df): |
| wb = table.new_batch_write_builder() |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(df) |
| c.commit(w.prepare_commit()) |
| w.close() |
| c.close() |
| |
| append(pd.DataFrame({'f0': [1, 1, 2], 'f1': ['a', 'b', 'c']})) |
| |
| K = 3 |
| wb = table.new_batch_write_builder().overwrite() # dynamic: filter from data (f0=1) |
| w = wb.new_write() |
| c = wb.new_commit() |
| w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']})) |
| provider = self._run_with_conflicts( |
| c, w.prepare_commit(), K, |
| lambda i: append(pd.DataFrame({'f0': [99], 'f1': [f'x{i}']}))) |
| |
| self.assertEqual(provider.full_scan_count, 1) # target f0=1 untouched -> reuse |
| self.assertEqual(provider.delta_probe_count, K) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
