[python] Check the bucket count of existing files when writing

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._bucket_count_mismatches: 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, bucket):
+        previous = self._bucket_count_mismatches.get(partition, {}).get(bucket)
+        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()
@@ -218,7 +234,9 @@ def _create_data_writer(self, partition: Tuple, bucket: int, options: CoreOption
 
         def max_seq_number():
             default = -1 if self.table.is_primary_key_table else 1
-            return self._seq_number_stats(partition).get(bucket, default)
+            stats = self._seq_number_stats(partition)
+            self._check_bucket_count(partition, bucket)
+            return stats.get(bucket, default)
 
         # Dedicated Blob files are an append-table layout. PK tables require
         # managed packs and references attached to their key-value Parquet files.
@@ -467,6 +485,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._bucket_count_mismatches.clear()
 
     def close(self):
         """Close all data writers and clean up resources."""
@@ -515,6 +534,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 +555,16 @@ def _load_seq_number_stats(self, partition: Tuple) -> dict:
         scanner.skip_level0 = False
 
         max_seq_numbers = {}
+        mismatches = {}
+        expected = self._runtime_total_buckets.get(partition, self.options.bucket())
         for entry in scanner.plan_files():
+            if entry.total_buckets != expected:
+                mismatches[entry.bucket] = 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._bucket_count_mismatches[partition] = mismatches
         return max_seq_numbers
 
 

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,45 @@ 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), ('partition', ['dt'], False), ('overwrite', ['dt'], True),
+    ])
+    @pytest.mark.python_write
+    def test_write_bucket_count(self, name, partition_keys, overwrite):
+        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': '2'}), 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', '4')], False)
+        table = self.catalog.get_table(identifier)
+        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(data)
+                commit.commit(writer.prepare_commit())
+            else:
+                with self.assertRaises(RuntimeError) as error:
+                    writer.write_arrow(data)
+                self.assertEqual(
+                    'Try to write {} with a new bucket num 4, but the previous bucket num is 2. '
+                    'Please switch to batch mode, and perform INSERT OVERWRITE to rescale '
+                    'current data layout first.'.format(
+                        'partition {dt=p}' if partition_keys else 'table'), str(error.exception))
+        finally:
+            writer.close()
+            commit.close()
+        if overwrite:
+            self._commit_arrow(table, data)
+        self.assertEqual(data.to_pylist(), self._read_sorted(table, 'id').to_pylist())
+
     @pytest.mark.python_write
     def test_multi_prepare_commit_ao(self):
         schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
@@ -891,6 +931,37 @@ def test_postpone_file_store_write_validates_runtime_bucket_count(self):
         finally:
             write.abort()
 
+    @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:
+                writer.write(('p',), 3, data.to_batches()[0], 4)
+                with self.assertRaisesRegex(RuntimeError, 'new bucket num 4, but the previous bucket num is 2'):
+                    writer.write(('p',), 0, 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',

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._bucket_count_mismatches = {}
             writer._sequence_read_table = lambda: table
             assert writer._load_seq_number_stats(()) == {3: 9, 4: 7}
         elif operation == 'dynamic':