adaptive per-BE bucket assignment for random distribution load

sollhui · sollhui · commit adcf29ceb180 · 2026-04-21T11:55:45.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,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) {
diff --git a/be/src/exec/sink/vtablet_finder.h b/be/src/exec/sink/vtablet_finder.h
@@ -37,10 +37,21 @@ 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 };
-
-    OlapTabletFinder(VOlapTablePartitionParam* vpartition, FindTabletMode mode)
-            : _vpartition(vpartition), _find_tablet_mode(mode), _filter_bitmap(1024) {};
+    // 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,
+                     OlapTableLocationParam* location = nullptr)
+            : _vpartition(vpartition),
+              _find_tablet_mode(mode),
+              _location(location),
+              _filter_bitmap(1024) {};
 
     Status find_tablets(RuntimeState* state, Block* block, int rows,
                         std::vector<VOlapTablePartition*>& partitions,
@@ -64,11 +75,18 @@ 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;
+    // Reserved: was used for BE-side local-bucket lookup; FE now sends local_bucket_seqs directly
+    OlapTableLocationParam* _location = nullptr;
     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;
diff --git a/be/src/exec/sink/writer/vtablet_writer.cpp b/be/src/exec/sink/writer/vtablet_writer.cpp
@@ -1546,11 +1546,11 @@ 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));
-    _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode);
+    _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode, _location);
     RETURN_IF_ERROR(_vpartition->init());
 
     _state = state;
diff --git a/be/src/exec/sink/writer/vtablet_writer_v2.cpp b/be/src/exec/sink/writer/vtablet_writer_v2.cpp
@@ -166,11 +166,11 @@ 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));
-    _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode);
+    _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode, _location);
     RETURN_IF_ERROR(_vpartition->init());
 
     _state = state;
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) {
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-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
@@ -217,6 +217,7 @@ private static void setParamsForOlapTableSink(List<PipelineDistributedPlan> dist
             Map<DistributedPlanWorker, TPipelineFragmentParamsList> fragmentsGroupByWorker,
             CoordinatorContext coordinatorContext) {
         int numBackendsWithSink = 0;
+        int totalSinkInstances = 0;
         for (PipelineDistributedPlan distributedPlan : distributedPlans) {
             PlanFragment fragment = distributedPlan.getFragmentJob().getFragment();
             if (fragment.getSink() instanceof OlapTableSink) {
@@ -225,11 +226,13 @@ private static void setParamsForOlapTableSink(List<PipelineDistributedPlan> dist
                         .map(AssignedJob::getAssignedWorker)
                         .distinct()
                         .count();
+                totalSinkInstances += distributedPlan.getInstanceJobs().size();
             }
         }
 
         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,8 +245,130 @@ 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(),
+                            numBackendsWithSink, totalSinkInstances);
+                }
+            }
+        }
+    }
+
+    /**
+     * 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 on threshold.
+     *
+     * <p>The load_tablet_idx already stored in each TOlapTablePartition (set during
+     * OlapTableSink.createPartitionParam) is re-used as a rotation offset so that consecutive
+     * load tasks spread across different local buckets of the same BE.
+     *
+     * <p>The number of local buckets assigned per BE is capped by CalculateBucketNum:
+     * <pre>activeBuckets = min(numBackendsWithSink, totalBucketNum, totalSinkInstances)</pre>
+     * This prevents low-concurrency loads (e.g. stream load with a single instance) from
+     * scattering data across more buckets than necessary.
+     */
+    private static void assignRandomBucketPerBe(TPipelineFragmentParams fragmentParams, long beId,
+            int numBackendsWithSink, int totalSinkInstances) {
+        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);
+                    }
+                }
+            }
+
+            // CalculateBucketNum = min(numBackendsWithSink, totalBucketNum, totalSinkInstances)
+            // determines how many distinct buckets are written to across all sinks.
+            // Each BE gets at most ceil(activeBuckets / numBackendsWithSink) local buckets.
+            // Example: stream load (totalSinkInstances=1, numBackendsWithSink=1, numBuckets=6)
+            //   → activeBuckets = 1, maxPerBe = 1 → only 1 bucket, no switching
+            int activeBuckets = Math.min(numBackendsWithSink,
+                    Math.min(numBuckets, totalSinkInstances));
+            int maxLocalBucketsPerBe = Math.max(1,
+                    (activeBuckets + numBackendsWithSink - 1) / numBackendsWithSink);
+            if (localBuckets.size() > maxLocalBucketsPerBe) {
+                localBuckets = localBuckets.subList(0, maxLocalBucketsPerBe);
             }
+
+            int loadTabletIdx;
+            if (!localBuckets.isEmpty()) {
+                loadTabletIdx = localBuckets.get(tabletIndex % localBuckets.size());
+            } else {
+                // Fallback: BE has no local replica for this partition
+                loadTabletIdx = tabletIndex % numBuckets;
+            }
+
+            tPartition.setLoadTabletIdx(loadTabletIdx);
+            tPartition.setBucketBeId(beId);
+            tPartition.setLocalBucketSeqs(localBuckets);
         }
     }
 
diff --git a/gensrc/thrift/Descriptors.thrift b/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 {

Original file line number	Diff line number	Diff line change
`@@ -1546,11 +1546,11 @@ Status VTabletWriter::_init(RuntimeState* state, RuntimeProfile* profile) {`
`1546`	`1546`	`if (table_sink.__isset.load_to_single_tablet && table_sink.load_to_single_tablet) {`
`1547`	`1547`	`find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_SINK;`
`1548`	`1548`	`} else {`
`1549`		`- find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_BATCH;`
	`1549`	`+ find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_RANDOM_BUCKET;`
`1550`	`1550`	`}`
`1551`	`1551`	`}`
`1552`	`1552`	`_vpartition = _pool->add(new doris::VOlapTablePartitionParam(_schema, table_sink.partition));`
`1553`		`- _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode);`
	`1553`	`+ _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode, _location);`
`1554`	`1554`	`RETURN_IF_ERROR(_vpartition->init());`
`1555`	`1555`
`1556`	`1556`	`_state = state;`
Original file line number	Diff line number	Diff line change
`@@ -166,11 +166,11 @@ Status VTabletWriterV2::_init(RuntimeState* state, RuntimeProfile* profile) {`
`166`	`166`	`if (table_sink.__isset.load_to_single_tablet && table_sink.load_to_single_tablet) {`
`167`	`167`	`find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_SINK;`
`168`	`168`	`} else {`
`169`		`- find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_EVERY_BATCH;`
	`169`	`+ find_tablet_mode = OlapTabletFinder::FindTabletMode::FIND_TABLET_RANDOM_BUCKET;`
`170`	`170`	`}`
`171`	`171`	`}`
`172`	`172`	`_vpartition = _pool->add(new doris::VOlapTablePartitionParam(_schema, table_sink.partition));`
`173`		`- _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode);`
	`173`	`+ _tablet_finder = std::make_unique<OlapTabletFinder>(_vpartition, find_tablet_mode, _location);`
`174`	`174`	`RETURN_IF_ERROR(_vpartition->init());`
`175`	`175`
`176`	`176`	`_state = state;`