From ad75b9b1427eaa579fb3e22ef2eecc616ce02754 Mon Sep 17 00:00:00 2001
From: Kang <kxiao.tiger@gmail.com>
Date: Fri, 19 Apr 2024 13:04:21 +0800
Subject: [PATCH] [opt](auto bucket) add fe config autobucket_max_buckets
 (#33842)

---
 .../java/org/apache/doris/common/Config.java  |  8 ++++-
 .../doris/common/util/AutoBucketUtils.java    |  3 +-
 .../suites/autobucket/test_autobucket.groovy  | 33 ++++++++++++++++---
 3 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/fe/fe-common/src/main/java/org/apache/doris/common/Config.java b/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
index 8dff0e3f4d..3525d411b6 100644
--- a/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
+++ b/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
@@ -2341,12 +2341,18 @@ public class Config extends ConfigBase {
     })
     public static long analyze_record_limit = 20000;
 
-    @ConfField(description = {
+    @ConfField(mutable = true, description = {
             "Auto Buckets中最小的buckets数目",
             "min buckets of auto bucket"
     })
     public static int autobucket_min_buckets = 1;
 
+    @ConfField(mutable = true, description = {
+        "Auto Buckets中最大的buckets数目",
+        "max buckets of auto bucket"
+    })
+    public static int autobucket_max_buckets = 128;
+
     @ConfField(description = {"Arrow Flight Server中所有用户token的缓存上限，超过后LRU淘汰，默认值为512, "
             + "并强制限制小于 qe_max_connection/2, 避免`Reach limit of connections`, "
             + "因为arrow flight sql是无状态的协议，连接通常不会主动断开，"
diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/util/AutoBucketUtils.java b/fe/fe-core/src/main/java/org/apache/doris/common/util/AutoBucketUtils.java
index 294250fd21..55a07b89eb 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/common/util/AutoBucketUtils.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/common/util/AutoBucketUtils.java
@@ -20,6 +20,7 @@ package org.apache.doris.common.util;
 import org.apache.doris.catalog.DiskInfo;
 import org.apache.doris.catalog.DiskInfo.DiskState;
 import org.apache.doris.catalog.Env;
+import org.apache.doris.common.Config;
 import org.apache.doris.system.Backend;
 import org.apache.doris.system.SystemInfoService;
 
@@ -85,7 +86,7 @@ public class AutoBucketUtils {
     public static int getBucketsNum(long partitionSize) {
         int bucketsNumByPartitionSize = convertParitionSizeToBucketsNum(partitionSize);
         int bucketsNumByBE = getBucketsNumByBEDisks();
-        int bucketsNum = Math.min(128, Math.min(bucketsNumByPartitionSize, bucketsNumByBE));
+        int bucketsNum = Math.min(Config.autobucket_max_buckets, Math.min(bucketsNumByPartitionSize, bucketsNumByBE));
         int beNum = getBENum();
         logger.debug("AutoBucketsUtil: bucketsNumByPartitionSize {}, bucketsNumByBE {}, bucketsNum {}, beNum {}",
                 bucketsNumByPartitionSize, bucketsNumByBE, bucketsNum, beNum);
diff --git a/regression-test/suites/autobucket/test_autobucket.groovy b/regression-test/suites/autobucket/test_autobucket.groovy
index d3ba70d0df..615ede5675 100644
--- a/regression-test/suites/autobucket/test_autobucket.groovy
+++ b/regression-test/suites/autobucket/test_autobucket.groovy
@@ -40,7 +40,8 @@ suite("test_autobucket") {
 
     sql "drop table if exists autobucket_test"
 
-
+    // set min to 5
+    sql "ADMIN SET FRONTEND CONFIG ('autobucket_min_buckets' = '5')"
     sql "drop table if exists autobucket_test_min_buckets"
     result = sql """
         CREATE TABLE `autobucket_test_min_buckets` (
@@ -55,11 +56,35 @@ suite("test_autobucket") {
         )
         """
 
-    default_min_buckets = 1 // in Config.java
     result = sql "show partitions from autobucket_test_min_buckets"
     logger.info("${result}")
     // XXX: buckets at pos(8), next maybe impl by sql meta
-    assertEquals(Integer.valueOf(result.get(0).get(8)), default_min_buckets)
-
+    assertEquals(Integer.valueOf(result.get(0).get(8)), 5)
+    // set back to default
+    sql "ADMIN SET FRONTEND CONFIG ('autobucket_min_buckets' = '1')"
     sql "drop table if exists autobucket_test_min_buckets"
+
+    // set max to 4
+    sql "ADMIN SET FRONTEND CONFIG ('autobucket_max_buckets' = '4')"
+    sql "drop table if exists autobucket_test_max_buckets"
+    result = sql """
+        CREATE TABLE `autobucket_test_max_buckets` (
+          `user_id` largeint(40) NOT NULL
+        ) ENGINE=OLAP
+        DUPLICATE KEY(`user_id`)
+        COMMENT 'OLAP'
+        DISTRIBUTED BY HASH(`user_id`) BUCKETS AUTO
+        PROPERTIES (
+          "replication_allocation" = "tag.location.default: 1",
+          "estimate_partition_size" = "100000G"
+        )
+        """
+
+    result = sql "show partitions from autobucket_test_max_buckets"
+    logger.info("${result}")
+    // XXX: buckets at pos(8), next maybe impl by sql meta
+    assertEquals(Integer.valueOf(result.get(0).get(8)), 4)
+    // set back to default
+    sql "ADMIN SET FRONTEND CONFIG ('autobucket_max_buckets' = '128')"
+    sql "drop table if exists autobucket_test_max_buckets"
 }