[BugFix] Fix the problem of column mode partial update in cross cluster replication (backport #40692) (#40933)

mergify[bot] · xiangguangyxg · web-flow · commit a40e2f872d81 · 2024-02-07T16:13:50.000+08:00
Co-authored-by: xiangguangyxg &lt;110401425+xiangguangyxg@users.noreply.github.com&gt;
diff --git a/be/src/agent/publish_version.cpp b/be/src/agent/publish_version.cpp
@@ -124,7 +124,7 @@ void run_publish_version_task(ThreadPoolToken* token, const TPublishVersionReque
                 tablet_span->SetAttribute("txn_id", transaction_id);
                 tablet_span->SetAttribute("tablet_id", task.tablet_id);
                 tablet_span->SetAttribute("version", task.version);
-                if (!task.rowset) {
+                if (!is_replication_txn && !task.rowset) {
                     task.st = Status::NotFound(
                             fmt::format("rowset not found  of tablet: {}, txn_id: {}", task.tablet_id, task.txn_id));
                     LOG(WARNING) << task.st;
diff --git a/be/src/storage/delta_column_group.h b/be/src/storage/delta_column_group.h
@@ -64,6 +64,8 @@ class DeltaColumnGroup {
         }
         return column_files;
     }
+
+    std::vector<std::vector<uint32_t>>& column_ids() { return _column_ids; }
     const std::vector<std::vector<uint32_t>>& column_ids() const { return _column_ids; }
     int64_t version() const { return _version; }
 
diff --git a/be/src/storage/lake/replication_txn_manager.cpp b/be/src/storage/lake/replication_txn_manager.cpp
@@ -378,6 +378,10 @@ Status ReplicationTxnManager::replicate_remote_snapshot(const TReplicateSnapshot
 Status ReplicationTxnManager::convert_rowset_meta(const RowsetMeta& rowset_meta, TTransactionId transaction_id,
                                                   TxnLogPB::OpWrite* op_write,
                                                   std::unordered_map<std::string, std::string>* filename_map) {
+    if (rowset_meta.is_column_mode_partial_update()) {
+        return Status::NotSupported("Column mode partial update is not supported in shared-data mode");
+    }
+
     // Convert rowset metadata
     auto* rowset_metadata = op_write->mutable_rowset();
     rowset_metadata->set_id(rowset_meta.get_rowset_seg_id());
diff --git a/be/src/storage/replication_txn_manager.cpp b/be/src/storage/replication_txn_manager.cpp
@@ -519,7 +519,9 @@ Status ReplicationTxnManager::replicate_remote_snapshot(const TReplicateSnapshot
         WritableFileOptions opts{.sync_on_close = true, .mode = FileSystem::CREATE_OR_OPEN_WITH_TRUNCATE};
         ASSIGN_OR_RETURN(auto output_file, fs::new_writable_file(opts, tablet_snapshot_dir_path + file_name));
 
-        if (HasSuffixString(file_name, ".dat") && !column_unique_id_map.empty()) {
+        if (!column_unique_id_map.empty() &&
+            (HasSuffixString(file_name, ".dat") || HasSuffixString(file_name, ".upt") ||
+             HasSuffixString(file_name, ".cols"))) {
             return std::make_unique<SegmentStreamConverter>(file_name, file_size, std::move(output_file),
                                                             &column_unique_id_map);
         }
@@ -559,7 +561,10 @@ static Status convert_rowset_meta_pb(RowsetMetaPB* rowset_meta_pb,
 Status ReplicationTxnManager::convert_snapshot_for_none_primary(
         const std::string& tablet_snapshot_path, std::unordered_map<uint32_t, uint32_t>* column_unique_id_map,
         const TReplicateSnapshotRequest& request) {
-    std::string src_header_file_path = tablet_snapshot_path + std::to_string(request.src_tablet_id) + ".hdr";
+    std::string src_tablet_id_path = tablet_snapshot_path + std::to_string(request.src_tablet_id);
+    std::string src_header_file_path = src_tablet_id_path + ".hdr";
+    std::string src_dcgs_snapshot_file_path = src_tablet_id_path + ".dcgs_snapshot";
+
     TabletMeta tablet_meta;
     RETURN_IF_ERROR(tablet_meta.create_from_file(src_header_file_path));
 
@@ -589,6 +594,33 @@ Status ReplicationTxnManager::convert_snapshot_for_none_primary(
         }
     }
 
+    if (fs::path_exist(src_dcgs_snapshot_file_path)) {
+        DeltaColumnGroupSnapshotPB dcg_snapshot_pb;
+        RETURN_IF_ERROR(DeltaColumnGroupListHelper::parse_snapshot(src_dcgs_snapshot_file_path, dcg_snapshot_pb));
+        for (auto& tablet_id : *dcg_snapshot_pb.mutable_tablet_id()) {
+            tablet_id = request.tablet_id;
+        }
+        for (auto& dcg_list : *dcg_snapshot_pb.mutable_dcg_lists()) {
+            for (auto& dcg : *dcg_list.mutable_dcgs()) {
+                for (auto& dcg_column_ids : *dcg.mutable_column_ids()) {
+                    RETURN_IF_ERROR(ReplicationUtils::convert_column_unique_ids(dcg_column_ids.mutable_column_ids(),
+                                                                                *column_unique_id_map));
+                }
+            }
+        }
+
+        std::string dcgs_snapshot_file_path =
+                tablet_snapshot_path + std::to_string(request.tablet_id) + ".dcgs_snapshot";
+        RETURN_IF_ERROR(DeltaColumnGroupListHelper::save_snapshot(dcgs_snapshot_file_path, dcg_snapshot_pb));
+
+        if (request.tablet_id != request.src_tablet_id) {
+            auto status = fs::delete_file(src_dcgs_snapshot_file_path);
+            if (!status.ok()) {
+                LOG(WARNING) << "Failed to delete file: " << src_dcgs_snapshot_file_path << ", " << status;
+            }
+        }
+    }
+
     RETURN_IF_ERROR(SnapshotManager::instance()->convert_rowset_ids(tablet_snapshot_path, request.tablet_id,
                                                                     request.schema_hash));
     return Status::OK();
@@ -616,6 +648,14 @@ Status ReplicationTxnManager::convert_snapshot_for_primary(const std::string& ta
         RETURN_IF_ERROR(convert_rowset_meta_pb(&rowset_meta, column_unique_id_map, request));
     }
 
+    for (auto& [segment_id, dcg_list] : snapshot_meta.delta_column_groups()) {
+        for (auto& dcg : dcg_list) {
+            for (auto& dcg_column_ids : dcg->column_ids()) {
+                RETURN_IF_ERROR(ReplicationUtils::convert_column_unique_ids(&dcg_column_ids, *column_unique_id_map));
+            }
+        }
+    }
+
     RETURN_IF_ERROR(snapshot_meta.serialize_to_file(snapshot_meta_file_path));
 
     RETURN_IF_ERROR(SnapshotManager::instance()->assign_new_rowset_id(&snapshot_meta, tablet_snapshot_path));
diff --git a/be/src/storage/replication_utils.cpp b/be/src/storage/replication_utils.cpp
@@ -336,14 +336,9 @@ StatusOr<std::string> ReplicationUtils::download_remote_snapshot_file(
 
 Status ReplicationUtils::convert_rowset_txn_meta(RowsetTxnMetaPB* rowset_txn_meta,
                                                  const std::unordered_map<uint32_t, uint32_t>& column_unique_id_map) {
-    for (auto& column_unique_id : *rowset_txn_meta->mutable_partial_update_column_unique_ids()) {
-        auto iter = column_unique_id_map.find(column_unique_id);
-        if (iter == column_unique_id_map.end()) {
-            LOG(ERROR) << "Column not found, column unique id: " << column_unique_id;
-            return Status::InternalError("Column not found");
-        }
-        column_unique_id = iter->second;
-    }
+    RETURN_IF_ERROR(convert_column_unique_ids(rowset_txn_meta->mutable_partial_update_column_unique_ids(),
+                                              column_unique_id_map));
+
     if (rowset_txn_meta->has_auto_increment_partial_update_column_uid()) {
         auto iter = column_unique_id_map.find(rowset_txn_meta->auto_increment_partial_update_column_uid());
         if (iter == column_unique_id_map.end()) {
diff --git a/be/src/storage/replication_utils.h b/be/src/storage/replication_utils.h
@@ -45,7 +45,7 @@ class ReplicationUtils {
                                                                TSchemaHash remote_schema_hash,
                                                                const std::string& file_name, uint64_t timeout_sec);
 
-    static constexpr uint32_t kFakeColumnUniqueId = -1;
+    static constexpr uint32_t kFakeColumnUniqueId = INT_MAX;
 
     template <typename T>
     static void calc_column_unique_id_map(const T& source_columns, const T& target_columns,
@@ -68,7 +68,7 @@ class ReplicationUtils {
         }
 
         if (need_convert) {
-            column_unique_id_map->emplace(kFakeColumnUniqueId, 0);
+            column_unique_id_map->emplace(kFakeColumnUniqueId, INT_MAX);
         } else {
             column_unique_id_map->clear();
         }
@@ -86,10 +86,29 @@ class ReplicationUtils {
                 column.set_unique_id(iter->second);
             } else {
                 uint32_t column_unique_id = --column_unique_id_map->operator[](kFakeColumnUniqueId);
-                column.set_unique_id(column_unique_id);
                 column_unique_id_map->emplace(column.unique_id(), column_unique_id);
+                column.set_unique_id(column_unique_id);
+            }
+        }
+    }
+
+    template <typename T>
+    static Status convert_column_unique_ids(T* column_unique_ids,
+                                            const std::unordered_map<uint32_t, uint32_t>& column_unique_id_map) {
+        if (column_unique_id_map.empty()) {
+            return Status::OK();
+        }
+
+        for (auto& column_unique_id : *column_unique_ids) {
+            auto iter = column_unique_id_map.find(column_unique_id);
+            if (iter == column_unique_id_map.end()) {
+                LOG(ERROR) << "Column not found, column unique id: " << column_unique_id;
+                return Status::InternalError("Column not found");
             }
+            column_unique_id = iter->second;
         }
+
+        return Status::OK();
     }
 
     static Status convert_rowset_txn_meta(RowsetTxnMetaPB* rowset_txn_meta,
diff --git a/be/test/CMakeLists.txt b/be/test/CMakeLists.txt
@@ -279,6 +279,7 @@ set(EXEC_FILES
         ./storage/push_handler_test.cpp
         ./storage/range_test.cpp
         ./storage/replication_txn_manager_test.cpp
+        ./storage/replication_utils_test.cpp
         ./storage/segment_stream_converter_test.cpp
         ./storage/row_source_mask_test.cpp
         ./storage/union_iterator_test.cpp
diff --git a/be/test/storage/replication_utils_test.cpp b/be/test/storage/replication_utils_test.cpp
@@ -0,0 +1,46 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "storage/replication_utils.h"
+
+#include <gtest/gtest.h>
+
+namespace starrocks {
+
+class ReplicationUtilsTest : public testing::Test {
+public:
+    ReplicationUtilsTest() = default;
+    ~ReplicationUtilsTest() override = default;
+
+    void SetUp() override {}
+    void TearDown() override {}
+};
+
+TEST_F(ReplicationUtilsTest, test_convert_column_unique_ids) {
+    std::vector<uint32_t> column_unique_ids = {1, 2};
+    std::unordered_map<uint32_t, uint32_t> column_unique_id_map = {{1, 10}, {2, 20}};
+
+    auto status = ReplicationUtils::convert_column_unique_ids(&column_unique_ids, column_unique_id_map);
+    EXPECT_TRUE(status.ok()) << status;
+
+    column_unique_id_map.erase(1);
+    status = ReplicationUtils::convert_column_unique_ids(&column_unique_ids, column_unique_id_map);
+    EXPECT_FALSE(status.ok()) << status;
+
+    column_unique_id_map.clear();
+    status = ReplicationUtils::convert_column_unique_ids(&column_unique_ids, column_unique_id_map);
+    EXPECT_TRUE(status.ok()) << status;
+}
+
+} // namespace starrocks

Original file line number	Diff line number	Diff line change
`@@ -64,6 +64,8 @@ class DeltaColumnGroup {`
`64`	`64`	`}`
`65`	`65`	`return column_files;`
`66`	`66`	`}`
	`67`	`+`
	`68`	`+ std::vector<std::vector<uint32_t>>& column_ids() { return _column_ids; }`
`67`	`69`	`const std::vector<std::vector<uint32_t>>& column_ids() const { return _column_ids; }`
`68`	`70`	`int64_t version() const { return _version; }`
`69`	`71`