[server] Limit max bucket num per table partition

fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java

@@ -356,7 +356,7 @@ CompletableFuture<List<PartitionInfo>> listPartitionInfos(
      *   <li>{@link TooManyPartitionsException} if the number of partitions is larger than the
      *       maximum number of partitions of one table, see {@link ConfigOptions#MAX_PARTITION_NUM}.
      *   <li>{@link TooManyBucketsException} if the number of buckets is larger than the maximum
-     *       number of buckets of one table, see {@link ConfigOptions#MAX_BUCKET_NUM}.
+     *       number of buckets of one partition, see {@link ConfigOptions#MAX_BUCKET_NUM}.
      * </ul>
      *
      * @param tablePath The table path of the table.

fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java

@@ -1835,12 +1835,8 @@ private List<FsPathAndFileName> toFsPathAndFileNames(KvSnapshotHandle kvSnapshot
                 .collect(Collectors.toList());
     }
 
-    /**
-     * Test that creating a partitioned table with bucket count exceeding the maximum throws
-     * TooManyBucketsException.
-     */
     @Test
-    public void testAddTooManyBuckets() throws Exception {
+    public void testBucketLimitForPartitionedTableAppliesPerPartition() throws Exception {
         // Already set low maximum bucket limit to 30 for this test in ClientToServerITCaseBase
         String dbName = DEFAULT_TABLE_PATH.getDatabaseName();
         TableDescriptor partitionedTable =
@@ -1857,57 +1853,89 @@ public void testAddTooManyBuckets() throws Exception {
         TablePath tablePath = TablePath.of(dbName, "test_add_too_many_buckets_table");
         admin.createTable(tablePath, partitionedTable, true).get();
 
-        // Add 3 partitions (3 * 10 = 30 buckets, which is the limit)
-        for (int i = 0; i < 3; i++) {
+        // Add 4 partitions. The total bucket count is 40, which is above the configured limit,
+        // but each partition only has 10 buckets.
+        for (int i = 0; i < 4; i++) {
             admin.createPartition(tablePath, newPartitionSpec("age", String.valueOf(i)), false)
                     .get();
         }
+        assertThat(admin.listPartitionInfos(tablePath).get()).hasSize(4);
+
+        TableDescriptor tooManyBucketsPartitionedTable =
+                TableDescriptor.builder()
+                        .schema(
+                                Schema.newBuilder()
+                                        .column("id", DataTypes.STRING())
+                                        .column("name", DataTypes.STRING())
+                                        .column("age", DataTypes.STRING())
+                                        .build())
+                        .distributedBy(40, "id")
+                        .partitionedBy("age")
+                        .build();
+        TablePath tooManyBucketsTablePath =
+                TablePath.of(dbName, "test_too_many_buckets_partitioned");
 
-        // Try to add one more partition, exceeding the bucket limit (4 * 10 > 30)
         assertThatThrownBy(
                         () ->
-                                admin.createPartition(
-                                                tablePath, newPartitionSpec("age", "4"), false)
+                                admin.createTable(
+                                                tooManyBucketsTablePath,
+                                                tooManyBucketsPartitionedTable,
+                                                false)
                                         .get())
                 .cause()
                 .isInstanceOf(TooManyBucketsException.class)
-                .hasMessageContaining("exceeding the maximum of 30 buckets");
+                .hasMessageContaining(
+                        "Bucket count 40 exceeds the maximum limit 30 for a non-partitioned table or partition.");
     }
 
-    /**
-     * Test that creating a non-partitioned table with bucket count exceeding the maximum throws
-     * TooManyBucketsException.
-     */
     @Test
     public void testBucketLimitForNonPartitionedTable() throws Exception {
-        // Set a low maximum bucket limit for this test
-        // (Assuming the configuration is already set to 30 in ClientToServerITCaseBase)
         String dbName = DEFAULT_TABLE_PATH.getDatabaseName();
 
-        // Create a non-partitioned table with 40 buckets (exceeding limit of 30)
-        TableDescriptor nonPartitionedTable =
+        TableDescriptor maxBucketsNonPartitionedTable =
                 TableDescriptor.builder()
                         .schema(
                                 Schema.newBuilder()
                                         .column("id", DataTypes.STRING())
                                         .column("name", DataTypes.STRING())
                                         .column("value", DataTypes.STRING())
                                         .build())
-                        .distributedBy(40, "id") // 40 buckets exceeds the limit of 30
-                        .build(); // No partitionedBy call makes this non-partitioned
+                        .distributedBy(30, "id")
+                        .build();
+        TablePath maxBucketsTablePath = TablePath.of(dbName, "test_max_buckets_non_partitioned");
 
-        TablePath tablePath = TablePath.of(dbName, "test_too_many_buckets_non_partitioned");
+        admin.createTable(maxBucketsTablePath, maxBucketsNonPartitionedTable, false).get();
 
-        // Creating this table should throw TooManyBucketsException
-        assertThatThrownBy(() -> admin.createTable(tablePath, nonPartitionedTable, false).get())
+        TableDescriptor tooManyBucketsNonPartitionedTable =
+                TableDescriptor.builder()
+                        .schema(
+                                Schema.newBuilder()
+                                        .column("id", DataTypes.STRING())
+                                        .column("name", DataTypes.STRING())
+                                        .column("value", DataTypes.STRING())
+                                        .build())
+                        .distributedBy(31, "id")
+                        .build();
+        TablePath tooManyBucketsTablePath =
+                TablePath.of(dbName, "test_too_many_buckets_non_partitioned");
+
+        assertThatThrownBy(
+                        () ->
+                                admin.createTable(
+                                                tooManyBucketsTablePath,
+                                                tooManyBucketsNonPartitionedTable,
+                                                false)
+                                        .get())
                 .cause()
                 .isInstanceOf(TooManyBucketsException.class)
-                .hasMessageContaining("exceeds the maximum limit");
+                .hasMessageContaining(
+                        "Bucket count 31 exceeds the maximum limit 30 for a non-partitioned table or partition.");
     }
 
     @Test
-    public void testDefaultBucketLimitForNonPartitionedTable() throws Exception {
-        assertThat(ConfigOptions.MAX_BUCKET_NUM.defaultValue()).isEqualTo(20000);
+    public void testDefaultPartitionAndBucketLimitsForNonPartitionedTable() throws Exception {
+        assertThat(ConfigOptions.MAX_PARTITION_NUM.defaultValue()).isEqualTo(1000);
+        assertThat(ConfigOptions.MAX_BUCKET_NUM.defaultValue()).isEqualTo(4096);
 
         FlussClusterExtension defaultLimitCluster =
                 FlussClusterExtension.builder().setNumOfTabletServers(1).build();
@@ -1925,7 +1953,7 @@ public void testDefaultBucketLimitForNonPartitionedTable() throws Exception {
                                             .column("id", DataTypes.STRING())
                                             .column("name", DataTypes.STRING())
                                             .build())
-                            .distributedBy(30000, "id")
+                            .distributedBy(4097, "id")
                             .build();
 
             defaultLimitAdmin.createDatabase(dbName, DatabaseDescriptor.EMPTY, false).get();
@@ -1936,7 +1964,8 @@ public void testDefaultBucketLimitForNonPartitionedTable() throws Exception {
                                             .get())
                     .cause()
                     .isInstanceOf(TooManyBucketsException.class)
-                    .hasMessageContaining("Bucket count 30000 exceeds the maximum limit 20000.");
+                    .hasMessageContaining(
+                            "Bucket count 4097 exceeds the maximum limit 4096 for a non-partitioned table or partition.");
         } finally {
             defaultLimitCluster.close();
         }

fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java

@@ -290,11 +290,12 @@ public class ConfigOptions {
     public static final ConfigOption<Integer> MAX_BUCKET_NUM =
             key("max.bucket.num")
                     .intType()
-                    .defaultValue(20000)
+                    .defaultValue(4096)
                     .withDescription(
-                            "The maximum number of buckets that can be created for a table. "
-                                    + "The default value is 20000. "
-                                    + "This default is capped to reduce the risk that the table assignment znode exceeds "
+                            "The maximum number of buckets that can be created for a non-partitioned table "
+                                    + "or for each partition of a partitioned table. "
+                                    + "The default value is 4096. "
+                                    + "This default is capped to reduce the risk that an assignment znode exceeds "
                                     + "ZooKeeper's packet size limit.");
 
     /**

fluss-common/src/main/java/org/apache/fluss/exception/TooManyBucketsException.java

@@ -20,7 +20,8 @@
 import org.apache.fluss.annotation.PublicEvolving;
 
 /**
- * This exception is thrown if the number of table buckets is exceed max.bucket.num.
+ * This exception is thrown if the number of buckets for a non-partitioned table or partition
+ * exceeds max.bucket.num.
  *
  * @since 0.7
  */

fluss-server/src/main/java/org/apache/fluss/server/coordinator/AutoPartitionManager.java

@@ -441,7 +441,7 @@ private void createPartitions(
             } catch (TooManyBucketsException t) {
                 LOG.warn(
                         "Auto partitioning skip to create partition {} for table [{}], "
-                                + "because exceed the maximum number of buckets.",
+                                + "because exceed the maximum number of buckets per partition.",
                         partition,
                         tablePath);
             } catch (Exception e) {

fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java

@@ -828,9 +828,8 @@ public void createPartition(
                             partition.getPartitionQualifiedName(), tablePath));
         }
 
-        final int partitionNumber;
         try {
-            partitionNumber = zookeeperClient.getPartitionNumber(tablePath);
+            int partitionNumber = zookeeperClient.getPartitionNumber(tablePath);
             if (partitionNumber + 1 > maxPartitionNum) {
                 throw new TooManyPartitionsException(
                         String.format(
@@ -847,24 +846,12 @@ public void createPartition(
                     e);
         }
 
-        try {
-            int bucketCount = partitionAssignment.getBucketAssignments().size();
-            // currently, every partition has the same bucket count
-            int totalBuckets = bucketCount * (partitionNumber + 1);
-            if (totalBuckets > maxBucketNum) {
-                throw new TooManyBucketsException(
-                        String.format(
-                                "Adding partition '%s' would result in %d total buckets for table %s, exceeding the maximum of %d buckets.",
-                                partition.getPartitionName(),
-                                totalBuckets,
-                                tablePath,
-                                maxBucketNum));
-            }
-        } catch (TooManyBucketsException e) {
-            throw e;
-        } catch (Exception e) {
-            throw new FlussRuntimeException(
-                    String.format("Failed to check total bucket count for table %s", tablePath), e);
+        int bucketCount = partitionAssignment.getBucketAssignments().size();
+        if (bucketCount > maxBucketNum) {
+            throw new TooManyBucketsException(
+                    String.format(
+                            "Partition '%s' has %d buckets for table %s, exceeding the maximum of %d buckets per partition.",
+                            partition.getPartitionName(), bucketCount, tablePath, maxBucketNum));
         }
 
         try {

fluss-server/src/main/java/org/apache/fluss/server/utils/TableDescriptorValidation.java

@@ -253,7 +253,7 @@ private static void checkDistribution(TableDescriptor tableDescriptor, int maxBu
         if (bucketCount > maxBucketNum) {
             throw new TooManyBucketsException(
                     String.format(
-                            "Bucket count %s exceeds the maximum limit %s.",
+                            "Bucket count %s exceeds the maximum limit %s for a non-partitioned table or partition.",
                             bucketCount, maxBucketNum));
         }
         List<String> bucketKeys = tableDescriptor.getTableDistribution().get().getBucketKeys();

fluss-server/src/test/java/org/apache/fluss/server/coordinator/AutoPartitionManagerTest.java

@@ -582,14 +582,14 @@ void testDayPartitionDropShouldNotBeDelayedByJitter() throws Exception {
     }
 
     /**
-     * Test if AutoPartionManager.createPartition adheres to maxBucketLimit while adding new
-     * parition automatically, skip if it breaches limit.
+     * Test if AutoPartitionManager.createPartition applies maxBucketLimit per partition while
+     * adding new partition automatically.
      */
     @Test
-    void testMaxBucketNum() throws Exception {
+    void testMaxBucketNumPerPartition() throws Exception {
 
         int bucketCountPerPartition = 10;
-        int maxBucketNum = 30; // Allow only 3 partitions with 10 buckets each
+        int maxBucketNum = 30;
 
         Configuration config = new Configuration();
         config.set(ConfigOptions.MAX_BUCKET_NUM, maxBucketNum);
@@ -626,16 +626,15 @@ void testMaxBucketNum() throws Exception {
         periodicExecutor.triggerNonPeriodicScheduledTask();
 
         int partitionsNum = zookeeperClient.getPartitionNumber(tablePath);
-        // Only 3 partitions should be created (3 * 10 = 30 buckets) out of the 4 requested
-        assertThat(partitionsNum).isEqualTo(3);
+        // All 4 requested partitions should be created because each partition is below the limit.
+        assertThat(partitionsNum).isEqualTo(4);
 
         // Advance time to trigger another auto-partition cycle
         clock.advanceTime(Duration.ofDays(1).plusHours(23));
         periodicExecutor.triggerPeriodicScheduledTasks();
 
-        // Check partitions again - should still have only 3 due to bucket limit
         partitionsNum = zookeeperClient.getPartitionNumber(tablePath);
-        assertThat(partitionsNum).isEqualTo(3);
+        assertThat(partitionsNum).isEqualTo(5);
     }
 
     @Test