refactor to support total buckets upper bound

Author: liyubin117

docs/layouts/shortcodes/generated/core_configuration.html

@@ -291,10 +291,10 @@
             <td>Initial buckets for a partition in assigner operator for dynamic bucket mode.</td>
         </tr>
         <tr>
-            <td><h5>dynamic-bucket.max-buckets-per-assigner</h5></td>
+            <td><h5>dynamic-bucket.max-buckets</h5></td>
             <td style="word-wrap: break-word;">-1</td>
             <td>Integer</td>
-            <td>Max buckets per assigner operator 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 greater than 0 (fixed upper bound).</td>
         </tr>
         <tr>
             <td><h5>dynamic-bucket.target-row-num</h5></td>

paimon-common/src/main/java/org/apache/paimon/CoreOptions.java

@@ -1076,12 +1076,12 @@ public class CoreOptions implements Serializable {
                     .withDescription(
                             "Initial buckets for a partition in assigner operator for dynamic bucket mode.");
 
-    public static final ConfigOption<Integer> DYNAMIC_BUCKET_MAX_BUCKETS_PER_ASSIGNER =
-            key("dynamic-bucket.max-buckets-per-assigner")
+    public static final ConfigOption<Integer> DYNAMIC_BUCKET_MAX_BUCKETS =
+            key("dynamic-bucket.max-buckets")
                     .intType()
                     .defaultValue(-1)
                     .withDescription(
-                            "Max buckets per assigner operator for a partition in dynamic bucket mode, It should "
+                            "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).");
 
     public static final ConfigOption<Integer> DYNAMIC_BUCKET_ASSIGNER_PARALLELISM =
@@ -2227,8 +2227,8 @@ public Integer dynamicBucketInitialBuckets() {
         return options.get(DYNAMIC_BUCKET_INITIAL_BUCKETS);
     }
 
-    public Integer dynamicBucketMaxBucketsPerAssigner() {
-        return options.get(DYNAMIC_BUCKET_MAX_BUCKETS_PER_ASSIGNER);
+    public Integer dynamicBucketMaxBuckets() {
+        return options.get(DYNAMIC_BUCKET_MAX_BUCKETS);
     }
 
     public Integer dynamicBucketAssignerParallelism() {

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

@@ -24,9 +24,13 @@
 import org.apache.paimon.utils.Int2ShortHashMap;
 import org.apache.paimon.utils.IntIterator;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.io.EOFException;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
@@ -39,6 +43,7 @@
 
 /** Bucket Index Per Partition. */
 public class PartitionIndex {
+    private static final Logger LOG = LoggerFactory.getLogger(PartitionIndex.class);
 
     public final Int2ShortHashMap hash2Bucket;
 
@@ -81,15 +86,9 @@ public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum) {
             Long number = entry.getValue();
             if (number < targetBucketRowNumber) {
                 entry.setValue(number + 1);
-                return cacheBucketAndGet(hash, bucket);
+                return cacheBucketAndGet(hash2Bucket, hash, bucket);
             } else {
                 iterator.remove();
-                if (-1 != maxBucketsNum && totalBucket.size() == maxBucketsNum) {
-                    return cacheBucketAndGet(
-                            hash,
-                            KeyAndBucketExtractor.bucketWithUpperBound(
-                                    totalBucket, hash, maxBucketsNum));
-                }
             }
         }
 
@@ -99,7 +98,7 @@ public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum) {
                 if (bucketFilter.test(i) && !totalBucket.contains(i)) {
                     nonFullBucketInformation.put(i, 1L);
                     totalBucket.add(i);
-                    return cacheBucketAndGet(hash, i);
+                    return cacheBucketAndGet(hash2Bucket, hash, i);
                 }
             }
 
@@ -110,7 +109,9 @@ public int assign(int hash, IntPredicate bucketFilter, int maxBucketsNum) {
                             "Too more bucket %s, you should increase target bucket row number %s.",
                             maxBucket, targetBucketRowNumber));
         } else {
+            // exceed buckets upper bound
             return cacheBucketAndGet(
+                    hash2Bucket,
                     hash,
                     KeyAndBucketExtractor.bucketWithUpperBound(totalBucket, hash, maxBucketsNum));
         }
@@ -149,8 +150,46 @@ public static PartitionIndex loadIndex(
         return new PartitionIndex(mapBuilder.build(), buckets, targetBucketRowNumber);
     }
 
-    private int cacheBucketAndGet(int hash, int bucket) {
+    public static int cacheBucketAndGet(Int2ShortHashMap hash2Bucket, int hash, int bucket) {
         hash2Bucket.put(hash, (short) bucket);
         return bucket;
     }
+
+    public static int[] getMaxBucketsPerAssigner(int maxBuckets, int assigners) {
+        int[] maxBucketsArr = new int[assigners];
+        if (-1 == maxBuckets) {
+            Arrays.fill(maxBucketsArr, -1);
+            return maxBucketsArr;
+        }
+        if (0 >= maxBuckets) {
+            throw new IllegalArgumentException(
+                    "Max-buckets should either be equal to -1 (unlimited), or it must be greater than 0 (fixed upper bound).");
+        }
+        int avg = maxBuckets / assigners;
+        int remainder = maxBuckets % assigners;
+        for (int i = 0; i < assigners; i++) {
+            maxBucketsArr[i] = avg;
+            if (remainder > 0) {
+                maxBucketsArr[i]++;
+                remainder--;
+            }
+        }
+        LOG.info(
+                "After distributing max-buckets {} to {} assigners evenly, maxBuckets layout: {}.",
+                maxBuckets,
+                assigners,
+                Arrays.toString(maxBucketsArr));
+        return maxBucketsArr;
+    }
+
+    public static int getSpecifiedMaxBuckets(int[] maxBucketsArr, int assignerId) {
+        int length = maxBucketsArr.length;
+        if (length == 0) {
+            throw new IllegalStateException("maxBuckets layout should exists!");
+        } else if (assignerId < length) {
+            return maxBucketsArr[assignerId];
+        } else {
+            return -1 == maxBucketsArr[0] ? -1 : 0;
+        }
+    }
 }

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

@@ -28,6 +28,8 @@
 import java.util.Map;
 import java.util.Set;
 
+import static org.apache.paimon.index.PartitionIndex.cacheBucketAndGet;
+
 /** When we need to overwrite the table, we should use this to avoid loading index. */
 public class SimpleHashBucketAssigner implements BucketAssigner {
 
@@ -89,11 +91,11 @@ public int assign(int hash) {
             Long num = bucketInformation.computeIfAbsent(currentBucket, i -> 0L);
             if (num >= targetBucketRowNumber) {
                 if (-1 != maxBucketsNum && bucketInformation.size() >= maxBucketsNum) {
-                    int bucket =
+                    return cacheBucketAndGet(
+                            hash2Bucket,
+                            hash,
                             KeyAndBucketExtractor.bucketWithUpperBound(
-                                    bucketInformation.keySet(), hash, maxBucketsNum);
-                    hash2Bucket.put(hash, (short) bucket);
-                    return bucket;
+                                    bucketInformation.keySet(), hash, maxBucketsNum));
                 } else {
                     loadNewBucket();
                 }

paimon-core/src/main/java/org/apache/paimon/table/sink/KeyAndBucketExtractor.java

@@ -27,7 +27,7 @@
 import java.util.Set;
 import java.util.concurrent.ThreadLocalRandom;
 
-import static org.apache.paimon.CoreOptions.DYNAMIC_BUCKET_MAX_BUCKETS_PER_ASSIGNER;
+import static org.apache.paimon.CoreOptions.DYNAMIC_BUCKET_MAX_BUCKETS;
 import static org.apache.paimon.utils.Preconditions.checkArgument;
 
 /**
@@ -66,7 +66,7 @@ static int bucketWithUpperBound(Set<Integer> bucketsSet, int hashcode, int maxBu
                 "Assign record (hashcode '{}') to new bucket exceed upper bound '{}' defined in '{}', Stop creating new buckets.",
                 hashcode,
                 maxBucketsNum,
-                DYNAMIC_BUCKET_MAX_BUCKETS_PER_ASSIGNER.key());
+                DYNAMIC_BUCKET_MAX_BUCKETS.key());
         return bucketsSet.stream()
                 .skip(ThreadLocalRandom.current().nextInt(maxBucketsNum))
                 .findFirst()

paimon-core/src/test/java/org/apache/paimon/index/HashBucketAssignerTest.java

@@ -37,6 +37,8 @@
 import java.util.Arrays;
 import java.util.Collections;
 
+import static org.apache.paimon.index.PartitionIndex.getMaxBucketsPerAssigner;
+import static org.apache.paimon.index.PartitionIndex.getSpecifiedMaxBuckets;
 import static org.apache.paimon.io.DataFileTestUtils.row;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -136,7 +138,7 @@ public void testAssignWithUpperBound() {
         assertThatThrownBy(() -> assigner.assign(row(1), 1))
                 .hasMessageContaining("This is a bug, record assign id");
 
-        // exceed upper bound
+        // exceed buckets upper bound
         // partition 1
         int hash = 18;
         for (int i = 0; i < 200; i++) {
@@ -151,7 +153,78 @@ public void testAssignWithUpperBound() {
         }
     }
 
-    @ParameterizedTest(name = "maxBucket: {0}")
+    @Test
+    public void testMultiAssigners() {
+        int[] maxBucketsArr = getMaxBucketsPerAssigner(4, 2);
+        Assertions.assertThat(maxBucketsArr).isEqualTo(new int[] {2, 2});
+
+        maxBucketsArr = getMaxBucketsPerAssigner(8, 3);
+        Assertions.assertThat(maxBucketsArr).isEqualTo(new int[] {3, 3, 2});
+
+        maxBucketsArr = getMaxBucketsPerAssigner(3, 2);
+        Assertions.assertThat(maxBucketsArr).isEqualTo(new int[] {2, 1});
+
+        Assertions.assertThat(getSpecifiedMaxBuckets(maxBucketsArr, 0)).isEqualTo(2);
+        Assertions.assertThat(getSpecifiedMaxBuckets(maxBucketsArr, 1)).isEqualTo(1);
+        Assertions.assertThat(getSpecifiedMaxBuckets(maxBucketsArr, 2)).isEqualTo(0);
+
+        maxBucketsArr = getMaxBucketsPerAssigner(-1, 2);
+        Assertions.assertThat(maxBucketsArr).isEqualTo(new int[] {-1, -1});
+
+        Assertions.assertThat(getSpecifiedMaxBuckets(maxBucketsArr, 0)).isEqualTo(-1);
+        Assertions.assertThat(getSpecifiedMaxBuckets(maxBucketsArr, 1)).isEqualTo(-1);
+        Assertions.assertThat(getSpecifiedMaxBuckets(maxBucketsArr, 2)).isEqualTo(-1);
+
+        assertThatThrownBy(() -> getMaxBucketsPerAssigner(-10, 2))
+                .hasMessageContaining(
+                        "Max-buckets should either be equal to -1 (unlimited), or it must be greater than 0 (fixed upper bound).");
+    }
+
+    @Test
+    public void testAssignWithUpperBoundMultiAssigners() {
+        int[] maxBucketsArr = getMaxBucketsPerAssigner(3, 2);
+        Assertions.assertThat(maxBucketsArr).isEqualTo(new int[] {2, 1});
+
+        HashBucketAssigner assigner0 = createAssigner(2, 2, 0, maxBucketsArr[0]);
+        HashBucketAssigner assigner1 = createAssigner(2, 2, 1, maxBucketsArr[1]);
+
+        // assigner0: assign
+        assertThat(assigner0.assign(row(1), 0)).isEqualTo(0);
+        assertThat(assigner0.assign(row(1), 2)).isEqualTo(0);
+        assertThat(assigner0.assign(row(1), 4)).isEqualTo(0);
+        assertThat(assigner0.assign(row(1), 6)).isEqualTo(0);
+        assertThat(assigner0.assign(row(1), 8)).isEqualTo(0);
+
+        // assigner0: full
+        assertThat(assigner0.assign(row(1), 10)).isEqualTo(2);
+        assertThat(assigner0.assign(row(1), 12)).isEqualTo(2);
+        assertThat(assigner0.assign(row(1), 14)).isEqualTo(2);
+        assertThat(assigner0.assign(row(1), 16)).isEqualTo(2);
+        assertThat(assigner0.assign(row(1), 18)).isEqualTo(2);
+
+        // assigner0: exceed buckets upper bound
+        int hash = 18;
+        for (int i = 0; i < 200; i++) {
+            int bucket = assigner0.assign(row(2), hash += 2);
+            Assertions.assertThat(bucket).isIn(0, 2);
+        }
+
+        // assigner1: assign
+        assertThat(assigner1.assign(row(1), 1)).isEqualTo(1);
+        assertThat(assigner1.assign(row(1), 3)).isEqualTo(1);
+        assertThat(assigner1.assign(row(1), 5)).isEqualTo(1);
+        assertThat(assigner1.assign(row(1), 7)).isEqualTo(1);
+        assertThat(assigner1.assign(row(1), 9)).isEqualTo(1);
+
+        // assigner1: exceed buckets upper bound
+        hash = 9;
+        for (int i = 0; i < 200; i++) {
+            int bucket = assigner1.assign(row(2), hash += 2);
+            Assertions.assertThat(bucket).isIn(1);
+        }
+    }
+
+    @ParameterizedTest(name = "maxBuckets: {0}")
     @ValueSource(ints = {-1, 1, 2})
     public void testPartitionCopy(int maxBucketsNum) {
         HashBucketAssigner assigner = createAssigner(1, 1, 0, maxBucketsNum);
@@ -229,7 +302,7 @@ public void testAssignRestoreWithUpperBound() {
         assertThat(assigner0.assign(row(1), 11)).isEqualTo(0);
         assertThat(assigner0.assign(row(1), 14)).isEqualTo(0);
         assertThat(assigner2.assign(row(1), 16)).isEqualTo(2);
-        // exceed buckets limits
+        // exceed buckets upper bound
         assertThat(assigner0.assign(row(1), 17)).isEqualTo(0);
     }
 

paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java

@@ -25,6 +25,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import static org.apache.paimon.index.PartitionIndex.getMaxBucketsPerAssigner;
 import static org.apache.paimon.io.DataFileTestUtils.row;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -83,7 +84,46 @@ public void testAssignWithUpperBound() {
         }
     }
 
-    @ParameterizedTest(name = "maxBucket: {0}")
+    @Test
+    public void testAssignWithUpperBoundMultiAssigners() {
+        int[] maxBucketsArr = getMaxBucketsPerAssigner(3, 2);
+        Assertions.assertThat(maxBucketsArr).isEqualTo(new int[] {2, 1});
+
+        SimpleHashBucketAssigner simpleHashBucketAssigner0 =
+                new SimpleHashBucketAssigner(2, 0, 100, maxBucketsArr[0]);
+        SimpleHashBucketAssigner simpleHashBucketAssigner1 =
+                new SimpleHashBucketAssigner(2, 1, 100, maxBucketsArr[1]);
+
+        BinaryRow binaryRow = BinaryRow.EMPTY_ROW;
+        int hash = 0;
+
+        for (int i = 0; i < 100; i++) {
+            int bucket = simpleHashBucketAssigner0.assign(binaryRow, hash++);
+            Assertions.assertThat(bucket).isEqualTo(0);
+        }
+
+        for (int i = 0; i < 100; i++) {
+            int bucket = simpleHashBucketAssigner1.assign(binaryRow, hash++);
+            Assertions.assertThat(bucket).isEqualTo(1);
+        }
+
+        for (int i = 0; i < 100; i++) {
+            int bucket = simpleHashBucketAssigner0.assign(binaryRow, hash++);
+            Assertions.assertThat(bucket).isEqualTo(2);
+        }
+
+        // exceed upper bound
+        for (int i = 0; i < 200; i++) {
+            int bucket = simpleHashBucketAssigner0.assign(binaryRow, hash++);
+            Assertions.assertThat(bucket).isIn(0, 2);
+        }
+        for (int i = 0; i < 200; i++) {
+            int bucket = simpleHashBucketAssigner1.assign(binaryRow, hash++);
+            Assertions.assertThat(bucket).isIn(1);
+        }
+    }
+
+    @ParameterizedTest(name = "maxBuckets: {0}")
     @ValueSource(ints = {-1, 1, 2})
     public void testAssignWithSameHash(int maxBucketsNum) {
         SimpleHashBucketAssigner simpleHashBucketAssigner =
@@ -105,7 +145,7 @@ public void testAssignWithSameHash(int maxBucketsNum) {
         }
     }
 
-    @ParameterizedTest(name = "maxBucket: {0}")
+    @ParameterizedTest(name = "maxBuckets: {0}")
     @ValueSource(ints = {-1, 1, 2})
     public void testPartitionCopy(int maxBucketsNum) {
         SimpleHashBucketAssigner assigner = new SimpleHashBucketAssigner(1, 0, 5, maxBucketsNum);