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':