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

Author: sollhui

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,39 @@ 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 same as FIND_TABLET_EVERY_BATCH
+        _vpartition->find_tablets(block, qualified_rows, partitions, tablet_index,
+                                  &_partition_to_tablet_map);
+
+        int64_t bytes_per_row = rows > 0 ? block->bytes() / rows : 0;
+        for (auto& [partition, _] : _partition_to_tablet_map) {
+            if (partition->indexes.empty()) {
+                continue;
+            }
+            int64_t cur_bucket = partition->load_tablet_idx % partition->num_buckets;
+            int64_t cur_tablet_id = partition->indexes[0].tablets[cur_bucket];
+            _tablet_written_bytes[cur_tablet_id] += bytes_per_row * 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,
+            // which handles both single-replica and multi-replica correctly.
+            const auto& seqs = partition->local_bucket_seqs;
+            if (_tablet_written_bytes[cur_tablet_id] >= BUCKET_SWITCH_THRESHOLD_BYTES &&
+                !seqs.empty()) {
+                auto it = std::find(seqs.begin(), seqs.end(), static_cast<int32_t>(cur_bucket));
+                if (it != seqs.end()) {
+                    ++it;
+                    if (it == seqs.end()) {
+                        it = seqs.begin();
+                    }
+                    partition->load_tablet_idx = *it;
+                }
+            }
+        }
+        _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) {

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,16 @@ 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 -> bytes written so far; used to decide when to rotate to the next local bucket
+    std::map<int64_t, int64_t> _tablet_written_bytes;
 
     int64_t _num_filtered_rows = 0;
     int64_t _num_immutable_partition_filtered_rows = 0;

be/src/exec/sink/writer/vtablet_writer.cpp

@@ -1546,7 +1546,7 @@ Status VTabletWriter::_init(RuntimeState* state, RuntimeProfile* profile) {
         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;
+            find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_RANDOM_BUCKET;
         }
     }
     _vpartition = _pool->add(new doris::VOlapTablePartitionParam(_schema, table_sink.partition));

be/src/exec/sink/writer/vtablet_writer_v2.cpp

@@ -166,7 +166,7 @@ Status VTabletWriterV2::_init(RuntimeState* state, RuntimeProfile* profile) {
         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;
+            find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_RANDOM_BUCKET;
         }
     }
     _vpartition = _pool->add(new doris::VOlapTablePartitionParam(_schema, table_sink.partition));

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) {

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

fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java

@@ -230,6 +230,7 @@ private static void setParamsForOlapTableSink(List<PipelineDistributedPlan> dist
 
         ConnectContext connectContext = coordinatorContext.connectContext;
         for (Entry<DistributedPlanWorker, TPipelineFragmentParamsList> kv : fragmentsGroupByWorker.entrySet()) {
+            DistributedPlanWorker worker = kv.getKey();
             TPipelineFragmentParamsList fragments = kv.getValue();
             for (TPipelineFragmentParams fragmentParams : fragments.getParamsList()) {
                 if (fragmentParams.getFragment().getOutputSink().getType() == TDataSinkType.OLAP_TABLE_SINK) {
@@ -242,11 +243,116 @@ private static void setParamsForOlapTableSink(List<PipelineDistributedPlan> dist
                     fragmentParams.setNumLocalSink(fragmentParams.getLocalParams().size());
                     LOG.info("num local sink for backend {} is {}", fragmentParams.getBackendId(),
                             fragmentParams.getNumLocalSink());
+
+                    // Assign per-BE starting bucket and bucket_be_id for random distribution
+                    assignRandomBucketPerBe(fragmentParams, worker.id());
                 }
             }
         }
     }
 
+    /**
+     * For random-distribution partitions, override load_tablet_idx and set bucket_be_id /
+     * local_bucket_seqs so that each BE starts writing to one of its own local buckets and
+     * can rotate within them once per-tablet bytes exceed the threshold.
+     *
+     * <p>local_bucket_seqs always contains all local buckets for this BE.  The starting bucket
+     * rotates across consecutive load tasks via tabletIndex; within a single load, the BE
+     * switches to the next entry in the list only when the 200 MB threshold is reached.
+     * Small loads (e.g. stream load) naturally stay on one bucket because the threshold is
+     * never triggered.
+     */
+    private static void assignRandomBucketPerBe(TPipelineFragmentParams fragmentParams, long beId) {
+        org.apache.doris.thrift.TDataSink outputSink = fragmentParams.getFragment().getOutputSink();
+        org.apache.doris.thrift.TOlapTableSink olapSink = outputSink.getOlapTableSink();
+        if (olapSink == null || olapSink.getPartition() == null
+                || !olapSink.getPartition().isSetPartitions()) {
+            return;
+        }
+
+        // Check whether any partition uses random distribution (identified by load_tablet_idx being set)
+        boolean hasRandomPartition = olapSink.getPartition().getPartitions().stream()
+                .anyMatch(org.apache.doris.thrift.TOlapTablePartition::isSetLoadTabletIdx);
+        if (!hasRandomPartition) {
+            return;
+        }
+
+        // Build tablet_id -> BE IDs index from location
+        org.apache.doris.thrift.TOlapTableLocationParam location = olapSink.getLocation();
+        if (location == null || !location.isSetTablets()) {
+            return;
+        }
+        Map<Long, List<Long>> tabletToBeIds = new java.util.HashMap<>();
+        for (org.apache.doris.thrift.TTabletLocation tabletLoc : location.getTablets()) {
+            tabletToBeIds.put(tabletLoc.getTabletId(), tabletLoc.getNodeIds());
+        }
+
+        // Deep-copy only the OlapTableSink so that per-BE modifications are isolated
+        org.apache.doris.thrift.TOlapTableSink sinkCopy = olapSink.deepCopy();
+        outputSink.setOlapTableSink(sinkCopy);
+
+        for (org.apache.doris.thrift.TOlapTablePartition tPartition
+                : sinkCopy.getPartition().getPartitions()) {
+            if (!tPartition.isSetLoadTabletIdx()) {
+                continue; // hash distribution – leave as-is
+            }
+
+            int tabletIndex = (int) tPartition.getLoadTabletIdx();
+            List<org.apache.doris.thrift.TOlapTableIndexTablets> indexes = tPartition.getIndexes();
+            if (indexes == null || indexes.isEmpty()) {
+                continue;
+            }
+            List<Long> tablets = indexes.get(0).getTablets();
+            int numBuckets = tPartition.getNumBuckets();
+
+            // Collect bucket indices assigned to this BE.
+            //
+            // For single-replica tablets, each bucket has exactly one entry in node_ids,
+            // so the bucket naturally belongs to one BE.
+            //
+            // For multi-replica tablets, multiple BEs host the same bucket. To avoid all BEs
+            // writing to the same bucket(s), we assign each bucket to exactly one BE using a
+            // deterministic hash rule: sort node_ids, then give bucket[i] to the BE at position
+            // (i % numReplicas) in the sorted list. This spreads buckets evenly across replica BEs
+            // and is computed consistently on every BE without coordination.
+            List<Integer> localBuckets = new ArrayList<>();
+            for (int bucketIdx = 0; bucketIdx < tablets.size(); bucketIdx++) {
+                List<Long> beIds = tabletToBeIds.get(tablets.get(bucketIdx));
+                if (beIds == null || beIds.isEmpty()) {
+                    continue;
+                }
+                if (beIds.size() == 1) {
+                    // Single replica: the only BE is the owner
+                    if (beIds.get(0).equals(beId)) {
+                        localBuckets.add(bucketIdx);
+                    }
+                } else {
+                    // Multi-replica: assign via hash so each bucket has exactly one owner BE
+                    List<Long> sorted = beIds.stream()
+                            .sorted()
+                            .collect(java.util.stream.Collectors.toList());
+                    int pos = sorted.indexOf(beId);
+                    if (pos >= 0 && bucketIdx % sorted.size() == pos) {
+                        localBuckets.add(bucketIdx);
+                    }
+                }
+            }
+
+            if (!localBuckets.isEmpty()) {
+                int loadTabletIdx = localBuckets.get(tabletIndex % localBuckets.size());
+                tPartition.setLoadTabletIdx(loadTabletIdx);
+                tPartition.setBucketBeId(beId);
+                // Send the full local bucket list; BE rotates through it in order when the
+                // per-tablet write threshold is reached.
+                tPartition.setLocalBucketSeqs(localBuckets);
+            } else {
+                // Fallback: BE has no local replica for this partition
+                tPartition.setLoadTabletIdx(tabletIndex % numBuckets);
+                // Leave bucket_be_id / local_bucket_seqs unset → no switching
+            }
+        }
+    }
+
     private static Multiset<DistributedPlanWorker> computeInstanceNumPerWorker(
             List<PipelineDistributedPlan> distributedPlans) {
         Multiset<DistributedPlanWorker> workerCounter = LinkedHashMultiset.create();

gensrc/thrift/Descriptors.thrift

@@ -285,6 +285,11 @@ struct TOlapTablePartition {
     // tablet_id -> list of backend_ids that have version gaps (lastFailedVersion >= 0)
     // used by BE to exclude these backends from success counting in majority write
     14: optional map<i64, list<i64>> tablet_version_gap_backends
+    // FE-assigned BE ID for local bucket rotation in FIND_TABLET_RANDOM_BUCKET mode
+    15: optional i64 bucket_be_id
+    // FE-assigned local bucket indices for this BE; used in FIND_TABLET_RANDOM_BUCKET mode
+    // to rotate only within the buckets owned by this BE (handles both single and multi-replica)
+    16: optional list<i32> local_bucket_seqs
 }
 
 struct TOlapTablePartitionParam {