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 17e9ad33d2 [core] Fix dynamic bucket allocation across partitions
(#9193)
17e9ad33d2 is described below
commit 17e9ad33d2ab94f702f0ae58f5d5030949cca146
Author: Arnav Balyan <[email protected]>
AuthorDate: Thu Aug 13 19:29:40 2026 +0530
[core] Fix dynamic bucket allocation across partitions (#9193)
---
.../org/apache/paimon/index/HashBucketAssigner.java | 6 +-----
.../org/apache/paimon/index/PartitionIndex.java | 8 +++++---
.../paimon/index/SimpleHashBucketAssigner.java | 21 +++++++--------------
.../apache/paimon/index/HashBucketAssignerTest.java | 21 +++++++++++++++++++++
.../paimon/index/SimpleHashBucketAssignerTest.java | 14 +++++---------
5 files changed, 39 insertions(+), 31 deletions(-)
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 93e05d6c94..71cbb2211b 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
@@ -46,7 +46,6 @@ public class HashBucketAssigner implements BucketAssigner {
private final int assignId;
private final long targetBucketRowNumber;
private final int maxBucketsNum;
- private int maxBucketId;
private final Map<BinaryRow, PartitionIndex> partitionIndex;
@@ -88,13 +87,10 @@ public class HashBucketAssigner implements BucketAssigner {
this.partitionIndex.put(partition, index);
}
- int assigned = index.assign(hash, this::isMyBucket, maxBucketsNum,
maxBucketId);
+ int assigned = index.assign(hash, this::isMyBucket, maxBucketsNum);
if (LOG.isDebugEnabled()) {
LOG.debug("Assign {} to the partition {} key hash {}", assigned,
partition, hash);
}
- if (assigned > maxBucketId) {
- maxBucketId = assigned;
- }
return assigned;
}
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 b0404b9d13..51add39ac3 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
@@ -49,6 +49,7 @@ public class PartitionIndex {
public final List<Integer> totalBucketArray;
private final long targetBucketRowNumber;
+ private boolean bucketUpperBoundReached;
public boolean accessed;
@@ -67,7 +68,7 @@ public class PartitionIndex {
this.accessed = true;
}
- public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum,
int maxBucketId) {
+ public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum) {
accessed = true;
// 1. is it a key that has appeared before
@@ -92,7 +93,7 @@ public class PartitionIndex {
}
int globalMaxBucketId = (maxBucketsNum == -1 ? Short.MAX_VALUE :
maxBucketsNum) - 1;
- if (totalBucketSet.isEmpty() || maxBucketId < globalMaxBucketId) {
+ if (!bucketUpperBoundReached) {
// 3. create a new bucket
for (int i = 0; i <= globalMaxBucketId; i++) {
if (bucketFilter.test(i) && !totalBucketSet.contains(i)) {
@@ -107,8 +108,9 @@ public class PartitionIndex {
throw new RuntimeException(
String.format(
"Too more bucket %s, you should increase
target bucket row number %s.",
- maxBucketId, targetBucketRowNumber));
+ globalMaxBucketId, targetBucketRowNumber));
}
+ bucketUpperBoundReached = true;
}
// 4. exceed buckets upper bound
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 a0bc5b427e..964244415d 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
@@ -37,7 +37,6 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
private final int assignId;
private final long targetBucketRowNumber;
private final int maxBucketsNum;
- private int maxBucketId;
private final Map<BinaryRow, SimplePartitionIndex> partitionIndex;
@@ -58,11 +57,7 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
index = new SimplePartitionIndex();
this.partitionIndex.put(partition, index);
}
- int assigned = index.assign(hash);
- if (assigned > maxBucketId) {
- maxBucketId = assigned;
- }
- return assigned;
+ return index.assign(hash);
}
@Override
@@ -82,6 +77,7 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
private final Map<Integer, Long> bucketInformation;
private final List<Integer> bucketList;
private int currentBucket;
+ private boolean bucketUpperBoundReached;
private SimplePartitionIndex() {
bucketInformation = new LinkedHashMap<>();
@@ -104,11 +100,8 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
});
if (num >= targetBucketRowNumber) {
- if (-1 == maxBucketsNum
- || bucketInformation.isEmpty()
- || maxBucketId < maxBucketsNum - 1) {
- loadNewBucket();
- } else {
+ if (bucketUpperBoundReached || !loadNewBucket()) {
+ bucketUpperBoundReached = true;
currentBucket = ListUtils.pickRandomly(bucketList);
}
}
@@ -131,16 +124,16 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
return currentBucket;
}
- private void loadNewBucket() {
+ private boolean loadNewBucket() {
for (int i = 0; i < Short.MAX_VALUE; i++) {
if (isMyBucket(i) && !bucketInformation.containsKey(i)) {
// The new bucketId may still be larger than the upper
bound
if (-1 == maxBucketsNum || i <= maxBucketsNum - 1) {
currentBucket = i;
- return;
+ return true;
}
// No need to enter the next iteration when upper bound
exceeded
- return;
+ return false;
}
}
throw new RuntimeException(
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/HashBucketAssignerTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/HashBucketAssignerTest.java
index 4c5415d45d..d750dc83f3 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/HashBucketAssignerTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/HashBucketAssignerTest.java
@@ -35,10 +35,13 @@ import org.junit.jupiter.params.provider.ValueSource;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
import static org.apache.paimon.io.DataFileTestUtils.row;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.Assertions.entry;
/** Test for {@link HashBucketAssigner}. */
public class HashBucketAssignerTest extends PrimaryKeyTableTestBase {
@@ -150,6 +153,24 @@ public class HashBucketAssignerTest extends
PrimaryKeyTableTestBase {
}
}
+ @Test
+ public void testEachPartitionUsesAllBucketsWithUpperBound() {
+ HashBucketAssigner assigner = createAssigner(1, 1, 0, 4);
+
+ Map<Integer, Integer> firstPartition = new HashMap<>();
+ for (int hash = 0; hash < 20; hash++) {
+ firstPartition.merge(assigner.assign(row(1), hash), 1,
Integer::sum);
+ }
+ assertThat(firstPartition).containsOnly(entry(0, 5), entry(1, 5),
entry(2, 5), entry(3, 5));
+
+ Map<Integer, Integer> secondPartition = new HashMap<>();
+ for (int hash = 0; hash < 20; hash++) {
+ secondPartition.merge(assigner.assign(row(2), hash), 1,
Integer::sum);
+ }
+ assertThat(secondPartition)
+ .containsOnly(entry(0, 5), entry(1, 5), entry(2, 5), entry(3,
5));
+ }
+
@Test
public void testAssignWithUpperBoundMultiAssigners() {
HashBucketAssigner assigner0 = createAssigner(2, 2, 0, 3);
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 70061941f6..a8b2df1e52 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
@@ -29,6 +29,7 @@ import java.util.Map;
import static org.apache.paimon.io.DataFileTestUtils.row;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.entry;
/** Tests for {@link SimpleHashBucketAssigner}. */
public class SimpleHashBucketAssignerTest {
@@ -56,24 +57,19 @@ public class SimpleHashBucketAssignerTest {
@Test
public void testSecondPartitionAfterTheCapIsExhausted() {
- // maxBucketId is shared across partitions, so by the time a later
partition starts, the
- // create-a-new-bucket branch is already closed for it and it goes
straight to
- // pickRandomly. Its own first bucket must be in the pool or that call
has nothing to
- // choose from.
SimpleHashBucketAssigner assigner = new SimpleHashBucketAssigner(1, 0,
100, 4);
- for (int hash = 0; hash < 500; hash++) {
+ for (int hash = 0; hash < 400; hash++) {
assigner.assign(row(1), hash);
}
Map<Integer, Integer> rowsPerBucket = new HashMap<>();
- for (int hash = 0; hash < 300; hash++) {
+ for (int hash = 0; hash < 400; hash++) {
rowsPerBucket.merge(assigner.assign(row(2), hash), 1,
Integer::sum);
}
- assertThat(rowsPerBucket.keySet()).allMatch(bucket -> bucket >= 0 &&
bucket < 4);
-
assertThat(rowsPerBucket.values().stream().mapToInt(Integer::intValue).sum())
- .isEqualTo(300);
+ assertThat(rowsPerBucket)
+ .containsOnly(entry(0, 100), entry(1, 100), entry(2, 100),
entry(3, 100));
}
@Test
