test

Author: sollhui

be/src/exec/sink/vtablet_finder.cpp

@@ -80,7 +80,10 @@ void AdaptiveRandomBucketState::init_partition(VOlapTablePartition* part) {
 
     PartitionState state;
     state.partition_id = part->id;
-    state.tablets.assign(part->indexes[0].tablets.begin(), part->indexes[0].tablets.end());
+    state.index_tablets.reserve(part->indexes.size());
+    for (const auto& index : part->indexes) {
+        state.index_tablets.emplace_back(index.tablets.begin(), index.tablets.end());
+    }
     state.local_bucket_seqs = part->local_bucket_seqs;
 
     auto it = std::find(state.local_bucket_seqs.begin(), state.local_bucket_seqs.end(),
@@ -89,12 +92,13 @@ void AdaptiveRandomBucketState::init_partition(VOlapTablePartition* part) {
         state.seq_pos = static_cast<int32_t>(it - state.local_bucket_seqs.begin());
     }
     int32_t current_bucket = state.local_bucket_seqs[state.seq_pos];
-    state.current_tablet_id = state.tablets[current_bucket];
 
-    for (const auto& index : part->indexes) {
+    for (int32_t index_idx = 0; index_idx < part->indexes.size(); ++index_idx) {
+        const auto& index = part->indexes[index_idx];
         for (int32_t bucket = 0; bucket < index.tablets.size(); ++bucket) {
             _tablet_to_partition[index.tablets[bucket]] = part->id;
             _tablet_to_bucket[index.tablets[bucket]] = bucket;
+            _tablet_to_index[index.tablets[bucket]] = index_idx;
         }
     }
     _partition_states.emplace(part->id, std::move(state));
@@ -113,6 +117,25 @@ int64_t AdaptiveRandomBucketState::current_bucket(int64_t partition_id) {
     return it->second.local_bucket_seqs[it->second.seq_pos];
 }
 
+int64_t AdaptiveRandomBucketState::current_tablet_by_tablet(int64_t tablet_id) {
+    std::lock_guard<std::mutex> lock(_mutex);
+    auto partition_it = _tablet_to_partition.find(tablet_id);
+    if (partition_it == _tablet_to_partition.end()) {
+        return tablet_id;
+    }
+    auto state_it = _partition_states.find(partition_it->second);
+    if (state_it == _partition_states.end()) {
+        return tablet_id;
+    }
+    auto index_it = _tablet_to_index.find(tablet_id);
+    if (index_it == _tablet_to_index.end()) {
+        return tablet_id;
+    }
+    auto& state = state_it->second;
+    int32_t current_bucket = state.local_bucket_seqs[state.seq_pos];
+    return state.index_tablets[index_it->second][current_bucket];
+}
+
 void AdaptiveRandomBucketState::rotate_by_tablet(int64_t tablet_id) {
     if (!config::enable_adaptive_random_bucket_load_bucket_rotation) {
         return;
@@ -142,7 +165,6 @@ void AdaptiveRandomBucketState::rotate_by_tablet(int64_t tablet_id) {
               << " shared rotate bucket " << current_bucket << " -> " << next_bucket
               << " after tablet=" << tablet_id << " memtable flushed";
     state.seq_pos = next_pos;
-    state.current_tablet_id = state.tablets[next_bucket];
 }
 
 // Computes local_bucket_seqs for a random-distribution partition and updates

be/src/exec/sink/vtablet_finder.h

@@ -47,22 +47,23 @@ class AdaptiveRandomBucketState {
 
     void init_partition(VOlapTablePartition* part);
     int64_t current_bucket(int64_t partition_id);
+    int64_t current_tablet_by_tablet(int64_t tablet_id);
     void rotate_by_tablet(int64_t tablet_id);
 
 private:
     struct PartitionState {
         int64_t partition_id = -1;
-        std::vector<int64_t> tablets;
+        std::vector<std::vector<int64_t>> index_tablets;
         std::vector<int32_t> local_bucket_seqs;
         int32_t seq_pos = 0;
-        int64_t current_tablet_id = -1;
     };
 
     std::mutex _mutex;
     UniqueId _load_id;
     std::unordered_map<int64_t, PartitionState> _partition_states;
     std::unordered_map<int64_t, int64_t> _tablet_to_partition;
     std::unordered_map<int64_t, int32_t> _tablet_to_bucket;
+    std::unordered_map<int64_t, int32_t> _tablet_to_index;
 };
 
 class AdaptiveRandomBucketStatePool {
@@ -113,6 +114,13 @@ class OlapTabletFinder {
         }
     }
 
+    int64_t current_adaptive_random_bucket_tablet(int64_t tablet_id) {
+        if (_adaptive_random_bucket_state == nullptr) {
+            return tablet_id;
+        }
+        return _adaptive_random_bucket_state->current_tablet_by_tablet(tablet_id);
+    }
+
     Status find_tablets(RuntimeState* state, Block* block, int rows,
                         std::vector<VOlapTablePartition*>& partitions,
                         std::vector<uint32_t>& tablet_index, std::vector<bool>& skip,

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

@@ -816,6 +816,28 @@ Status VNodeChannel::add_block(Block* block, const Payload* payload) {
     return Status::OK();
 }
 
+void VNodeChannel::rewrite_tablet_ids_by_adaptive_random_bucket(
+        PTabletWriterAddBlockRequest* request) {
+    if (!_parent->_tablet_finder->is_adaptive_random_bucket() || request->eos()) {
+        return;
+    }
+    int64_t rewrite_count = 0;
+    for (int i = 0; i < request->tablet_ids_size(); ++i) {
+        int64_t old_tablet_id = request->tablet_ids(i);
+        int64_t current_tablet_id =
+                _parent->_tablet_finder->current_adaptive_random_bucket_tablet(old_tablet_id);
+        if (current_tablet_id != old_tablet_id) {
+            request->set_tablet_ids(i, current_tablet_id);
+            rewrite_count++;
+        }
+    }
+    if (rewrite_count > 0) {
+        LOG(INFO) << "FIND_TABLET_RANDOM_BUCKET: load_id=" << print_id(_parent->_load_id)
+                  << ", node_id=" << _node_id << " rewrite pending tablet ids, rows="
+                  << request->tablet_ids_size() << ", rewrite_count=" << rewrite_count;
+    }
+}
+
 static void injection_full_gc_fn() {
     MemoryReclamation::revoke_process_memory("injection_full_gc_fn");
 }
@@ -931,6 +953,7 @@ void VNodeChannel::try_send_pending_block(RuntimeState* state) {
     auto request = std::move(send_block.second); // doesn't need to be saved in heap
 
     // tablet_ids has already set when add row
+    rewrite_tablet_ids_by_adaptive_random_bucket(request.get());
     request->set_packet_seq(_next_packet_seq);
     auto block = mutable_block->to_block();
     CHECK(block.rows() == request->tablet_ids_size())

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

@@ -266,6 +266,8 @@ class VNodeChannel {
     // when there's pending block found by try_send_and_fetch_status(), we will awake a thread to send it.
     void try_send_pending_block(RuntimeState* state);
 
+    void rewrite_tablet_ids_by_adaptive_random_bucket(PTabletWriterAddBlockRequest* request);
+
     void clear_all_blocks();
 
     // two ways to stop channel: