[Enhancement] Primary key consistency test supports PK parallel execution (#63698)

luohaha · web-flow · commit 4918a1fadcef · 2025-10-09T10:22:07.000+08:00
Signed-off-by: luohaha &lt;18810541851@163.com&gt;
diff --git a/be/src/storage/sstable/table_builder.cpp b/be/src/storage/sstable/table_builder.cpp
@@ -66,7 +66,6 @@ TableBuilder::TableBuilder(const Options& options, WritableFile* file) : rep_(ne
 }
 
 TableBuilder::~TableBuilder() {
-    assert(rep_->closed); // Catch errors where caller forgot to call Finish()
     delete rep_->filter_block;
     delete rep_;
 }
diff --git a/be/test/storage/lake/lake_primary_key_consistency_test.cpp b/be/test/storage/lake/lake_primary_key_consistency_test.cpp
@@ -99,15 +99,15 @@ class Replayer {
     }
 
     bool check(const ChunkPtr& chunk) {
-        std::map<int, std::pair<int, int>> tmp_chunk;
+        std::map<std::string, std::pair<int, int>> tmp_chunk;
         for (int i = 0; i < chunk->num_rows(); i++) {
-            if (tmp_chunk.count(chunk->columns()[0]->get(i).get_int32()) > 0) {
+            if (tmp_chunk.count(chunk->columns()[0]->get(i).get_slice().to_string()) > 0) {
                 // duplicate pk
-                LOG(ERROR) << "duplicate pk: " << chunk->columns()[0]->get(i).get_int32();
+                LOG(ERROR) << "duplicate pk: " << chunk->columns()[0]->get(i).get_slice().to_string();
                 return false;
             }
-            tmp_chunk[chunk->columns()[0]->get(i).get_int32()] = {chunk->columns()[1]->get(i).get_int32(),
-                                                                  chunk->columns()[2]->get(i).get_int32()};
+            tmp_chunk[chunk->columns()[0]->get(i).get_slice().to_string()] = {chunk->columns()[1]->get(i).get_int32(),
+                                                                              chunk->columns()[2]->get(i).get_int32()};
         }
         if (tmp_chunk.size() != _replayer_index.size()) {
             LOG(ERROR) << "inconsistency row number, actual : " << tmp_chunk.size()
@@ -146,24 +146,24 @@ class Replayer {
             if (log.op == ReplayerOP::UPSERT) {
                 // Upsert
                 for (int i = 0; i < chunk->num_rows(); i++) {
-                    _replayer_index[chunk->columns()[0]->get(i).get_int32()] = {
+                    _replayer_index[chunk->columns()[0]->get(i).get_slice().to_string()] = {
                             chunk->columns()[1]->get(i).get_int32(), chunk->columns()[2]->get(i).get_int32()};
                 }
             } else if (log.op == ReplayerOP::ERASE) {
                 // Delete
                 for (int i = 0; i < chunk->num_rows(); i++) {
-                    _replayer_index.erase(chunk->columns()[0]->get(i).get_int32());
+                    _replayer_index.erase(chunk->columns()[0]->get(i).get_slice().to_string());
                 }
             } else if (log.op == ReplayerOP::PARTIAL_UPSERT || log.op == ReplayerOP::PARTIAL_UPDATE) {
                 // Partial update
                 for (int i = 0; i < chunk->num_rows(); i++) {
-                    auto iter = _replayer_index.find(chunk->columns()[0]->get(i).get_int32());
+                    auto iter = _replayer_index.find(chunk->columns()[0]->get(i).get_slice().to_string());
                     if (iter != _replayer_index.end()) {
-                        _replayer_index[chunk->columns()[0]->get(i).get_int32()] = {
+                        _replayer_index[chunk->columns()[0]->get(i).get_slice().to_string()] = {
                                 chunk->columns()[1]->get(i).get_int32(), iter->second.second};
                     } else if (log.op == ReplayerOP::PARTIAL_UPSERT) {
                         // insert new record with default val
-                        _replayer_index[chunk->columns()[0]->get(i).get_int32()] = {
+                        _replayer_index[chunk->columns()[0]->get(i).get_slice().to_string()] = {
                                 chunk->columns()[1]->get(i).get_int32(), 0};
                     } else {
                         // do nothing
@@ -186,11 +186,11 @@ class Replayer {
                     return false;
                 };
                 for (int i = 0; i < chunk->num_rows(); i++) {
-                    auto iter = _replayer_index.find(chunk->columns()[0]->get(i).get_int32());
+                    auto iter = _replayer_index.find(chunk->columns()[0]->get(i).get_slice().to_string());
                     if (iter == _replayer_index.end() || is_condition_meet_fn(iter->second, i)) {
                         // update if condition meet or not found
                         // insert new record
-                        _replayer_index[chunk->columns()[0]->get(i).get_int32()] = {
+                        _replayer_index[chunk->columns()[0]->get(i).get_slice().to_string()] = {
                                 chunk->columns()[1]->get(i).get_int32(), chunk->columns()[2]->get(i).get_int32()};
                     }
                 }
@@ -206,7 +206,7 @@ class Replayer {
     // logs for replay.
     std::vector<ReplayEntry> _redo_logs;
     // c0 -> <c1, c2>
-    std::map<int, std::pair<int, int>> _replayer_index;
+    std::map<std::string, std::pair<int, int>> _replayer_index;
 };
 
 class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterface<PrimaryKeyParam> {
@@ -216,8 +216,8 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
         _tablet_metadata->set_enable_persistent_index(true);
         _tablet_metadata->set_persistent_index_type(GetParam().persistent_index_type);
 
-        _slots.emplace_back(0, "c0", TypeDescriptor{LogicalType::TYPE_INT});
-        _partial_slots.emplace_back(0, "c0", TypeDescriptor{LogicalType::TYPE_INT});
+        _slots.emplace_back(0, "c0", TypeDescriptor{LogicalType::TYPE_VARCHAR});
+        _partial_slots.emplace_back(0, "c0", TypeDescriptor{LogicalType::TYPE_VARCHAR});
         _slots.emplace_back(1, "c1", TypeDescriptor{LogicalType::TYPE_INT});
         _partial_slots.emplace_back(1, "c1", TypeDescriptor{LogicalType::TYPE_INT});
         _slots.emplace_back(2, "c2", TypeDescriptor{LogicalType::TYPE_INT});
@@ -269,13 +269,16 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
         config::enable_pindex_minor_compaction = false;
         _old_enable_pk_strict_memcheck = config::enable_pk_strict_memcheck;
         config::enable_pk_strict_memcheck = false;
+        _old_pk_parallel_execution_threshold_bytes = config::pk_parallel_execution_threshold_bytes;
+        config::pk_parallel_execution_threshold_bytes = 1;
     }
 
     void TearDown() override {
         (void)fs::remove_all(kTestGroupPath);
         config::l0_max_mem_usage = _old_l0_size;
         config::write_buffer_size = _old_memtable_size;
         config::enable_pk_strict_memcheck = _old_enable_pk_strict_memcheck;
+        config::pk_parallel_execution_threshold_bytes = _old_pk_parallel_execution_threshold_bytes;
     }
 
     std::shared_ptr<TabletMetadataPB> generate_tablet_metadata(KeysType keys_type) {
@@ -287,7 +290,7 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
         //
         //  | column | type | KEY | NULL |
         //  +--------+------+-----+------+
-        //  |   c0   |  INT | YES |  NO  |
+        //  |   c0   |  STRING | YES |  NO  |
         //  |   c1   |  INT | NO  |  NO  |
         //  |   c2   |  INT | NO  |  NO  |
         auto schema = metadata->mutable_schema();
@@ -299,9 +302,10 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
         {
             c0->set_unique_id(next_id());
             c0->set_name("c0");
-            c0->set_type("INT");
+            c0->set_type("VARCHAR");
             c0->set_is_key(true);
             c0->set_is_nullable(false);
+            c0->set_length(3200);
         }
         auto c1 = schema->add_column();
         {
@@ -327,14 +331,22 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     std::pair<ChunkPtr, std::vector<uint32_t>> gen_upsert_data(bool is_upsert) {
         const size_t chunk_size = (size_t)_random_generator->random_n();
         std::vector<std::vector<int>> cols(3);
+        std::vector<std::string> key_col_str;
+        std::vector<Slice> key_col;
         std::vector<uint8_t> v3(chunk_size, is_upsert ? TOpType::UPSERT : TOpType::DELETE);
         _random_generator->random_cols(chunk_size, &cols);
+        for (size_t i = 0; i < chunk_size; i++) {
+            key_col_str.emplace_back(std::to_string(cols[0][i]));
+        }
+        for (const auto& s : key_col_str) {
+            key_col.emplace_back(Slice(s));
+        }
 
-        auto c0 = Int32Column::create();
+        auto c0 = BinaryColumn::create();
         auto c1 = Int32Column::create();
         auto c2 = Int32Column::create();
         auto c3 = Int8Column::create();
-        c0->append_numbers(cols[0].data(), cols[0].size() * sizeof(int));
+        c0->append_strings(key_col.data(), key_col.size());
         c1->append_numbers(cols[1].data(), cols[1].size() * sizeof(int));
         c2->append_numbers(cols[2].data(), cols[2].size() * sizeof(int));
         c3->append_numbers(v3.data(), v3.size() * sizeof(uint8_t));
@@ -350,12 +362,20 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     std::pair<ChunkPtr, std::vector<uint32_t>> gen_partial_update_data() {
         const size_t chunk_size = (size_t)_random_generator->random_n();
         std::vector<std::vector<int>> cols(2);
+        std::vector<std::string> key_col_str;
+        std::vector<Slice> key_col;
         std::vector<uint8_t> v3(chunk_size, TOpType::UPSERT);
         _random_generator->random_cols(chunk_size, &cols);
+        for (size_t i = 0; i < chunk_size; i++) {
+            key_col_str.emplace_back(std::to_string(cols[0][i]));
+        }
+        for (const auto& s : key_col_str) {
+            key_col.emplace_back(Slice(s));
+        }
 
-        auto c0 = Int32Column::create();
+        auto c0 = BinaryColumn::create();
         auto c1 = Int32Column::create();
-        c0->append_numbers(cols[0].data(), cols[0].size() * sizeof(int));
+        c0->append_strings(key_col.data(), key_col.size());
         c1->append_numbers(cols[1].data(), cols[1].size() * sizeof(int));
         auto indexes = std::vector<uint32_t>(chunk_size);
         for (uint32_t i = 0; i < chunk_size; i++) {
@@ -364,6 +384,29 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
         return {std::make_shared<Chunk>(Columns{std::move(c0), std::move(c1)}, _slot_cid_map), std::move(indexes)};
     }
 
+    // 5% chance to force index memtable flush
+    std::unique_ptr<ConfigResetGuard<int64_t>> random_force_index_mem_flush() {
+        std::unique_ptr<ConfigResetGuard<int64_t>> force_flush_guard;
+        uint32_t r = _random_generator->random() % 100;
+        if (r < 5) {
+            // 5% chance to force index memtable flush
+            force_flush_guard = std::make_unique<ConfigResetGuard<int64_t>>(&config::l0_max_mem_usage, 1);
+        }
+        return force_flush_guard;
+    }
+
+    // 20% chance to enable pk parallel execution
+    std::unique_ptr<ConfigResetGuard<bool>> random_pk_parallel_execution() {
+        std::unique_ptr<ConfigResetGuard<bool>> pk_parallel_execution_guard;
+        uint32_t r = _random_generator->random() % 100;
+        if (r < 20) {
+            // 20% chance to enable pk parallel execution
+            pk_parallel_execution_guard =
+                    std::make_unique<ConfigResetGuard<bool>>(&config::enable_pk_parallel_execution, true);
+        }
+        return pk_parallel_execution_guard;
+    }
+
     ChunkPtr read(int64_t tablet_id, int64_t version) {
         ASSIGN_OR_ABORT(auto metadata, _tablet_mgr->get_tablet_metadata(tablet_id, version));
         auto reader = std::make_shared<TabletReader>(_tablet_mgr.get(), metadata, *_schema);
@@ -383,6 +426,8 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     }
 
     Status upsert_op() {
+        std::unique_ptr<ConfigResetGuard<int64_t>> force_index_mem_flush_guard = random_force_index_mem_flush();
+        std::unique_ptr<ConfigResetGuard<bool>> pk_parallel_execution_guard = random_pk_parallel_execution();
         auto txn_id = next_id();
         ASSIGN_OR_ABORT(auto delta_writer, DeltaWriterBuilder()
                                                    .set_tablet_manager(_tablet_mgr.get())
@@ -418,6 +463,8 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     }
 
     Status partial_update_op(PartialUpdateMode mode) {
+        std::unique_ptr<ConfigResetGuard<int64_t>> force_index_mem_flush_guard = random_force_index_mem_flush();
+        std::unique_ptr<ConfigResetGuard<bool>> pk_parallel_execution_guard = random_pk_parallel_execution();
         auto txn_id = next_id();
         ASSIGN_OR_ABORT(auto delta_writer, DeltaWriterBuilder()
                                                    .set_tablet_manager(_tablet_mgr.get())
@@ -453,6 +500,8 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     }
 
     Status condition_update() {
+        std::unique_ptr<ConfigResetGuard<int64_t>> force_index_mem_flush_guard = random_force_index_mem_flush();
+        std::unique_ptr<ConfigResetGuard<bool>> pk_parallel_execution_guard = random_pk_parallel_execution();
         auto txn_id = next_id();
         // c2 as merge_condition
         std::string merge_condition = "c2";
@@ -484,6 +533,8 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     }
 
     Status upsert_with_batch_pub_op() {
+        std::unique_ptr<ConfigResetGuard<int64_t>> force_index_mem_flush_guard = random_force_index_mem_flush();
+        std::unique_ptr<ConfigResetGuard<bool>> pk_parallel_execution_guard = random_pk_parallel_execution();
         size_t batch_cnt = std::max(_random_generator->random() % MaxBatchCnt, (size_t)1);
         std::vector<int64_t> txn_ids;
         for (int i = 0; i < batch_cnt; i++) {
@@ -525,6 +576,8 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     }
 
     Status delete_op() {
+        std::unique_ptr<ConfigResetGuard<int64_t>> force_index_mem_flush_guard = random_force_index_mem_flush();
+        std::unique_ptr<ConfigResetGuard<bool>> pk_parallel_execution_guard = random_pk_parallel_execution();
         auto chunk_index = gen_upsert_data(false);
         auto txn_id = next_id();
         ASSIGN_OR_ABORT(auto delta_writer, DeltaWriterBuilder()
@@ -550,6 +603,8 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     }
 
     Status compact_op() {
+        std::unique_ptr<ConfigResetGuard<int64_t>> force_index_mem_flush_guard = random_force_index_mem_flush();
+        std::unique_ptr<ConfigResetGuard<bool>> pk_parallel_execution_guard = random_pk_parallel_execution();
         auto txn_id = next_id();
         auto task_context = std::make_unique<CompactionTaskContext>(txn_id, _tablet_metadata->id(), _version, false,
                                                                     false, nullptr);
@@ -658,6 +713,7 @@ class LakePrimaryKeyConsistencyTest : public TestBase, testing::WithParamInterfa
     int64_t _old_l0_size = 0;
     int64_t _old_memtable_size = 0;
     bool _old_enable_pk_strict_memcheck = false;
+    int64_t _old_pk_parallel_execution_threshold_bytes = 0;
 };
 
 TEST_P(LakePrimaryKeyConsistencyTest, test_local_pk_consistency) {

Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,6 @@ TableBuilder::TableBuilder(const Options& options, WritableFile* file) : rep_(ne`
`66`	`66`	`}`
`67`	`67`
`68`	`68`	`TableBuilder::~TableBuilder() {`
`69`		`- assert(rep_->closed); // Catch errors where caller forgot to call Finish()`
`70`	`69`	`delete rep_->filter_block;`
`71`	`70`	`delete rep_;`
`72`	`71`	`}`