[python] Fix bucket registration for dynamic bucket overflow

paimon-python/pypaimon/tests/write/simple_hash_bucket_assigner_test.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import unittest
+from unittest.mock import patch
 
 from pypaimon.write.row_key_extractor import SimpleHashBucketAssigner
 
@@ -55,6 +56,44 @@ def test_assign_with_same_hash(self):
                 for b in buckets[100:]:
                     self.assertEqual(b, 0)
 
+    def test_register_each_bucket_once(self):
+        for num_assigners, assign_id, max_buckets, expected in [
+            (1, 0, 1, [0, 0, 0, 0, 0, 0]),
+            (1, 0, 3, [0, 0, 1, 1, 2, 2]),
+            (1, 0, -1, [0, 0, 1, 1, 2, 2]),
+            (2, 1, 6, [1, 1, 3, 3, 5, 5]),
+        ]:
+            with self.subTest(num_assigners=num_assigners, assign_id=assign_id,
+                              max_buckets=max_buckets):
+                assigner = SimpleHashBucketAssigner(num_assigners, assign_id, 2, max_buckets)
+                for h, expected_bucket in enumerate(expected):
+                    self.assertEqual(assigner.assign((), h), expected_bucket)
+                    index = assigner._partition_index[()]
+                    self.assertCountEqual(index.bucket_list, index.bucket_information)
+
+    def test_overflow_uses_all_registered_buckets(self):
+        assigner = SimpleHashBucketAssigner(1, 0, 2, 3)
+        initial = [assigner.assign((), h) for h in range(6)]
+        self.assertEqual(initial, [0, 0, 1, 1, 2, 2])
+        index = assigner._partition_index[()]
+
+        with patch('pypaimon.write.row_key_extractor.random.choice') as choice:
+            for h, selected in enumerate([0, 1, 2, 0, 1, 2], start=6):
+                choice.return_value = selected
+                self.assertEqual(assigner.assign((), h), selected)
+                choice.assert_called_with([0, 1, 2])
+                self.assertCountEqual(index.bucket_list, [0, 1, 2])
+            self.assertEqual(choice.call_count, 6)
+            self.assertEqual(index.bucket_information, {0: 4, 1: 4, 2: 4})
+            self.assertEqual(assigner.max_bucket_id, 2)
+
+            choice.reset_mock()
+            repeated = [assigner.assign((), h) for h in range(12)]
+            self.assertEqual(repeated, initial + [0, 1, 2, 0, 1, 2])
+            choice.assert_not_called()
+            self.assertEqual(index.bucket_information, {0: 4, 1: 4, 2: 4})
+            self.assertCountEqual(index.bucket_list, [0, 1, 2])
+
 
 if __name__ == '__main__':
     unittest.main()

paimon-python/pypaimon/write/row_key_extractor.py

@@ -353,7 +353,6 @@ def assign(
             return assigned, max(max_bucket_id, assigned)
 
         if self.current_bucket not in self.bucket_information:
-            self.bucket_list.append(self.current_bucket)
             self.bucket_information[self.current_bucket] = 0
         num = self.bucket_information[self.current_bucket]
 
@@ -383,6 +382,7 @@ def _load_new_bucket(
             ):
                 if max_buckets_num == -1 or i <= max_buckets_num - 1:
                     self.current_bucket = i
+                    self.bucket_list.append(i)
                     return
                 return
         raise RuntimeError(