adaptive random bucket load: per-BE bucket assignment with threshold-based rotation

sollhui · sollhui · commit 0e8fec709e17 · 2026-04-22T00:02:11.000+08:00
diff --git a/be/src/exec/sink/vtablet_finder.cpp b/be/src/exec/sink/vtablet_finder.cpp
@@ -22,6 +22,7 @@
 #include <gen_cpp/FrontendService_types.h>
 #include <glog/logging.h>
 
+#include <algorithm>
 #include <string>
 #include <utility>
 
@@ -82,8 +83,82 @@ Status OlapTabletFinder::find_tablets(RuntimeState* state, Block* block, int row
 
     if (_find_tablet_mode == FindTabletMode::FIND_TABLET_EVERY_ROW) {
         _vpartition->find_tablets(block, qualified_rows, partitions, tablet_index);
+    } else if (_find_tablet_mode == FindTabletMode::FIND_TABLET_RANDOM_BUCKET) {
+        // Use per-batch caching: all rows in the same partition go to the same tablet this batch.
+        _vpartition->find_tablets(block, qualified_rows, partitions, tablet_index,
+                                  &_partition_to_tablet_map);
+
+        // Count actual rows per partition so we don't credit each partition with the
+        // full block size when a block spans multiple partitions.
+        flat_hash_map<VOlapTablePartition*, int64_t> partition_row_counts;
+        for (uint32_t qi : qualified_rows) {
+            partition_row_counts[partitions[qi]]++;
+        }
+
+        int64_t bytes_per_row = rows > 0 ? block->bytes() / rows : 0;
+        for (auto& [partition, _] : _partition_to_tablet_map) {
+            if (partition->indexes.empty()) {
+                continue;
+            }
+            const auto& seqs = partition->local_bucket_seqs;
+
+            // Maintain _partition_seq_pos so rotation is O(1) after the first batch.
+            // On first encounter, find the current bucket's position in seqs (one-time O(n)).
+            // On subsequent batches we use the cached position directly.
+            int32_t& seq_pos = _partition_seq_pos[partition]; // default-initialised to 0
+            if (seq_pos == 0 && !seqs.empty()) {
+                // Sync seq_pos with the FE-assigned starting bucket. If load_tablet_idx happens
+                // to already be seqs[0] this is a no-op (seq_pos stays 0), otherwise we find
+                // the correct starting position once.
+                int32_t cur_bucket_val =
+                        static_cast<int32_t>(partition->load_tablet_idx % partition->num_buckets);
+                if (seqs[0] != cur_bucket_val) {
+                    auto it = std::find(seqs.begin(), seqs.end(), cur_bucket_val);
+                    if (it != seqs.end()) {
+                        seq_pos = static_cast<int32_t>(it - seqs.begin());
+                    }
+                }
+            }
+
+            // Determine the current bucket:
+            //   - adaptive mode: use the cached position in local_bucket_seqs (O(1))
+            //   - fallback (no seqs, or dynamically added partition without seqs): use
+            //     load_tablet_idx if valid, otherwise fall back to fast_rand() to match
+            //     the FIND_TABLET_EVERY_BATCH behaviour for partitions that were added
+            //     at runtime (e.g. auto-partition, overwrite replace) and did not go
+            //     through assignRandomBucketPerBe.
+            int64_t cur_bucket;
+            if (!seqs.empty()) {
+                cur_bucket = seqs[seq_pos];
+            } else if (partition->load_tablet_idx >= 0) {
+                cur_bucket = partition->load_tablet_idx % partition->num_buckets;
+            } else {
+                cur_bucket =
+                        static_cast<int64_t>(butil::fast_rand()) % partition->num_buckets;
+            }
+            int64_t cur_tablet_id = partition->indexes[0].tablets[cur_bucket];
+
+            // Credit only the rows that actually went to this partition.
+            int64_t part_rows = 0;
+            if (auto it = partition_row_counts.find(partition); it != partition_row_counts.end()) {
+                part_rows = it->second;
+            }
+            _tablet_written_bytes[cur_tablet_id] += bytes_per_row * part_rows;
+
+            // Switch to the next local bucket when the threshold is reached.
+            // local_bucket_seqs carries the exact set of buckets assigned to this BE by the FE.
+            if (!seqs.empty() &&
+                _tablet_written_bytes[cur_tablet_id] >= BUCKET_SWITCH_THRESHOLD_BYTES) {
+                seq_pos = (seq_pos + 1) % static_cast<int32_t>(seqs.size());
+                partition->load_tablet_idx = seqs[seq_pos];
+                VLOG(2) << "FIND_TABLET_RANDOM_BUCKET: partition=" << partition->id
+                        << " rotate bucket " << cur_bucket << " -> " << seqs[seq_pos]
+                        << " after " << _tablet_written_bytes[cur_tablet_id] << " bytes";
+            }
+        }
+        _partition_to_tablet_map.clear();
     } else {
-        // for random distribution
+        // FIND_TABLET_EVERY_BATCH / FIND_TABLET_EVERY_SINK
         _vpartition->find_tablets(block, qualified_rows, partitions, tablet_index,
                                   &_partition_to_tablet_map);
         if (_find_tablet_mode == FindTabletMode::FIND_TABLET_EVERY_BATCH) {
diff --git a/be/src/exec/sink/vtablet_finder.h b/be/src/exec/sink/vtablet_finder.h
@@ -37,7 +37,14 @@ class OlapTabletFinder {
     // FIND_TABLET_EVERY_SINK is used for random distribution info when load_to_single_tablet set to true,
     // which indicates that we should only compute tablet index in the corresponding partition once for the
     // whole time in olap table sink
-    enum FindTabletMode { FIND_TABLET_EVERY_ROW, FIND_TABLET_EVERY_BATCH, FIND_TABLET_EVERY_SINK };
+    // FIND_TABLET_RANDOM_BUCKET is used for random distribution info where FE assigns a starting
+    // bucket per BE; BE rotates within its local buckets once per-tablet bytes exceed a threshold
+    enum FindTabletMode {
+        FIND_TABLET_EVERY_ROW,
+        FIND_TABLET_EVERY_BATCH,
+        FIND_TABLET_EVERY_SINK,
+        FIND_TABLET_RANDOM_BUCKET
+    };
 
     OlapTabletFinder(VOlapTablePartitionParam* vpartition, FindTabletMode mode)
             : _vpartition(vpartition), _find_tablet_mode(mode), _filter_bitmap(1024) {};
@@ -64,11 +71,19 @@ class OlapTabletFinder {
 
     Bitmap& filter_bitmap() { return _filter_bitmap; }
 
+    // Threshold for switching to the next local bucket in FIND_TABLET_RANDOM_BUCKET mode.
+    static constexpr int64_t BUCKET_SWITCH_THRESHOLD_BYTES = 200LL * 1024 * 1024;
+
 private:
     VOlapTablePartitionParam* _vpartition = nullptr;
     FindTabletMode _find_tablet_mode;
     std::map<VOlapTablePartition*, int64_t> _partition_to_tablet_map;
     flat_hash_set<int64_t> _partition_ids;
+    // tablet_id -> cumulative bytes written; used to decide when to rotate to the next local bucket
+    flat_hash_map<int64_t, int64_t> _tablet_written_bytes;
+    // partition -> current index into local_bucket_seqs; initialized lazily on first batch.
+    // Storing the position (not the bucket value) makes rotation O(1) after the first lookup.
+    flat_hash_map<VOlapTablePartition*, int32_t> _partition_seq_pos;
 
     int64_t _num_filtered_rows = 0;
     int64_t _num_immutable_partition_filtered_rows = 0;
diff --git a/be/src/exec/sink/writer/vtablet_writer.cpp b/be/src/exec/sink/writer/vtablet_writer.cpp
@@ -1538,12 +1538,25 @@ Status VTabletWriter::_init(RuntimeState* state, RuntimeProfile* profile) {
     if (table_sink.__isset.send_batch_parallelism && table_sink.send_batch_parallelism > 1) {
         _send_batch_parallelism = table_sink.send_batch_parallelism;
     }
-    // if distributed column list is empty, we can ensure that tablet is with random distribution info
-    // and if load_to_single_tablet is set and set to true, we should find only one tablet in one partition
-    // for the whole olap table sink
+    // If distributed column list is empty, the table uses random distribution.
+    // Mode priority (highest to lowest):
+    //   1. FIND_TABLET_RANDOM_BUCKET: FE set local_bucket_seqs on at least one partition,
+    //      meaning enable_adaptive_random_bucket_load is ON. This overrides load_to_single_tablet
+    //      because the local-bucket strategy already constrains writes to a small set of tablets.
+    //   2. FIND_TABLET_EVERY_SINK: load_to_single_tablet=true (legacy single-tablet mode).
+    //   3. FIND_TABLET_EVERY_BATCH: default round-robin per batch.
     auto find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_ROW;
     if (table_sink.partition.distributed_columns.empty()) {
-        if (table_sink.__isset.load_to_single_tablet && table_sink.load_to_single_tablet) {
+        bool use_adaptive_random_bucket = false;
+        for (const auto& part : table_sink.partition.partitions) {
+            if (part.__isset.local_bucket_seqs && !part.local_bucket_seqs.empty()) {
+                use_adaptive_random_bucket = true;
+                break;
+            }
+        }
+        if (use_adaptive_random_bucket) {
+            find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_RANDOM_BUCKET;
+        } else if (table_sink.__isset.load_to_single_tablet && table_sink.load_to_single_tablet) {
             find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_SINK;
         } else {
             find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_BATCH;
diff --git a/be/src/exec/sink/writer/vtablet_writer_v2.cpp b/be/src/exec/sink/writer/vtablet_writer_v2.cpp
@@ -158,12 +158,25 @@ Status VTabletWriterV2::_init(RuntimeState* state, RuntimeProfile* profile) {
     _location = _pool->add(new OlapTableLocationParam(table_sink.location));
     _nodes_info = _pool->add(new DorisNodesInfo(table_sink.nodes_info));
 
-    // if distributed column list is empty, we can ensure that tablet is with random distribution info
-    // and if load_to_single_tablet is set and set to true, we should find only one tablet in one partition
-    // for the whole olap table sink
+    // If distributed column list is empty, the table uses random distribution.
+    // Mode priority (highest to lowest):
+    //   1. FIND_TABLET_RANDOM_BUCKET: FE set local_bucket_seqs on at least one partition,
+    //      meaning enable_adaptive_random_bucket_load is ON. This overrides load_to_single_tablet
+    //      because the local-bucket strategy already constrains writes to a small set of tablets.
+    //   2. FIND_TABLET_EVERY_SINK: load_to_single_tablet=true (legacy single-tablet mode).
+    //   3. FIND_TABLET_EVERY_BATCH: default round-robin per batch.
     auto find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_ROW;
     if (table_sink.partition.distributed_columns.empty()) {
-        if (table_sink.__isset.load_to_single_tablet && table_sink.load_to_single_tablet) {
+        bool use_adaptive_random_bucket = false;
+        for (const auto& part : table_sink.partition.partitions) {
+            if (part.__isset.local_bucket_seqs && !part.local_bucket_seqs.empty()) {
+                use_adaptive_random_bucket = true;
+                break;
+            }
+        }
+        if (use_adaptive_random_bucket) {
+            find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_RANDOM_BUCKET;
+        } else if (table_sink.__isset.load_to_single_tablet && table_sink.load_to_single_tablet) {
             find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_SINK;
         } else {
             find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_BATCH;
diff --git a/be/src/storage/tablet_info.cpp b/be/src/storage/tablet_info.cpp
@@ -700,10 +700,15 @@ Status VOlapTablePartitionParam::generate_partition_from(const TOlapTablePartiti
     part_result = _obj_pool.add(new VOlapTablePartition(&_partition_block));
     part_result->id = t_part.id;
     part_result->is_mutable = t_part.is_mutable;
-    // only load_to_single_tablet = true will set load_tablet_idx
     if (t_part.__isset.load_tablet_idx) {
         part_result->load_tablet_idx = t_part.load_tablet_idx;
     }
+    if (t_part.__isset.bucket_be_id) {
+        part_result->bucket_be_id = t_part.bucket_be_id;
+    }
+    if (t_part.__isset.local_bucket_seqs) {
+        part_result->local_bucket_seqs = t_part.local_bucket_seqs;
+    }
 
     if (_is_in_partition) {
         for (const auto& keys : t_part.in_keys) {
@@ -790,6 +795,15 @@ Status VOlapTablePartitionParam::add_partitions(
         }
 
         part->num_buckets = t_part.num_buckets;
+        if (t_part.__isset.load_tablet_idx) {
+            part->load_tablet_idx = t_part.load_tablet_idx;
+        }
+        if (t_part.__isset.bucket_be_id) {
+            part->bucket_be_id = t_part.bucket_be_id;
+        }
+        if (t_part.__isset.local_bucket_seqs) {
+            part->local_bucket_seqs = t_part.local_bucket_seqs;
+        }
         auto num_indexes = _schema->indexes().size();
         if (t_part.indexes.size() != num_indexes) {
             return Status::InternalError(
@@ -862,6 +876,15 @@ Status VOlapTablePartitionParam::replace_partitions(
         }
 
         part->num_buckets = t_part.num_buckets;
+        if (t_part.__isset.load_tablet_idx) {
+            part->load_tablet_idx = t_part.load_tablet_idx;
+        }
+        if (t_part.__isset.bucket_be_id) {
+            part->bucket_be_id = t_part.bucket_be_id;
+        }
+        if (t_part.__isset.local_bucket_seqs) {
+            part->local_bucket_seqs = t_part.local_bucket_seqs;
+        }
         auto num_indexes = _schema->indexes().size();
         if (t_part.indexes.size() != num_indexes) {
             return Status::InternalError(
diff --git a/be/src/storage/tablet_info.h b/be/src/storage/tablet_info.h
@@ -164,6 +164,11 @@ struct VOlapTablePartition {
     bool is_mutable;
     // -1 indicates partition with hash distribution
     int64_t load_tablet_idx = -1;
+    // BE ID assigned by FE for local bucket rotation; -1 means not set
+    int64_t bucket_be_id = -1;
+    // Bucket indices (0-based) that FE assigned to this BE for rotation.
+    // Empty means not set (fallback to full round-robin).
+    std::vector<int32_t> local_bucket_seqs;
     int total_replica_num = 0;
     int load_required_replica_num = 0;
     // tablet_id -> set of backend_ids that have version gaps
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
@@ -3610,4 +3610,15 @@ public static int metaServiceRpcRetryTimes() {
                     + "obtaining partition version information when calculating the delete bitmap. Enabled "
                     + "by default."})
     public static boolean calc_delete_bitmap_get_versions_waiting_for_pending_txns = true;
+
+    @ConfField(mutable = true, masterOnly = true, description = {
+            "Whether to enable adaptive random bucket load. When enabled, the FE assigns each BE only the "
+                    + "buckets whose primary replica it hosts (local_bucket_seqs), and the BE rotates across "
+                    + "those buckets once per-tablet write volume exceeds the threshold (default 200 MB). "
+                    + "This reduces import memory pressure and improves throughput for random-distribution "
+                    + "tables. When enabled, load_to_single_tablet is also subsumed by this strategy.",
+            "是否启用自适应随机桶导入。开启后 FE 为每个 BE 仅分配其持有主副本的桶（local_bucket_seqs），"
+                    + "BE 在单个 tablet 写入量超过阈值（默认 200 MB）后在本地桶之间轮转。"
+                    + "可降低导入内存压力并提升随机分桶表的吞吐量。开启后 load_to_single_tablet 也走新逻辑。"})
+    public static boolean enable_adaptive_random_bucket_load = true;
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java
diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/FrontendServiceImpl.java b/fe/fe-core/src/main/java/org/apache/doris/service/FrontendServiceImpl.java
diff --git a/gensrc/thrift/Descriptors.thrift b/gensrc/thrift/Descriptors.thrift
