From 661aeb705d1727e2d91ee67b7f6e5958539e729a Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 27 Oct 2026 09:14:50 -0700 Subject: [PATCH] [python] Check the bucket count of existing files when writing PyPaimon restored the sequence numbers of a partition without checking the bucket count its existing files were written with. After the bucket option was changed, for example from 2 to 4, writes to an existing partition used the new bucket count, so primary keys already stored under the old layout were written again to other buckets and the table returned duplicate keys. As in Java AbstractFileStoreWrite, writing to a partition whose existing files have a different total bucket number now fails with Java's message, asking to rescale with INSERT OVERWRITE first. As in Java, the check runs whenever a writer of a bucket is created, also for writers that do not restore sequence numbers, such as Blob writers. It uses the existing per-partition restore scan, and overwrite, which does not restore existing files, is not affected. Java's writer only checks the files of the bucket it writes, so a write that only reaches buckets with no files yet, such as bucket 3 after changing from 2 to 4 buckets, is only caught by its commit conflict check, which a default batch commit does not run. PyPaimon checks every file of the partition instead and, like Java's commit check (ConflictDetection.checkBucketKeepSame), skips files without a fixed bucket count, such as the postpone files left after changing a postpone table to a fixed bucket count. --- .../python/write/file_store_write.py | 31 +++++ .../pypaimon/tests/write/table_write_test.py | 130 ++++++++++++++++++ .../tests/write/trusted_scan_metadata_test.py | 4 + 3 files changed, 165 insertions(+) diff --git a/paimon-python/pypaimon/execution/python/write/file_store_write.py b/paimon-python/pypaimon/execution/python/write/file_store_write.py index af5bb5a127d4..527486a055f2 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_write.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_write.py @@ -65,6 +65,7 @@ def __init__(self, table, commit_user): self._ever_written = False self._prepared_writer_lifecycle = None self.max_seq_numbers: dict = {} + self._previous_total_buckets: dict = {} self.restore_snapshot_id = None self.write_cols = None self.blob_consumer = None @@ -196,6 +197,21 @@ def _check_runtime_bucket(self, partition, bucket, total_buckets): ) self._runtime_total_buckets[partition] = total_buckets + def _check_bucket_count(self, partition): + previous = self._previous_total_buckets.get(partition) + if previous is None: + return + part_info = "table" + if self.table.partition_keys: + values = self.table.path_factory()._canonical_partition(partition) + part_info = "partition {" + ", ".join( + "{}={}".format(key, value) for key, value in zip(self.table.partition_keys, values)) + "}" + expected = self._runtime_total_buckets.get(partition, self.options.bucket()) + raise RuntimeError( + "Try to write {} with a new bucket num {}, but the previous bucket num is {}. " + "Please switch to batch mode, and perform INSERT OVERWRITE to rescale current data layout first.".format( + part_info, expected, previous)) + def _create_data_writer(self, partition: Tuple, bucket: int, options: CoreOptions) -> DataWriter: row_limit = options.target_file_row_num() max_value = CoreOptions.TARGET_FILE_ROW_NUM.default_value() @@ -216,6 +232,11 @@ def _create_data_writer(self, partition: Tuple, bucket: int, options: CoreOption "based file rolling only for data-evolution append and postpone tables; " "unset it or write with Java/Flink/Spark.") + # As in Java, every writer of a restored bucket checks the bucket count. + if bucket >= 0 and self.table.bucket_mode() != BucketMode.BUCKET_UNAWARE: + self._seq_number_stats(partition) + self._check_bucket_count(partition) + def max_seq_number(): default = -1 if self.table.is_primary_key_table else 1 return self._seq_number_stats(partition).get(bucket, default) @@ -467,6 +488,7 @@ def _release_prepared_files(self): # bucket opened after this checkpoint must restore from the latest # snapshot rather than reuse another bucket's earlier partition scan. self.max_seq_numbers.clear() + self._previous_total_buckets.clear() def close(self): """Close all data writers and clean up resources.""" @@ -515,6 +537,9 @@ def _sequence_read_table(self): return self.table._copy_with_snapshot(self.table.snapshot_manager().get_latest_snapshot()) def _load_seq_number_stats(self, partition: Tuple) -> dict: + # Overwrite ignores previous files, like Java's ignorePreviousFiles. + if self.restore_snapshot_id == 0: + return {} read_builder = self._sequence_read_table().new_read_builder() predicate_builder = read_builder.new_predicate_builder() sub_predicates = [] @@ -533,11 +558,17 @@ def _load_seq_number_stats(self, partition: Tuple) -> dict: scanner.skip_level0 = False max_seq_numbers = {} + previous = None + expected = self._runtime_total_buckets.get(partition, self.options.bucket()) for entry in scanner.plan_files(): + # Check the whole partition, skipping postpone and dynamic files, like Java. + if previous is None and 0 < entry.total_buckets != expected: + previous = entry.total_buckets current_seq_num = entry.file.max_sequence_number existing_max = max_seq_numbers.get(entry.bucket, -1) if current_seq_num > existing_max: max_seq_numbers[entry.bucket] = current_seq_num + self._previous_total_buckets[partition] = previous return max_seq_numbers diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index dabfa5898a41..a2798a95699d 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -28,6 +28,7 @@ from pypaimon.execution.python.read.table_scan import PythonTableScan from pypaimon import CatalogFactory, Schema from pypaimon.schema.arrow_schema import read_arrow_schema +from pypaimon.schema.schema_change import SchemaChange import pyarrow as pa import pytest from parameterized import parameterized @@ -635,6 +636,85 @@ def test_write_row_dynamic_bucket_primary_key_table(self): self.assertEqual( expected.sort_by(sort_keys), self._read_sorted(table, sort_keys)) + @parameterized.expand([ + ('table', [], False, 2, 4), ('partition', ['dt'], False, 2, 4), + ('overwrite', ['dt'], True, 2, 4), ('fewer_buckets', ['dt'], False, 4, 2), + ]) + @pytest.mark.python_write + def test_write_bucket_count(self, name, partition_keys, overwrite, previous, new): + identifier = 'default.bucket_count_' + name + self.catalog.create_table(identifier, Schema.from_pyarrow_schema( + self.postpone_pa_schema, partition_keys=partition_keys, + primary_keys=['id'] + partition_keys, options={'bucket': str(previous)}), False) + table = self.catalog.get_table(identifier) + data = pa.Table.from_pydict({ + 'id': list(range(8)), 'dt': ['p'] * 8, 'value': ['old'] * 8, + }, schema=self.postpone_pa_schema) + self._commit_arrow(table, data) + self.catalog.alter_table(identifier, [SchemaChange.set_option('bucket', str(new))], False) + table = self.catalog.get_table(identifier) + rows_to_write = [data] + if not overwrite and new > previous: + buckets = create_row_key_extractor(table).extract_partition_bucket_batch(data.to_batches()[0])[1] + # Rows of buckets that have no files yet. + new_bucket_rows = data.filter(pa.array([bucket >= previous for bucket in buckets])) + self.assertGreater(new_bucket_rows.num_rows, 0) + rows_to_write.append(new_bucket_rows) + for rows in rows_to_write: + builder = table.new_batch_write_builder() + if overwrite: + builder.overwrite({'dt': 'p'}) + writer, commit = builder.new_write(), builder.new_commit() + try: + if overwrite: + writer.write_arrow(rows) + commit.commit(writer.prepare_commit()) + else: + with self.assertRaises(RuntimeError) as error: + writer.write_arrow(rows) + self.assertEqual( + 'Try to write {} with a new bucket num {}, but the previous bucket num is {}. ' + 'Please switch to batch mode, and perform INSERT OVERWRITE to rescale ' + 'current data layout first.'.format( + 'partition {dt=p}' if partition_keys else 'table', new, previous), + str(error.exception)) + finally: + writer.close() + commit.close() + if overwrite: + self._commit_arrow(table, data) + expected = data + if partition_keys and not overwrite: + # A new partition uses the new bucket count. + new_partition = pa.Table.from_pydict({ + 'id': list(range(8)), 'dt': ['q'] * 8, 'value': ['new'] * 8, + }, schema=self.postpone_pa_schema) + self._commit_arrow(table, new_partition) + expected = pa.concat_tables([data, new_partition]) + sort_keys = [('dt', 'ascending'), ('id', 'ascending')] + self.assertEqual(expected.to_pylist(), self._read_sorted(table, sort_keys).to_pylist()) + + @pytest.mark.python_write + def test_write_bucket_count_with_blob_writer(self): + from pypaimon.execution.python.write.file_store_write import FileStoreWrite + + identifier = 'default.bucket_count_blob_writer' + self.catalog.create_table(identifier, Schema.from_pyarrow_schema( + self.postpone_pa_schema, partition_keys=['dt'], options={'bucket': '2', 'bucket-key': 'id'}), False) + data = pa.Table.from_pydict({ + 'id': list(range(8)), 'dt': ['p'] * 8, 'value': ['old'] * 8, + }, schema=self.postpone_pa_schema) + self._commit_arrow(self.catalog.get_table(identifier), data) + self.catalog.alter_table(identifier, [SchemaChange.set_option('bucket', '4')], False) + writer = self.catalog.get_table(identifier).new_batch_write_builder().new_write() + try: + # Writers that do not restore sequence numbers, such as Blob writers, are checked too. + with patch.object(FileStoreWrite, '_has_blob_columns', return_value=True): + with self.assertRaisesRegex(RuntimeError, 'new bucket num 4, but the previous bucket num is 2'): + writer.write_arrow(data) + finally: + writer.close() + @pytest.mark.python_write def test_multi_prepare_commit_ao(self): schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) @@ -891,6 +971,56 @@ def test_postpone_file_store_write_validates_runtime_bucket_count(self): finally: write.abort() + @pytest.mark.python_write + def test_write_fixed_bucket_after_postpone(self): + table = self._create_postpone_table( + 'default.fixed_bucket_after_postpone', self.postpone_pa_schema, + partition_keys=['dt'], primary_keys=['id', 'dt']) + data = pa.Table.from_pydict({ + 'id': list(range(8)), 'dt': ['p'] * 8, 'value': ['old'] * 8, + }, schema=self.postpone_pa_schema) + self._commit_arrow(table, data) + self.catalog.alter_table( + 'default.fixed_bucket_after_postpone', [SchemaChange.set_option('bucket', '2')], False) + # Postpone files do not count, as in Java. + self._commit_arrow(self.catalog.get_table('default.fixed_bucket_after_postpone'), data) + self.catalog.alter_table( + 'default.fixed_bucket_after_postpone', [SchemaChange.set_option('bucket', '4')], False) + # Nor do they exempt the fixed bucket files of the same partition. + with self.assertRaisesRegex(RuntimeError, 'new bucket num 4, but the previous bucket num is 2'): + self._commit_arrow(self.catalog.get_table('default.fixed_bucket_after_postpone'), data) + + @pytest.mark.python_write + def test_postpone_checks_restored_bucket_count(self): + from pypaimon.execution.python.write.file_store_write import PostponeFixedBucketFileStoreWrite + + table = self._create_postpone_table( + 'default.postpone_restored_bucket_count', self.postpone_pa_schema, + partition_keys=['dt'], primary_keys=['id', 'dt'], options={ + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': '2', + }) + data = pa.Table.from_pydict({ + 'id': list(range(8)), 'dt': ['p'] * 8, 'value': ['old'] * 8, + }, schema=self.postpone_pa_schema) + self._commit_arrow(table, data, fixed_bucket=True) + writer = PostponeFixedBucketFileStoreWrite(table, 'test-user') + try: + with patch.object(writer, '_load_seq_number_stats', wraps=writer._load_seq_number_stats) as restore: + for bucket in (3, 0): + with self.assertRaisesRegex(RuntimeError, 'new bucket num 4, but the previous bucket num is 2'): + writer.write(('p',), bucket, data.to_batches()[0], 4) + restore.assert_called_once_with(('p',)) + finally: + writer.abort() + # Overwrite ignores the existing layout. + writer = PostponeFixedBucketFileStoreWrite(table, 'test-user') + writer.restore_snapshot_id = 0 + try: + writer.write(('p',), 0, data.to_batches()[0], 4) + finally: + writer.abort() + def test_postpone_batch_write_builder_keeps_postpone_mode(self): table = self._create_postpone_table( 'default.test_postpone_default_builder', diff --git a/paimon-python/pypaimon/tests/write/trusted_scan_metadata_test.py b/paimon-python/pypaimon/tests/write/trusted_scan_metadata_test.py index 68566cd204e3..045e1af343e8 100644 --- a/paimon-python/pypaimon/tests/write/trusted_scan_metadata_test.py +++ b/paimon-python/pypaimon/tests/write/trusted_scan_metadata_test.py @@ -61,6 +61,10 @@ def plan_files(): if operation == 'sequence': writer = object.__new__(FileStoreWrite) writer.table = table + writer.options = Mock(bucket=lambda: 8) + writer.restore_snapshot_id = None + writer._runtime_total_buckets = {} + writer._previous_total_buckets = {} writer._sequence_read_table = lambda: table assert writer._load_seq_number_stats(()) == {3: 9, 4: 7} elif operation == 'dynamic':