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 d817014975 Fix: partition index bucket overflow (#9977)
d817014975 is described below
commit d817014975d57db1c4b835519db28c2bb4856386
Author: Akash Reddy Jammula <[email protected]>
AuthorDate: Sat Sep 19 22:47:43 2026 -0700
Fix: partition index bucket overflow (#9977)
---
docs/docs/primary-key-table/data-distribution.md | 2 +-
docs/generated/core_configuration.html | 2 +-
.../main/java/org/apache/paimon/CoreOptions.java | 6 +-
.../apache/paimon/index/HashBucketAssigner.java | 1 +
.../org/apache/paimon/index/PartitionIndex.java | 33 ++++++++--
.../paimon/index/SimpleHashBucketAssigner.java | 7 ++-
.../apache/paimon/index/PartitionIndexTest.java | 70 ++++++++++++++++++++++
.../paimon/index/SimpleHashBucketAssignerTest.java | 11 ++++
.../pypaimon/common/options/core_options.py | 3 +-
paimon-python/pypaimon/index/dynamic_bucket.py | 23 ++++++-
.../pypaimon/tests/write/dynamic_bucket_test.py | 62 +++++++++++++++++--
.../write/simple_hash_bucket_assigner_test.py | 15 +++++
paimon-python/pypaimon/write/row_key_extractor.py | 9 ++-
13 files changed, 226 insertions(+), 18 deletions(-)
diff --git a/docs/docs/primary-key-table/data-distribution.md
b/docs/docs/primary-key-table/data-distribution.md
index 1ce17b35e6..71c5b05bc3 100644
--- a/docs/docs/primary-key-table/data-distribution.md
+++ b/docs/docs/primary-key-table/data-distribution.md
@@ -70,7 +70,7 @@ Do not configure `bucket-key` in this mode.
| --- | --- |
| `dynamic-bucket.target-row-num` | Target row count per bucket |
| `dynamic-bucket.initial-buckets` | Initial bucket count |
-| `dynamic-bucket.max-buckets` | Maximum bucket count |
+| `dynamic-bucket.max-buckets` | Maximum bucket count; use `-1` or a value
from `1` to `32768` |
:::warning Concurrent writers
diff --git a/docs/generated/core_configuration.html
b/docs/generated/core_configuration.html
index 4c8538b4f1..63bcf12a49 100644
--- a/docs/generated/core_configuration.html
+++ b/docs/generated/core_configuration.html
@@ -690,7 +690,7 @@ under the License.
<td><h5>dynamic-bucket.max-buckets</h5></td>
<td style="word-wrap: break-word;">-1</td>
<td>Integer</td>
- <td>Max buckets for a partition in dynamic bucket mode, It should
either be equal to -1 (unlimited), or it must be greater than 0 (fixed upper
bound).</td>
+ <td>Max buckets for a partition in dynamic bucket mode, It should
either be equal to -1 (unlimited), or it must be between 1 and 32768 (fixed
upper bound).</td>
</tr>
<tr>
<td><h5>dynamic-bucket.target-row-num</h5></td>
diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
index 3842623860..e561c8efc6 100644
--- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
+++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
@@ -1741,6 +1741,8 @@ public class CoreOptions implements Serializable {
"If the bucket is -1, for primary key table, is
dynamic bucket mode, "
+ "this option controls the target row
number for one bucket.");
+ public static final int MAX_DYNAMIC_BUCKETS = Short.MAX_VALUE + 1;
+
@Immutable
public static final ConfigOption<Integer> DYNAMIC_BUCKET_INITIAL_BUCKETS =
key("dynamic-bucket.initial-buckets")
@@ -1755,7 +1757,9 @@ public class CoreOptions implements Serializable {
.defaultValue(-1)
.withDescription(
"Max buckets for a partition in dynamic bucket
mode, It should "
- + "either be equal to -1 (unlimited), or
it must be greater than 0 (fixed upper bound).");
+ + "either be equal to -1 (unlimited), or
it must be between 1 and "
+ + MAX_DYNAMIC_BUCKETS
+ + " (fixed upper bound).");
public static final ConfigOption<Integer>
DYNAMIC_BUCKET_ASSIGNER_PARALLELISM =
key("dynamic-bucket.assigner-parallelism")
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
b/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
index 71cbb2211b..10bd7771ec 100644
--- a/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
+++ b/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
@@ -58,6 +58,7 @@ public class HashBucketAssigner implements BucketAssigner {
int assignId,
long targetBucketRowNumber,
int maxBucketsNum) {
+ PartitionIndex.validateMaxBuckets(maxBucketsNum);
this.snapshotManager = snapshotManager;
this.commitUser = commitUser;
this.indexFileHandler = indexFileHandler;
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/PartitionIndex.java
b/paimon-core/src/main/java/org/apache/paimon/index/PartitionIndex.java
index 51add39ac3..7a0a0802eb 100644
--- a/paimon-core/src/main/java/org/apache/paimon/index/PartitionIndex.java
+++ b/paimon-core/src/main/java/org/apache/paimon/index/PartitionIndex.java
@@ -36,7 +36,10 @@ import java.util.Map;
import java.util.Set;
import java.util.function.IntPredicate;
+import static org.apache.paimon.CoreOptions.DYNAMIC_BUCKET_MAX_BUCKETS;
+import static org.apache.paimon.CoreOptions.MAX_DYNAMIC_BUCKETS;
import static org.apache.paimon.index.HashIndexFile.HASH_INDEX;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
/** Bucket Index Per Partition. */
public class PartitionIndex {
@@ -69,6 +72,7 @@ public class PartitionIndex {
}
public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum) {
+ validateMaxBuckets(maxBucketsNum);
accessed = true;
// 1. is it a key that has appeared before
@@ -85,14 +89,14 @@ public class PartitionIndex {
Long number = entry.getValue();
if (number < targetBucketRowNumber) {
entry.setValue(number + 1);
- hash2Bucket.put(hash, (short) bucket.intValue());
+ hash2Bucket.put(hash, toBucketShort(bucket));
return bucket;
} else {
iterator.remove();
}
}
- int globalMaxBucketId = (maxBucketsNum == -1 ? Short.MAX_VALUE :
maxBucketsNum) - 1;
+ int globalMaxBucketId = (maxBucketsNum == -1 ? MAX_DYNAMIC_BUCKETS :
maxBucketsNum) - 1;
if (!bucketUpperBoundReached) {
// 3. create a new bucket
for (int i = 0; i <= globalMaxBucketId; i++) {
@@ -100,7 +104,7 @@ public class PartitionIndex {
nonFullBucketInformation.put(i, 1L);
totalBucketSet.add(i);
totalBucketArray.add(i);
- hash2Bucket.put(hash, (short) i);
+ hash2Bucket.put(hash, toBucketShort(i));
return i;
}
}
@@ -115,7 +119,7 @@ public class PartitionIndex {
// 4. exceed buckets upper bound
int bucket = ListUtils.pickRandomly(totalBucketArray);
- hash2Bucket.put(hash, (short) bucket);
+ hash2Bucket.put(hash, toBucketShort(bucket));
return bucket;
}
@@ -129,6 +133,7 @@ public class PartitionIndex {
Int2ShortHashMap.Builder mapBuilder = Int2ShortHashMap.builder();
Map<Integer, Long> buckets = new HashMap<>();
for (IndexManifestEntry file : files) {
+ short loadedBucket = toBucketShort(file.bucket());
try (IntIterator iterator =
indexFileHandler
.hashIndex(file.partition(), file.bucket())
@@ -137,7 +142,7 @@ public class PartitionIndex {
try {
int hash = iterator.next();
if (loadFilter.test(hash)) {
- mapBuilder.put(hash, (short) file.bucket());
+ mapBuilder.put(hash, loadedBucket);
}
if (bucketFilter.test(file.bucket())) {
buckets.compute(
@@ -154,4 +159,22 @@ public class PartitionIndex {
}
return new PartitionIndex(mapBuilder.build(), buckets,
targetBucketRowNumber);
}
+
+ static void validateMaxBuckets(int maxBucketsNum) {
+ checkArgument(
+ maxBucketsNum == -1 || (maxBucketsNum > 0 && maxBucketsNum <=
MAX_DYNAMIC_BUCKETS),
+ "'%s' must be -1 or between 1 and %s, but was %s.",
+ DYNAMIC_BUCKET_MAX_BUCKETS.key(),
+ MAX_DYNAMIC_BUCKETS,
+ maxBucketsNum);
+ }
+
+ static short toBucketShort(int bucket) {
+ checkArgument(
+ bucket >= 0 && bucket < MAX_DYNAMIC_BUCKETS,
+ "Dynamic bucket id must be between 0 and %s, but was %s.",
+ MAX_DYNAMIC_BUCKETS - 1,
+ bucket);
+ return (short) bucket;
+ }
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
b/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
index 964244415d..e0f0b56e85 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
@@ -30,6 +30,8 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
+import static org.apache.paimon.CoreOptions.MAX_DYNAMIC_BUCKETS;
+
/** When we need to overwrite the table, we should use this to avoid loading
index. */
public class SimpleHashBucketAssigner implements BucketAssigner {
@@ -42,6 +44,7 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
public SimpleHashBucketAssigner(
int numAssigners, int assignId, long targetBucketRowNumber, int
maxBucketsNum) {
+ PartitionIndex.validateMaxBuckets(maxBucketsNum);
this.numAssigners = numAssigners;
this.assignId = assignId;
this.targetBucketRowNumber = targetBucketRowNumber;
@@ -120,12 +123,12 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
}
return l + 1;
});
- hash2Bucket.put(hash, (short) currentBucket);
+ hash2Bucket.put(hash, PartitionIndex.toBucketShort(currentBucket));
return currentBucket;
}
private boolean loadNewBucket() {
- for (int i = 0; i < Short.MAX_VALUE; i++) {
+ for (int i = 0; i < MAX_DYNAMIC_BUCKETS; i++) {
if (isMyBucket(i) && !bucketInformation.containsKey(i)) {
// The new bucketId may still be larger than the upper
bound
if (-1 == maxBucketsNum || i <= maxBucketsNum - 1) {
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/PartitionIndexTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/PartitionIndexTest.java
new file mode 100644
index 0000000000..eadc062215
--- /dev/null
+++ b/paimon-core/src/test/java/org/apache/paimon/index/PartitionIndexTest.java
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.index;
+
+import org.apache.paimon.utils.Int2ShortHashMap;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.HashMap;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link PartitionIndex}. */
+class PartitionIndexTest {
+
+ @ParameterizedTest
+ @ValueSource(ints = {-1, 32768})
+ void testMaximumBucketCount(int maxBuckets) {
+ PartitionIndex index =
+ new PartitionIndex(new Int2ShortHashMap(), new HashMap<>(),
Long.MAX_VALUE);
+
+ int hash = 12345;
+ assertThat(index.assign(hash, bucket -> bucket == Short.MAX_VALUE,
maxBuckets))
+ .isEqualTo(Short.MAX_VALUE);
+ assertThat(index.assign(hash, bucket -> false,
maxBuckets)).isEqualTo(Short.MAX_VALUE);
+ }
+
+ @Test
+ void testRejectBucketCountAboveShortRange() {
+ PartitionIndex index =
+ new PartitionIndex(new Int2ShortHashMap(), new HashMap<>(),
Long.MAX_VALUE);
+
+ for (int i = 0; i < 2; i++) {
+ assertThatThrownBy(() -> index.assign(12345, bucket -> bucket >=
32768, 40000))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage(
+ "'dynamic-bucket.max-buckets' must be -1 or
between 1 and 32768, but was 40000.");
+ }
+ }
+
+ @Test
+ void testRejectExistingBucketAboveShortRange() {
+ HashMap<Integer, Long> buckets = new HashMap<>();
+ buckets.put(32768, 0L);
+ PartitionIndex index = new PartitionIndex(new Int2ShortHashMap(),
buckets, Long.MAX_VALUE);
+
+ assertThatThrownBy(() -> index.assign(12345, bucket -> false, 32768))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Dynamic bucket id must be between 0 and 32767,
but was 32768.");
+ }
+}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
index a8b2df1e52..c354eda5f5 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
@@ -86,6 +86,17 @@ public class SimpleHashBucketAssignerTest {
}
}
+ @ParameterizedTest
+ @ValueSource(ints = {-1, 32768})
+ public void testMaximumBucketId(int maxBucketsNum) {
+ SimpleHashBucketAssigner assigner =
+ new SimpleHashBucketAssigner(32768, 32767, 1, maxBucketsNum);
+ BinaryRow partition = BinaryRow.EMPTY_ROW;
+
+ assertThat(assigner.assign(partition, 1)).isEqualTo(Short.MAX_VALUE);
+ assertThat(assigner.assign(partition, 1)).isEqualTo(Short.MAX_VALUE);
+ }
+
@Test
public void testAssign() {
SimpleHashBucketAssigner simpleHashBucketAssigner =
diff --git a/paimon-python/pypaimon/common/options/core_options.py
b/paimon-python/pypaimon/common/options/core_options.py
index bb312781bf..801b73e531 100644
--- a/paimon-python/pypaimon/common/options/core_options.py
+++ b/paimon-python/pypaimon/common/options/core_options.py
@@ -240,7 +240,8 @@ class CoreOptions:
.int_type()
.default_value(-1)
.with_description(
- "In dynamic bucket mode, max buckets per partition. -1 means
unlimited."
+ "In dynamic bucket mode, max buckets per partition. It must be -1 "
+ "or between 1 and 32768."
)
)
diff --git a/paimon-python/pypaimon/index/dynamic_bucket.py
b/paimon-python/pypaimon/index/dynamic_bucket.py
index c5bbd08b7d..90a1c49230 100644
--- a/paimon-python/pypaimon/index/dynamic_bucket.py
+++ b/paimon-python/pypaimon/index/dynamic_bucket.py
@@ -32,6 +32,7 @@ HASH_INDEX = "HASH"
_ADD = 0
_DELETE = 1
SHORT_MAX_VALUE = 32767
+MAX_DYNAMIC_BUCKETS = SHORT_MAX_VALUE + 1
_SNAPSHOT_UNSET = object()
@@ -64,6 +65,22 @@ def is_my_bucket(bucket: int, num_assigners: int, assign_id:
int) -> bool:
return bucket % num_assigners == assign_id % num_assigners
+def validate_max_buckets(max_buckets_num: int) -> None:
+ if max_buckets_num != -1 and not 1 <= max_buckets_num <=
MAX_DYNAMIC_BUCKETS:
+ raise ValueError(
+ "'dynamic-bucket.max-buckets' must be -1 or between 1 and "
+ f"{MAX_DYNAMIC_BUCKETS}, but was {max_buckets_num}."
+ )
+
+
+def validate_bucket_id(bucket: int) -> None:
+ if not 0 <= bucket < MAX_DYNAMIC_BUCKETS:
+ raise ValueError(
+ "Dynamic bucket id must be between 0 and "
+ f"{MAX_DYNAMIC_BUCKETS - 1}, but was {bucket}."
+ )
+
+
def _iter_hashes(table, entry: IndexManifestEntry) -> Iterator[int]:
meta = entry.index_file
path = meta.external_path
@@ -125,6 +142,7 @@ class _PartitionIndex:
max_buckets_num: int,
max_bucket_id: int,
) -> Tuple[int, bool]:
+ validate_max_buckets(max_buckets_num)
if key_hash in self.hash_to_bucket:
return self.hash_to_bucket[key_hash], False
@@ -136,7 +154,7 @@ class _PartitionIndex:
del self.non_full_bucket_information[bucket]
global_max_bucket_id = (
- SHORT_MAX_VALUE if max_buckets_num == -1 else max_buckets_num
+ MAX_DYNAMIC_BUCKETS if max_buckets_num == -1 else max_buckets_num
) - 1
if not self.total_bucket_set or max_bucket_id < global_max_bucket_id:
for bucket in range(global_max_bucket_id + 1):
@@ -176,6 +194,7 @@ class HashBucketAssigner:
ignore_existing: bool = False,
snapshot=_SNAPSHOT_UNSET,
) -> None:
+ validate_max_buckets(max_buckets_num)
if not 0 <= assign_id < num_channels:
raise ValueError(
f"assign_id must be in [0, {num_channels}), got {assign_id}"
@@ -286,6 +305,7 @@ class HashBucketAssigner:
)
entries_by_bucket = {}
for entry in entries:
+ validate_bucket_id(entry.bucket)
previous = entries_by_bucket.get(entry.bucket)
if previous is not None:
raise RuntimeError(
@@ -347,6 +367,7 @@ class HashBucketAssigner:
)
remaining = set(requested_hashes)
for entry in entries:
+ validate_bucket_id(entry.bucket)
for key_hash in _iter_hashes(self.table, entry):
if key_hash not in remaining:
continue
diff --git a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
index ffe6c88319..be175f1e39 100644
--- a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
+++ b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
@@ -31,6 +31,7 @@ from pypaimon.index.dynamic_bucket import (
_iter_hashes,
compute_assigner,
to_signed_int32,
+ validate_bucket_id,
)
from pypaimon.index.index_file_handler import IndexFileHandler
from pypaimon.index.index_file_meta import IndexFileMeta
@@ -191,12 +192,65 @@ class DynamicBucketTest(unittest.TestCase):
bucket = index.assign(
key_hash=1,
- bucket_filter=lambda candidate: candidate == 32766,
+ bucket_filter=lambda candidate: candidate == 32767,
max_buckets_num=-1,
- max_bucket_id=32765,
+ max_bucket_id=32766,
)
- self.assertEqual(bucket, 32766)
+ self.assertEqual(bucket, 32767)
+
+ def test_rejects_bucket_count_above_java_short_range(self):
+ index = _PartitionIndex({}, {}, 1)
+
+ with self.assertRaisesRegex(
+ ValueError,
+ "'dynamic-bucket.max-buckets' must be -1 or between 1 and 32768",
+ ):
+ index.assign(
+ key_hash=1,
+ bucket_filter=lambda candidate: candidate >= 32768,
+ max_buckets_num=40000,
+ max_bucket_id=32767,
+ )
+
+ def test_rejects_bucket_id_above_java_short_range(self):
+ with self.assertRaisesRegex(
+ ValueError,
+ "Dynamic bucket id must be between 0 and 32767, but was 32768",
+ ):
+ validate_bucket_id(32768)
+
+ def test_restore_rejects_bucket_id_above_java_short_range(self):
+ assigner = HashBucketAssigner(
+ table=Mock(),
+ num_channels=1,
+ num_assigners=1,
+ assign_id=0,
+ target_bucket_row_number=1,
+ max_buckets_num=-1,
+ snapshot=Mock(),
+ )
+ entry = Mock()
+ entry.bucket = 32768
+
+ with patch(
+ 'pypaimon.index.dynamic_bucket.IndexFileHandler'
+ ) as handler:
+ handler.return_value.scan.return_value = [entry]
+ with patch.object(
+ assigner, '_active_data_buckets', return_value=set()
+ ):
+ with self.assertRaisesRegex(
+ ValueError,
+ "Dynamic bucket id must be between 0 and 32767",
+ ):
+ assigner._load_partition((), {1})
+
+ with self.assertRaisesRegex(
+ ValueError,
+ "Dynamic bucket id must be between 0 and 32767",
+ ):
+ assigner._restore_requested_hashes((), {1}, {})
def test_assigner_rejects_record_owned_by_another_writer(self):
with tempfile.TemporaryDirectory() as root:
@@ -238,7 +292,7 @@ class DynamicBucketTest(unittest.TestCase):
key_hash=1,
bucket_filter=lambda _: False,
max_buckets_num=-1,
- max_bucket_id=32766,
+ max_bucket_id=32767,
)
def test_corrupt_hash_index_rejects_trailing_bytes(self):
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 6a7b78c2bd..0e27dde077 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
@@ -56,6 +56,21 @@ class SimpleHashBucketAssignerTest(unittest.TestCase):
for b in buckets[100:]:
self.assertEqual(b, 0)
+ def test_maximum_bucket_id(self):
+ for max_buckets in [-1, 32768]:
+ with self.subTest(max_buckets=max_buckets):
+ assigner = SimpleHashBucketAssigner(
+ 32768, 32767, 1, max_buckets)
+ self.assertEqual(assigner.assign((), 1), 32767)
+ self.assertEqual(assigner.assign((), 1), 32767)
+
+ def test_rejects_bucket_count_above_java_short_range(self):
+ with self.assertRaisesRegex(
+ ValueError,
+ "'dynamic-bucket.max-buckets' must be -1 or between 1 and 32768",
+ ):
+ SimpleHashBucketAssigner(1, 0, 100, 40000)
+
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]),
diff --git a/paimon-python/pypaimon/write/row_key_extractor.py
b/paimon-python/pypaimon/write/row_key_extractor.py
index 827b9f08b2..d5053e983f 100644
--- a/paimon-python/pypaimon/write/row_key_extractor.py
+++ b/paimon-python/pypaimon/write/row_key_extractor.py
@@ -26,7 +26,11 @@ import numpy as np
import pyarrow as pa
from pypaimon.common.options.core_options import CoreOptions
-from pypaimon.index.dynamic_bucket import SHORT_MAX_VALUE, is_my_bucket
+from pypaimon.index.dynamic_bucket import (
+ MAX_DYNAMIC_BUCKETS,
+ is_my_bucket,
+ validate_max_buckets,
+)
from pypaimon.schema.table_schema import TableSchema
from pypaimon.table.bucket_mode import BucketMode
from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer
@@ -376,7 +380,7 @@ class _SimplePartitionIndex:
def _load_new_bucket(
self, max_buckets_num: int, num_assigners: int, assign_id: int
) -> None:
- for i in range(SHORT_MAX_VALUE):
+ for i in range(MAX_DYNAMIC_BUCKETS):
if is_my_bucket(i, num_assigners, assign_id) and (
i not in self.bucket_information
):
@@ -392,6 +396,7 @@ class _SimplePartitionIndex:
class SimpleHashBucketAssigner:
def __init__(self, num_assigners, assign_id, target_bucket_row_number,
max_buckets_num):
+ validate_max_buckets(max_buckets_num)
self.num_assigners = num_assigners
self.assign_id = assign_id
self.target_bucket_row_number = target_bucket_row_number
