cache the maxBucketId

Author: liyubin117

paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java

@@ -46,6 +46,7 @@ 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;
 
@@ -87,11 +88,14 @@ public int assign(BinaryRow partition, int hash) {
             this.partitionIndex.put(partition, index);
         }
 
-        int assigned = index.assign(hash, this::isMyBucket, maxBucketsNum);
+        int assigned = index.assign(hash, this::isMyBucket, maxBucketsNum, maxBucketId);
         if (LOG.isDebugEnabled()) {
             LOG.debug(
                     "Assign " + assigned + " to the partition " + partition + " key hash " + hash);
         }
+        if (assigned > maxBucketId) {
+            maxBucketId = assigned;
+        }
         return assigned;
     }
 

paimon-core/src/main/java/org/apache/paimon/index/PartitionIndex.java

@@ -64,7 +64,7 @@ public PartitionIndex(
         this.accessed = true;
     }
 
-    public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum) {
+    public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum, int maxBucketId) {
         accessed = true;
 
         // 1. is it a key that has appeared before
@@ -88,10 +88,6 @@ public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum) {
             }
         }
 
-        int maxBucketId =
-                totalBucket.isEmpty()
-                        ? 0
-                        : totalBucket.stream().mapToInt(Integer::intValue).max().getAsInt();
         if (-1 == maxBucketsNum || totalBucket.isEmpty() || maxBucketId < maxBucketsNum - 1) {
             // 3. create a new bucket
             for (int i = 0; i < Short.MAX_VALUE; i++) {

paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java

@@ -35,6 +35,7 @@ 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;
 
@@ -55,7 +56,11 @@ public int assign(BinaryRow partition, int hash) {
             index = new SimplePartitionIndex();
             this.partitionIndex.put(partition, index);
         }
-        return index.assign(hash);
+        int assigned = index.assign(hash);
+        if (assigned > maxBucketId) {
+            maxBucketId = assigned;
+        }
+        return assigned;
     }
 
     @Override
@@ -89,13 +94,6 @@ public int assign(int hash) {
             Long num = bucketInformation.computeIfAbsent(currentBucket, i -> 0L);
 
             if (num >= targetBucketRowNumber) {
-                int maxBucketId =
-                        bucketInformation.isEmpty()
-                                ? 0
-                                : bucketInformation.keySet().stream()
-                                        .mapToInt(Integer::intValue)
-                                        .max()
-                                        .getAsInt();
                 if (-1 == maxBucketsNum
                         || bucketInformation.isEmpty()
                         || maxBucketId < maxBucketsNum - 1) {