This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new eaeac0b93c [python] Fix bucket registration for dynamic bucket
overflow (#9788)
eaeac0b93c is described below
commit eaeac0b93c499eb13798845f10591cb83e66eced
Author: QuakeWang <[email protected]>
AuthorDate: Mon Sep 14 10:52:10 2026 +0800
[python] Fix bucket registration for dynamic bucket overflow (#9788)
---
.../write/simple_hash_bucket_assigner_test.py | 39 ++++++++++++++++++++++
paimon-python/pypaimon/write/row_key_extractor.py | 2 +-
2 files changed, 40 insertions(+), 1 deletion(-)
diff --git
a/paimon-python/pypaimon/tests/write/simple_hash_bucket_assigner_test.py
b/paimon-python/pypaimon/tests/write/simple_hash_bucket_assigner_test.py
index 8bc19fd084..6a7b78c2bd 100644
--- a/paimon-python/pypaimon/tests/write/simple_hash_bucket_assigner_test.py
+++ b/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 @@ class SimpleHashBucketAssignerTest(unittest.TestCase):
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()
diff --git a/paimon-python/pypaimon/write/row_key_extractor.py
b/paimon-python/pypaimon/write/row_key_extractor.py
index 3313620d7d..827b9f08b2 100644
--- a/paimon-python/pypaimon/write/row_key_extractor.py
+++ b/paimon-python/pypaimon/write/row_key_extractor.py
@@ -353,7 +353,6 @@ class _SimplePartitionIndex:
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 @@ class _SimplePartitionIndex:
):
if max_buckets_num == -1 or i <= max_buckets_num - 1:
self.current_bucket = i
+ self.bucket_list.append(i)
return
return
raise RuntimeError(
