From 066d19df1e2215451080e0f1ae29c7db20e1f45e Mon Sep 17 00:00:00 2001 From: congqixia Date: Thu, 16 Jul 2026 14:22:39 +0800 Subject: [PATCH] enhance: move sealed segment PK state into runtime snapshot (#51395) Related to #51068 Store PK index slots and virtual PK offset maps in RuntimeResourceState. Build Storage V1 and V2 PK indexes through the unified translator and stage PK replacements until the final state publication. Make PK lookup, range search, bulk subscript, and delete filtering use a single published snapshot. Remove the sealed InsertRecord fallback and clear stale PK resources when external segments become empty. Add regression coverage for atomic PK replacement and zero-row external segment cleanup. Signed-off-by: Congqi Xia --- .../src/segcore/ChunkedSegmentSealedImpl.cpp | 390 ++++++++++-------- .../src/segcore/ChunkedSegmentSealedImpl.h | 79 ++-- internal/core/src/segcore/DeletedRecord.h | 24 +- .../core/src/segcore/DeletedRecordTest.cpp | 3 +- internal/core/src/segcore/SegmentSealed.h | 4 - internal/core/unittest/test_sealed.cpp | 156 ++++++- 6 files changed, 418 insertions(+), 238 deletions(-) diff --git a/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp b/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp index b931e99bb8..cc6c082e2f 100644 --- a/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp +++ b/internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp @@ -333,8 +333,10 @@ ChunkedSegmentSealedImpl::ReadTimestamp( } PinWrapper -ChunkedSegmentSealedImpl::PinPkIndex(milvus::OpContext* op_ctx) const { - auto slot = *pk_index_slot_.rlock(); +ChunkedSegmentSealedImpl::PinPkIndex( + const std::shared_ptr& runtime, + milvus::OpContext* op_ctx) const { + auto slot = runtime != nullptr ? runtime->pk_index_slot : nullptr; if (!slot) { return PinWrapper(nullptr); } @@ -424,13 +426,15 @@ ChunkedSegmentSealedImpl::GetJsonFlatIndexNestedPath( bool ChunkedSegmentSealedImpl::Contain(const PkType& pk) const { - auto schema_snapshot = CaptureSchemaSnapshot(); + auto snapshot = CapturePublishedState(); + auto schema_snapshot = snapshot->schema; + auto runtime = snapshot->runtime; // Zero-storage pk2offset (VirtualPKOffsetMap) resolves PKs by bit-extract. // Skips PinPkIndex + sorted-pk binary search on the virtual PK column. - if (insert_record_.pk2offset_is_zero_storage()) { - return insert_record_.contain(pk); + if (runtime != nullptr && runtime->virtual_pk2offset != nullptr) { + return runtime->virtual_pk2offset->contain(pk); } - auto pk_index = PinPkIndex(nullptr); + auto pk_index = PinPkIndex(runtime, nullptr); if (pk_index.get() != nullptr && pk_index.get()->has_pk2offset()) { return pk_index.get()->contain(pk); } @@ -439,7 +443,7 @@ ChunkedSegmentSealedImpl::Contain(const PkType& pk) const { auto pk_field_id = schema_snapshot->get_primary_field_id().value_or(FieldId(-1)); AssertInfo(pk_field_id.get() != -1, "Primary key is -1"); - auto pk_column = get_column(pk_field_id); + auto pk_column = get_column(runtime, pk_field_id); if (pk_column != nullptr) { auto num_chunks = pk_column->num_chunks(); auto all_chunks = pk_column->GetAllChunks(nullptr); @@ -476,7 +480,7 @@ ChunkedSegmentSealedImpl::Contain(const PkType& pk) const { } } } - return insert_record_.contain(pk); + return false; } bool @@ -566,48 +570,23 @@ ChunkedSegmentSealedImpl::init_storage_v1_timestamp_index( init_storage_v1_timestamp_index(std::move(timestamps), num_rows, nullptr); } -void -ChunkedSegmentSealedImpl::init_storage_v1_pk_index( - FieldId field_id, +std::shared_ptr> +ChunkedSegmentSealedImpl::BuildPkIndexSlot( const std::shared_ptr& column, DataType data_type, - bool is_replace) { - auto schema_snapshot = CaptureSchemaSnapshot(); - if (schema_snapshot->get_primary_field_id().value_or(FieldId(-1)) != - field_id) { - return; - } - // Build compressed offset->pk for FillPrimaryKeys fast path - insert_record_.build_offset2pk(data_type, column.get()); - - if (!is_sorted_by_pk_) { - AssertInfo(field_id.get() != -1, "Primary key is -1"); - if (!is_replace) { - AssertInfo(insert_record_.empty_pks(), - "primary key records already exists, current " - "field id {}", - field_id.get()); - insert_record_.insert_pks(data_type, column.get()); - insert_record_.seal_pks(); - } - } -} - -void -ChunkedSegmentSealedImpl::init_storage_v2_pk_index( - FieldId field_id, - const std::shared_ptr& column, - DataType data_type) { - auto schema_snapshot = CaptureSchemaSnapshot(); - if (schema_snapshot->get_primary_field_id().value_or(FieldId(-1)) != - field_id) { - return; - } + bool eager, + milvus::OpContext* op_ctx) const { std::unique_ptr> translator = std::make_unique( id_, column, data_type, is_sorted_by_pk_); - *pk_index_slot_.wlock() = - Manager::GetInstance().CreateCacheSlot(std::move(translator)); + auto slot = Manager::GetInstance().CreateCacheSlot(std::move(translator)); + if (eager) { + auto cell_holder = SemiInlineGet(slot->PinCells(op_ctx, {0})); + AssertInfo(cell_holder->get_cell_of(0) != nullptr, + "primary key index cache is corrupted, segment {}", + id_); + } + return slot; } bool @@ -999,6 +978,8 @@ ChunkedSegmentSealedImpl::CloneRuntimeResourceState( state->timestamps = current->timestamps; state->timestamp_index = current->timestamp_index; state->timestamp_index_slot = current->timestamp_index_slot; + state->pk_index_slot = current->pk_index_slot; + state->virtual_pk2offset = current->virtual_pk2offset; return state; } @@ -1393,6 +1374,8 @@ ChunkedSegmentSealedImpl::FreezeRuntimeResourceState( runtime->timestamps = current.timestamps; runtime->timestamp_index = current.timestamp_index; runtime->timestamp_index_slot = current.timestamp_index_slot; + runtime->pk_index_slot = current.pk_index_slot; + runtime->virtual_pk2offset = current.virtual_pk2offset; return ToConstRuntimeState(std::move(runtime)); } @@ -2214,18 +2197,28 @@ ChunkedSegmentSealedImpl::SynthesizeExternalSystemFields( } return nullptr; }; + auto pk_field_id = schema_snapshot->get_primary_field_id().value(); int64_t num_rows = segment_load_info.GetNumOfRows(); if (num_rows == 0) { + runtime->fields.erase(pk_field_id); + runtime->pk_index_slot.reset(); + runtime->virtual_pk2offset.reset(); + if (!schema_snapshot->IsExternalDataField(pk_field_id)) { + runtime->virtual_pk2offset = + std::make_shared(id_, 0); + } + auto timestamps = std::make_shared(); + timestamps->InitFromOwnedData({}); + runtime->timestamps = std::move(timestamps); + runtime->timestamp_index = std::make_shared(); + runtime->timestamp_index_slot.reset(); { std::unique_lock lck(mutex_); update_row_count(0); - // Initialize empty timestamps so system fields remain query-visible. - insert_record_.init_timestamps_from_owned({}, TimestampIndex()); } return; } - auto pk_field_id = schema_snapshot->get_primary_field_id().value(); if (!schema_snapshot->IsExternalDataField(pk_field_id)) { // 1. VirtualPKChunkedColumn for the synthetic primary key. // This is lazy; data is only materialized if DataOfChunk/Span is called. @@ -2237,7 +2230,9 @@ ChunkedSegmentSealedImpl::SynthesizeExternalSystemFields( // Virtual PK = (seg_id << 32) | offset, so pk to offset is a simple // bit-extract. This replaces the OffsetOrderedArray that would // otherwise store num_rows (pk, offset) pairs (~17 GB for 1B rows). - insert_record_.set_virtual_pk_offset_map(id_, num_rows); + runtime->virtual_pk2offset = + std::make_shared(id_, num_rows); + runtime->pk_index_slot.reset(); } else { AssertInfo( get_runtime_column(pk_field_id) != nullptr, @@ -2258,7 +2253,6 @@ ChunkedSegmentSealedImpl::SynthesizeExternalSystemFields( } else { // Synthetic timestamps: constant mode (all 0; rows always visible). // No data is materialized, saving ~8 GB for 1B-row external tables. - insert_record_.init_timestamps_constant(num_rows, 0); auto timestamps = std::make_shared(); timestamps->InitConstant(num_rows, 0); runtime->timestamps = timestamps; @@ -3926,38 +3920,36 @@ ChunkedSegmentSealedImpl::search_pks(BitsetType& bitset, if (pks.empty()) { return; } + auto snapshot = CapturePublishedState(); + auto runtime = snapshot->runtime; BitsetTypeView bitset_view(bitset); // See Contain() — same zero-storage pk2offset fast path. - if (insert_record_.pk2offset_is_zero_storage()) { + if (runtime != nullptr && runtime->virtual_pk2offset != nullptr) { for (auto& pk : pks) { - insert_record_.search_pk_range( - pk, proto::plan::OpType::Equal, bitset_view); + runtime->virtual_pk2offset->find_range( + pk, proto::plan::OpType::Equal, bitset_view, [](int64_t) { + return true; + }); } return; } if (!is_sorted_by_pk_) { - auto pk_index = PinPkIndex(nullptr); + auto pk_index = PinPkIndex(runtime, nullptr); auto* pk_cell = pk_index.get(); - AssertInfo(pk_cell != nullptr || !insert_record_.empty_pks(), + AssertInfo(pk_cell != nullptr && pk_cell->has_pk2offset(), "primary key index is not ready"); for (auto& pk : pks) { - if (pk_cell != nullptr) { - pk_cell->pk2offset().find_range( - pk, - proto::plan::OpType::Equal, - bitset_view, - [](int64_t offset) { return true; }); - } else { - insert_record_.search_pk_range( - pk, proto::plan::OpType::Equal, bitset_view); - } + pk_cell->pk2offset().find_range( + pk, + proto::plan::OpType::Equal, + bitset_view, + [](int64_t offset) { return true; }); } return; } - auto snapshot = CapturePublishedState(); auto schema_snapshot = snapshot->schema; auto pk_field_id = schema_snapshot->get_primary_field_id().value_or(FieldId(-1)); @@ -3992,9 +3984,9 @@ ChunkedSegmentSealedImpl::search_batch_pks( callback) const { // Helper to read a single timestamp by segment offset. // For import/CDC segments with commit_ts_ set: every row carries commit_ts_, - // so short-circuit without touching the raw timestamp column or insert_record_. + // so short-circuit without touching the raw timestamp column. // For StorageV2: pins the timestamp column and indexes into chunks. - // For StorageV1: reads from insert_record_ directly. + // PK lookup and timestamp data come from the same published runtime. auto snapshot = CapturePublishedState(); auto runtime = snapshot->runtime; auto effective_commit_ts = @@ -4008,14 +4000,14 @@ ChunkedSegmentSealedImpl::search_batch_pks( // Avoid the sorted-PK column scan below: external segments synthesize PKs // with VirtualPKChunkedColumn, which intentionally does not support // GetAllChunks(). - if (insert_record_.pk2offset_is_zero_storage()) { + if (runtime != nullptr && runtime->virtual_pk2offset != nullptr) { auto timestamp_hit = include_same_ts ? [](Timestamp lhs, Timestamp rhs) { return lhs <= rhs; } : [](Timestamp lhs, Timestamp rhs) { return lhs < rhs; }; for (size_t i = 0; i < pks.size(); i++) { auto timestamp = get_timestamp(i); - for (auto offset : insert_record_.pk2offset_->find(pks[i])) { + for (auto offset : runtime->virtual_pk2offset->find(pks[i])) { auto insert_ts = read_ts(offset); if (timestamp_hit(insert_ts, timestamp)) { callback(SegOffset(offset), timestamp); @@ -4027,17 +4019,18 @@ ChunkedSegmentSealedImpl::search_batch_pks( // handle unsorted case if (!is_sorted_by_pk_) { - auto pk_index = PinPkIndex(nullptr); + auto pk_index = PinPkIndex(runtime, nullptr); auto* pk_cell = pk_index.get(); + if (pk_cell == nullptr || !pk_cell->has_pk2offset()) { + return; + } auto timestamp_hit = include_same_ts ? [](Timestamp lhs, Timestamp rhs) { return lhs <= rhs; } : [](Timestamp lhs, Timestamp rhs) { return lhs < rhs; }; for (size_t i = 0; i < pks.size(); i++) { auto timestamp = get_timestamp(i); - auto offsets = pk_cell != nullptr - ? pk_cell->pk2offset().find(pks[i]) - : insert_record_.pk2offset_->find(pks[i]); + auto offsets = pk_cell->pk2offset().find(pks[i]); for (auto offset : offsets) { auto insert_ts = read_ts(offset); if (timestamp_hit(insert_ts, timestamp)) { @@ -4137,38 +4130,39 @@ ChunkedSegmentSealedImpl::pk_range(milvus::OpContext* op_ctx, proto::plan::OpType op, const PkType& pk, BitsetTypeView& bitset) const { + auto snapshot = CapturePublishedState(); + auto runtime = snapshot->runtime; // See Contain() — same zero-storage pk2offset fast path. - if (insert_record_.pk2offset_is_zero_storage()) { - insert_record_.search_pk_range(pk, op, bitset); + if (runtime != nullptr && runtime->virtual_pk2offset != nullptr) { + runtime->virtual_pk2offset->find_range( + pk, op, bitset, [](int64_t) { return true; }); return; } if (!is_sorted_by_pk_) { - auto pk_index = PinPkIndex(op_ctx); + auto pk_index = PinPkIndex(runtime, op_ctx); auto* pk_cell = pk_index.get(); - AssertInfo(pk_cell != nullptr || !insert_record_.empty_pks(), + AssertInfo(pk_cell != nullptr && pk_cell->has_pk2offset(), "primary key index is not ready"); - if (pk_cell != nullptr) { - pk_cell->pk2offset().find_range( - pk, op, bitset, [](int64_t offset) { return true; }); - } else { - insert_record_.search_pk_range(pk, op, bitset); - } + pk_cell->pk2offset().find_range( + pk, op, bitset, [](int64_t offset) { return true; }); return; } - search_sorted_pk_range(op_ctx, op, pk, bitset); + search_sorted_pk_range(op_ctx, op, pk, bitset, snapshot); } void -ChunkedSegmentSealedImpl::search_sorted_pk_range(milvus::OpContext* op_ctx, - proto::plan::OpType op, - const PkType& pk, - BitsetTypeView& bitset) const { - auto schema_snapshot = CaptureSchemaSnapshot(); +ChunkedSegmentSealedImpl::search_sorted_pk_range( + milvus::OpContext* op_ctx, + proto::plan::OpType op, + const PkType& pk, + BitsetTypeView& bitset, + const std::shared_ptr& snapshot) const { + auto schema_snapshot = snapshot->schema; auto pk_field_id = schema_snapshot->get_primary_field_id().value_or(FieldId(-1)); AssertInfo(pk_field_id.get() != -1, "Primary key is -1"); - auto pk_column = get_column(pk_field_id); + auto pk_column = get_column(snapshot->runtime, pk_field_id); AssertInfo(pk_column != nullptr, "primary key column not loaded"); switch (schema_snapshot->get_fields().at(pk_field_id).get_data_type()) { @@ -4196,46 +4190,50 @@ ChunkedSegmentSealedImpl::pk_binary_range(milvus::OpContext* op_ctx, const PkType& upper_pk, bool upper_inclusive, BitsetTypeView& bitset) const { + auto snapshot = CapturePublishedState(); + auto runtime = snapshot->runtime; // See Contain() — same zero-storage pk2offset fast path. - if (insert_record_.pk2offset_is_zero_storage()) { - insert_record_.search_pk_binary_range( - lower_pk, lower_inclusive, upper_pk, upper_inclusive, bitset); + if (runtime != nullptr && runtime->virtual_pk2offset != nullptr) { + auto lower_op = lower_inclusive ? proto::plan::OpType::GreaterEqual + : proto::plan::OpType::GreaterThan; + auto upper_op = upper_inclusive ? proto::plan::OpType::LessEqual + : proto::plan::OpType::LessThan; + BitsetType upper_result(bitset.size()); + auto upper_view = upper_result.view(); + runtime->virtual_pk2offset->find_range( + lower_pk, lower_op, bitset, [](int64_t) { return true; }); + runtime->virtual_pk2offset->find_range( + upper_pk, upper_op, upper_view, [](int64_t) { return true; }); + bitset &= upper_result; return; } if (!is_sorted_by_pk_) { - auto pk_index = PinPkIndex(op_ctx); + auto pk_index = PinPkIndex(runtime, op_ctx); auto* pk_cell = pk_index.get(); - AssertInfo(pk_cell != nullptr || !insert_record_.empty_pks(), + AssertInfo(pk_cell != nullptr && pk_cell->has_pk2offset(), "primary key index is not ready"); - if (pk_cell != nullptr) { - auto lower_op = lower_inclusive ? proto::plan::OpType::GreaterEqual - : proto::plan::OpType::GreaterThan; - auto upper_op = upper_inclusive ? proto::plan::OpType::LessEqual - : proto::plan::OpType::LessThan; - BitsetType upper_result(bitset.size()); - auto upper_view = upper_result.view(); - pk_cell->pk2offset().find_range( - lower_pk, lower_op, bitset, [](int64_t offset) { - return true; - }); - pk_cell->pk2offset().find_range( - upper_pk, upper_op, upper_view, [](int64_t offset) { - return true; - }); - bitset &= upper_result; - } else { - insert_record_.search_pk_binary_range( - lower_pk, lower_inclusive, upper_pk, upper_inclusive, bitset); - } + auto lower_op = lower_inclusive ? proto::plan::OpType::GreaterEqual + : proto::plan::OpType::GreaterThan; + auto upper_op = upper_inclusive ? proto::plan::OpType::LessEqual + : proto::plan::OpType::LessThan; + BitsetType upper_result(bitset.size()); + auto upper_view = upper_result.view(); + pk_cell->pk2offset().find_range( + lower_pk, lower_op, bitset, [](int64_t offset) { return true; }); + pk_cell->pk2offset().find_range( + upper_pk, upper_op, upper_view, [](int64_t offset) { + return true; + }); + bitset &= upper_result; return; } // For sorted segments, use binary search - auto schema_snapshot = CaptureSchemaSnapshot(); + auto schema_snapshot = snapshot->schema; auto pk_field_id = schema_snapshot->get_primary_field_id().value_or(FieldId(-1)); AssertInfo(pk_field_id.get() != -1, "Primary key is -1"); - auto pk_column = get_column(pk_field_id); + auto pk_column = get_column(runtime, pk_field_id); AssertInfo(pk_column != nullptr, "primary key column not loaded"); switch (schema_snapshot->get_fields().at(pk_field_id).get_data_type()) { @@ -4269,15 +4267,16 @@ ChunkedSegmentSealedImpl::pk_binary_range(milvus::OpContext* op_ctx, std::pair, bool> ChunkedSegmentSealedImpl::find_first_n(int64_t limit, const BitsetTypeView& bitset) const { + auto runtime = CaptureRuntimeResourceState(); + if (runtime != nullptr && runtime->virtual_pk2offset != nullptr) { + return runtime->virtual_pk2offset->find_first_n(limit, bitset); + } if (!is_sorted_by_pk_) { - auto pk_index = PinPkIndex(nullptr); + auto pk_index = PinPkIndex(runtime, nullptr); auto* pk_cell = pk_index.get(); - AssertInfo(pk_cell != nullptr || !insert_record_.empty_pks(), + AssertInfo(pk_cell != nullptr && pk_cell->has_pk2offset(), "primary key index is not ready"); - if (pk_cell != nullptr) { - return pk_cell->pk2offset().find_first_n(limit, bitset); - } - return insert_record_.pk2offset_->find_first_n(limit, bitset); + return pk_cell->pk2offset().find_first_n(limit, bitset); } if (limit == Unlimited || limit == NoLimit) { limit = num_rows_.value(); @@ -4313,9 +4312,18 @@ ChunkedSegmentSealedImpl::find_first_n_element( const BitsetTypeView& element_bitset, const IArrayOffsets* array_offsets, const std::optional& cursor) const { + auto snapshot = CapturePublishedState(); + auto runtime = snapshot->runtime; + if (runtime != nullptr && runtime->virtual_pk2offset != nullptr) { + return runtime->virtual_pk2offset->find_first_n_element( + limit, element_bitset, array_offsets, cursor); + } if (!is_sorted_by_pk_) { - // Not sorted by PK, use pk2offset_ to iterate in PK order - return insert_record_.pk2offset_->find_first_n_element( + auto pk_index = PinPkIndex(runtime, nullptr); + auto* pk_cell = pk_index.get(); + AssertInfo(pk_cell != nullptr && pk_cell->has_pk2offset(), + "primary key index is not ready"); + return pk_cell->pk2offset().find_first_n_element( limit, element_bitset, array_offsets, cursor); } @@ -4332,11 +4340,11 @@ ChunkedSegmentSealedImpl::find_first_n_element( // already been returned. std::optional cursor_doc_offset; if (cursor.has_value()) { - auto schema_snapshot = CaptureSchemaSnapshot(); + auto schema_snapshot = snapshot->schema; auto pk_field_id = schema_snapshot->get_primary_field_id().value_or(FieldId(-1)); AssertInfo(pk_field_id.get() != -1, "Primary key is -1"); - auto pk_column = get_column(pk_field_id); + auto pk_column = get_column(runtime, pk_field_id); AssertInfo(pk_column != nullptr, "primary key column not loaded"); switch (schema_snapshot->get_fields().at(pk_field_id).get_data_type()) { case DataType::INT64: @@ -4408,31 +4416,50 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl( mmap_descriptor_(storage::MmapManager::GetInstance() .GetMmapChunkManager() ->Register()), - insert_record_(*schema, MAX_ROW_COUNT), id_(segment_id), col_index_meta_(index_meta), is_sorted_by_pk_(is_sorted_by_pk), deleted_record_( - &insert_record_, - [this, schema]( - const std::vector& pks, - const Timestamp* timestamps, - const std::function& callback) { - if (commit_ts_ != 0) { + nullptr, + [this](const std::vector& pks, + const Timestamp* timestamps, + const std::function& callback) { + auto snapshot = CapturePublishedState(); + auto runtime = snapshot->runtime; + auto schema = snapshot->schema; + if (snapshot->commit_ts != 0) { // For import segments with commit_ts, row timestamps are - // overwritten to commit_ts. Skip the timestamp filter in - // PK search so that deletes with ts < commit_ts can still - // find the matching rows. Pass the original delete - // timestamp to the callback for correct storage. + // overwritten to commit_ts. Filter with the same published + // snapshot used for PK lookup, so pre-commit and + // same-timestamp deletes never reach DeletedRecord. + auto delete_is_after_insert = [&](size_t i) { + return timestamps[i] > snapshot->commit_ts; + }; if (!is_sorted_by_pk_) { - auto pk_index = PinPkIndex(nullptr); + if (runtime != nullptr && + runtime->virtual_pk2offset != nullptr) { + for (size_t i = 0; i < pks.size(); i++) { + if (!delete_is_after_insert(i)) { + continue; + } + for (auto offset : + runtime->virtual_pk2offset->find(pks[i])) { + callback(SegOffset(offset), timestamps[i]); + } + } + return; + } + auto pk_index = PinPkIndex(runtime, nullptr); auto* pk_cell = pk_index.get(); + if (pk_cell == nullptr || !pk_cell->has_pk2offset()) { + return; + } for (size_t i = 0; i < pks.size(); i++) { - auto offsets = - pk_cell != nullptr - ? pk_cell->pk2offset().find(pks[i]) - : insert_record_.pk2offset_->find(pks[i]); + if (!delete_is_after_insert(i)) { + continue; + } + auto offsets = pk_cell->pk2offset().find(pks[i]); for (auto offset : offsets) { callback(SegOffset(offset), timestamps[i]); } @@ -4441,7 +4468,6 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl( auto pk_field_id = schema->get_primary_field_id().value_or(FieldId(-1)); AssertInfo(pk_field_id.get() != -1, "Primary key is -1"); - auto runtime = CaptureRuntimeResourceState(); auto pk_column = get_column(runtime, pk_field_id); AssertInfo(pk_column != nullptr, "primary key column not loaded"); @@ -4461,6 +4487,9 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl( auto num_rows_until_chunk = pk_column->GetNumRowsUntilChunk(i); for (size_t j = 0; j < pks.size(); j++) { + if (!delete_is_after_insert(j)) { + continue; + } auto target = std::get(pks[j]); auto it = std::lower_bound( src, src + chunk_row_num, target); @@ -4485,6 +4514,9 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl( auto string_chunk = static_cast(pw.get()); for (size_t j = 0; j < pks.size(); ++j) { + if (!delete_is_after_insert(j)) { + continue; + } auto& target = std::get(pks[j]); auto offset = @@ -4521,11 +4553,6 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl( } }, segment_id) { - deleted_record_.set_get_insert_timestamp_func( - [this](int64_t row_id) -> Timestamp { - return ReadTimestamp( - row_id, CaptureRuntimeResourceState(), EffectiveCommitTs()); - }); auto load_info = std::make_shared( milvus::proto::segcore::SegmentLoadInfo(), schema); std::atomic_store(&published_state_, @@ -4956,9 +4983,7 @@ ChunkedSegmentSealedImpl::ClearData() { ngram_indexings_.withWLock([&](auto& ngram_indexings) { cancel_and_clear_ngram_indexings(ngram_indexings); }); - insert_record_.clear(); timestamp_index_slot_.wlock()->reset(); - pk_index_slot_.wlock()->reset(); fields_.wlock()->clear(); variable_fields_avg_size_.clear(); stats_.mem_size = 0; @@ -5863,23 +5888,17 @@ ChunkedSegmentSealedImpl::bulk_subscript(milvus::OpContext* op_ctx, // Fast path for int64 PK field: use compressed offset2pk index auto pk_field_id = snapshot->schema->get_primary_field_id(); - auto pk_index = PinPkIndex(op_ctx); + auto pk_index = PinPkIndex(snapshot->runtime, op_ctx); if (pk_field_id.has_value() && pk_field_id.value() == field_id && field_meta.get_data_type() == DataType::INT64 && - (pk_index.get() != nullptr ? pk_index.get()->has_int64_pk_index() - : insert_record_.has_int64_pk_index())) { + pk_index.get() != nullptr && pk_index.get()->has_int64_pk_index()) { auto ret = fill_with_empty(field_id, count); auto* output = ret->mutable_scalars() ->mutable_long_data() ->mutable_data() ->mutable_data(); - if (pk_index.get() != nullptr) { - pk_index.get()->bulk_get_int64_pks_by_offsets( - seg_offsets, count, output); - } else { - insert_record_.bulk_get_int64_pks_by_offsets( - seg_offsets, count, output); - } + pk_index.get()->bulk_get_int64_pks_by_offsets( + seg_offsets, count, output); return ret; } @@ -6184,7 +6203,9 @@ SegcoreError ChunkedSegmentSealedImpl::Delete(int64_t size, const IdArray* ids, const Timestamp* timestamps_raw) { - auto schema_snapshot = CaptureSchemaSnapshot(); + auto snapshot = CapturePublishedState(); + auto schema_snapshot = snapshot->schema; + auto runtime = snapshot->runtime; auto field_id = schema_snapshot->get_primary_field_id().value_or(FieldId(-1)); AssertInfo(field_id.get() != -1, "Primary key is -1"); @@ -6197,20 +6218,23 @@ ChunkedSegmentSealedImpl::Delete(int64_t size, for (int i = 0; i < size; i++) { ordering[i] = std::make_tuple(timestamps_raw[i], pks[i]); } - // if insert record is empty (may be only-load meta but not data for lru-cache at go side), - // filtering may cause the deletion lost, skip the filtering to avoid it. - auto pk_index = PinPkIndex(nullptr); - auto has_pk_index = pk_index.get() != nullptr ? !pk_index.get()->empty_pks() - : !insert_record_.empty_pks(); + // If PK state is unavailable (for example, only metadata is loaded by the + // Go-side cache), filtering could lose deletions, so preserve them all. + auto pk_index = PinPkIndex(runtime, nullptr); + auto virtual_pk2offset = + runtime != nullptr ? runtime->virtual_pk2offset : nullptr; + auto has_pk_index = + virtual_pk2offset != nullptr || + (pk_index.get() != nullptr && !pk_index.get()->empty_pks()); if (has_pk_index) { auto end = std::remove_if( ordering.begin(), ordering.end(), [&](const std::tuple& record) { - if (pk_index.get() != nullptr) { - return !pk_index.get()->contain(std::get<1>(record)); + if (virtual_pk2offset != nullptr) { + return !virtual_pk2offset->contain(std::get<1>(record)); } - return !insert_record_.contain(std::get<1>(record)); + return !pk_index.get()->contain(std::get<1>(record)); }); size = end - ordering.begin(); ordering.resize(size); @@ -6600,13 +6624,16 @@ ChunkedSegmentSealedImpl::load_field_data_common( } } - if (schema_snapshot->get_primary_field_id().value_or(FieldId(-1)) == - field_id) { - if (segment_load_info.GetStorageVersion() >= STORAGE_V2) { - init_storage_v2_pk_index(field_id, column, data_type); - } else { - init_storage_v1_pk_index(field_id, column, data_type, is_replace); - } + std::shared_ptr> pk_index_slot; + const bool is_primary_field = + schema_snapshot->get_primary_field_id().value_or(FieldId(-1)) == + field_id; + if (is_primary_field) { + pk_index_slot = + BuildPkIndexSlot(column, + data_type, + segment_load_info.GetStorageVersion() < STORAGE_V2, + op_ctx); } generate_interim_index(field_id, num_rows, column, op_ctx, committer); @@ -6641,6 +6668,11 @@ ChunkedSegmentSealedImpl::load_field_data_common( const PublishedSegmentState& state_snapshot) { prepare_array_offsets(target_runtime); + if (is_primary_field) { + target_runtime.pk_index_slot = pk_index_slot; + target_runtime.virtual_pk2offset.reset(); + } + if (is_replace) { if (old_column && !enable_mmap) { if (!is_proxy_column || diff --git a/internal/core/src/segcore/ChunkedSegmentSealedImpl.h b/internal/core/src/segcore/ChunkedSegmentSealedImpl.h index 65efabef5e..3e2123bc49 100644 --- a/internal/core/src/segcore/ChunkedSegmentSealedImpl.h +++ b/internal/core/src/segcore/ChunkedSegmentSealedImpl.h @@ -86,10 +86,6 @@ namespace milvus::segcore { -namespace storagev1translator { -class InsertRecordTranslator; -} - namespace storagev2translator { class TimestampIndexCell; class PkIndexCell; @@ -374,6 +370,9 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { std::shared_ptr timestamp_index; std::shared_ptr> timestamp_index_slot; + std::shared_ptr> + pk_index_slot; + std::shared_ptr virtual_pk2offset; }; struct PublishedSegmentState { @@ -423,11 +422,6 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { return stats_.mem_size.load() + deleted_record_.mem_size(); } - InsertRecord& - get_insert_record() override { - return insert_record_; - } - int64_t get_row_count() const override; @@ -452,10 +446,12 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { BitsetTypeView& bitset) const override; void - search_sorted_pk_range(milvus::OpContext* op_ctx, - proto::plan::OpType op, - const PkType& pk, - BitsetTypeView& bitset) const; + search_sorted_pk_range( + milvus::OpContext* op_ctx, + proto::plan::OpType op, + const PkType& pk, + BitsetTypeView& bitset, + const std::shared_ptr& snapshot) const; void pk_binary_range(milvus::OpContext* op_ctx, @@ -2147,7 +2143,8 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { std::optional effective_commit_ts) const; PinWrapper - PinPkIndex(milvus::OpContext* op_ctx) const; + PinPkIndex(const std::shared_ptr& runtime, + milvus::OpContext* op_ctx) const; void init_storage_v2_timestamp_index( @@ -2155,18 +2152,11 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { size_t num_rows, const std::string& warmup_policy = ""); - void - init_storage_v1_pk_index( - FieldId field_id, - const std::shared_ptr& column, - DataType data_type, - bool is_replace); - - void - init_storage_v2_pk_index( - FieldId field_id, - const std::shared_ptr& column, - DataType data_type); + std::shared_ptr> + BuildPkIndexSlot(const std::shared_ptr& column, + DataType data_type, + bool eager, + milvus::OpContext* op_ctx) const; private: std::unique_ptr @@ -2190,14 +2180,9 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { folly::Synchronized> scalar_indexings_; - // inserted fields data and row_ids, timestamps - InsertRecord insert_record_; folly::Synchronized< std::shared_ptr>> timestamp_index_slot_; - folly::Synchronized< - std::shared_ptr>> - pk_index_slot_; // deleted pks mutable DeletedRecord deleted_record_; @@ -2371,6 +2356,38 @@ class ChunkedSegmentSealedImpl : public SegmentSealed { committer.Publish(current, final_delta); } + template + void + TestStageLoadFieldDataThenPublish( + FieldId field_id, + const std::shared_ptr& column, + size_t num_rows, + DataType data_type, + const SchemaPtr& schema_snapshot, + std::shared_ptr runtime, + PublishedSegmentState* staged_state, + const std::shared_ptr& current, + StateDelta& final_delta, + Verifier&& verifier) { + StagedStateCommitter committer(*this, runtime.get(), staged_state); + load_field_data_common(field_id, + column, + num_rows, + data_type, + /*enable_mmap=*/false, + /*is_proxy_column=*/false, + *current->load_info, + schema_snapshot, + runtime.get(), + std::nullopt, + nullptr, + /*is_replace=*/true, + &committer); + verifier(); + final_delta.runtime = ToConstRuntimeState(std::move(runtime)); + committer.Publish(current, final_delta); + } + template void TestStageLoadIndexGenerateInterimThenPublish( diff --git a/internal/core/src/segcore/DeletedRecord.h b/internal/core/src/segcore/DeletedRecord.h index 06b4d2e087..7b784c593a 100644 --- a/internal/core/src/segcore/DeletedRecord.h +++ b/internal/core/src/segcore/DeletedRecord.h @@ -158,14 +158,11 @@ class DeletedRecord { // so only rows with insert_ts < delete_ts reach this callback. // This check is therefore redundant for normal production // callers and keeps direct callers/tests on the same boundary. - // For import segments with commit_timestamp, insert_ts is set - // to commit_ts so that pre-commit deletes (delete_ts < commit_ts) - // are correctly rejected — the row's data only becomes visible - // at commit_ts. + // Sealed-segment search callbacks perform the same check from + // one published snapshot before invoking this callback. Timestamp insert_ts = 0; - if (get_insert_timestamp_func_) { - insert_ts = get_insert_timestamp_func_(row_id); - } else if (!insert_record_->timestamps_.empty()) { + if (insert_record_ != nullptr && + !insert_record_->timestamps_.empty()) { insert_ts = insert_record_->timestamps_[row_id]; } if (insert_ts != 0 && delete_ts <= insert_ts) { @@ -409,15 +406,6 @@ class DeletedRecord { deleted_mask_.resize(row_count); } - // Set a callback to read insert timestamp for a given row_id. - // Used by StorageV2 lazy-init path where insert_record_.timestamps_ - // may be empty but timestamp data is available in the column. - void - set_get_insert_timestamp_func( - std::function func) { - get_insert_timestamp_func_ = std::move(func); - } - std::vector> get_snapshots() const { std::shared_lock lock(snap_lock_); @@ -456,10 +444,6 @@ class DeletedRecord { // estimated memory size of DeletedRecord, only used for sealed segment int64_t estimated_memory_size_{0}; - // Callback to read insert timestamp from column (StorageV2 lazy path). - // Used when insert_record_->timestamps_ is empty. - std::function get_insert_timestamp_func_; - // atomic snapshot for fast path query optimization // when query_timestamp >= snapshot.max_ts, we can directly use the bitset // without traversing the SkipList diff --git a/internal/core/src/segcore/DeletedRecordTest.cpp b/internal/core/src/segcore/DeletedRecordTest.cpp index 2868b9b92d..bdf503fefc 100644 --- a/internal/core/src/segcore/DeletedRecordTest.cpp +++ b/internal/core/src/segcore/DeletedRecordTest.cpp @@ -61,9 +61,8 @@ TEST(DeleteMVCC, common_case) { auto pks = dataset.get_col(pk); auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); ASSERT_EQ(c, segment->get_real_count()); - auto& insert_record = segment->get_insert_record(); DeletedRecord delete_record( - &insert_record, + nullptr, [&pks](const std::vector& delete_pks, const Timestamp* timestamps, std::function diff --git a/internal/core/src/segcore/SegmentSealed.h b/internal/core/src/segcore/SegmentSealed.h index 2187ebe649..b40c646e58 100644 --- a/internal/core/src/segcore/SegmentSealed.h +++ b/internal/core/src/segcore/SegmentSealed.h @@ -24,7 +24,6 @@ #include "index/JsonFlatIndex.h" #include "pb/index_cgo_msg.pb.h" #include "pb/segcore.pb.h" -#include "segcore/InsertRecord.h" #include "segcore/SegmentInterface.h" #include "segcore/Types.h" #include "index/NgramInvertedIndex.h" @@ -56,9 +55,6 @@ class SegmentSealed : public SegmentInternalInterface { const int64_t* ids, int64_t count) const = 0; - virtual InsertRecord& - get_insert_record() = 0; - virtual PinWrapper GetNgramIndex(milvus::OpContext* op_ctx, FieldId field_id) const override = 0; diff --git a/internal/core/unittest/test_sealed.cpp b/internal/core/unittest/test_sealed.cpp index ffdf328856..fb19794fd6 100644 --- a/internal/core/unittest/test_sealed.cpp +++ b/internal/core/unittest/test_sealed.cpp @@ -69,13 +69,13 @@ #include "query/PlanImpl.h" #include "query/Utils.h" #include "segcore/ChunkedSegmentSealedImpl.h" -#include "segcore/InsertRecord.h" #include "segcore/SegcoreConfig.h" #include "segcore/SegmentGrowing.h" #include "segcore/SegmentGrowingImpl.h" #include "segcore/SegmentLoadInfo.h" #include "segcore/SegmentSealed.h" #include "segcore/Types.h" +#include "segcore/storagev2translator/SystemIndexTranslator.h" #include "storage/FileManager.h" #include "storage/InsertData.h" #include "storage/PayloadReader.h" @@ -1433,7 +1433,6 @@ TEST(Sealed, DeleteCount) { auto dataset = DataGen(schema, N); auto segment = CreateSealedWithFieldDataLoaded(schema, dataset); - segment->get_insert_record().seal_pks(); int64_t c = 10; ASSERT_EQ(segment->get_deleted_count(), 0); @@ -4678,6 +4677,123 @@ TEST(SealedSegmentCowState, ReplaceScalarIndexStagesRuntimeUntilFinalPublish) { EXPECT_EQ(next->runtime->ngram_fields.count(payload), 0); } +TEST(SealedSegmentCowState, ReplacePkStateIsInvisibleUntilFinalPublish) { + auto schema = std::make_shared(); + auto pk = schema->AddDebugField("pk", DataType::INT64); + schema->set_primary_field_id(pk); + + auto old_dataset = DataGen(schema, 4); + auto segment = CreateSealedWithFieldDataLoaded(schema, old_dataset); + auto* sealed = dynamic_cast(segment.get()); + ASSERT_NE(sealed, nullptr); + + auto replacement_dataset = DataGen(schema, + 4, + /*seed=*/314159, + /*ts_offset=*/0, + /*repeat_count=*/1, + /*array_len=*/10, + /*group_count=*/1, + /*random_pk=*/true); + auto replacement_segment = + CreateSealedWithFieldDataLoaded(schema, replacement_dataset); + auto* replacement_sealed = + dynamic_cast(replacement_segment.get()); + ASSERT_NE(replacement_sealed, nullptr); + + auto old_pks = old_dataset.get_col(pk); + auto replacement_pks = replacement_dataset.get_col(pk); + auto old_only = + std::find_if(old_pks.begin(), old_pks.end(), [&](int64_t value) { + return std::find(replacement_pks.begin(), + replacement_pks.end(), + value) == replacement_pks.end(); + }); + auto replacement_only = std::find_if( + replacement_pks.begin(), replacement_pks.end(), [&](int64_t value) { + return std::find(old_pks.begin(), old_pks.end(), value) == + old_pks.end(); + }); + ASSERT_NE(old_only, old_pks.end()); + ASSERT_NE(replacement_only, replacement_pks.end()); + PkType old_pk = *old_only; + PkType replacement_pk = *replacement_only; + + auto current = sealed->TestGetPublishedStateSnapshot(); + ASSERT_NE(current->runtime->pk_index_slot, nullptr); + ASSERT_TRUE(sealed->Contain(old_pk)); + ASSERT_FALSE(sealed->Contain(replacement_pk)); + auto old_pk_index_slot = current->runtime->pk_index_slot; + + auto replacement_state = + replacement_sealed->TestGetPublishedStateSnapshot(); + auto replacement_column = replacement_state->runtime->fields.at(pk); + + auto stage_replacement = [&](bool fail_before_publish) { + auto runtime = sealed->TestCloneMutableRuntimeResourceState(); + ChunkedSegmentSealedImpl::StateDelta initial_delta; + initial_delta.schema = current->schema; + initial_delta.load_info = current->load_info; + initial_delta.runtime = sealed->TestFreezeRuntimeResourceState(runtime); + initial_delta.commit_ts = current->commit_ts; + auto staged = + sealed->TestBuildNextPublishedState(current, initial_delta); + + ChunkedSegmentSealedImpl::StateDelta final_delta; + final_delta.schema = current->schema; + final_delta.load_info = current->load_info; + final_delta.commit_ts = current->commit_ts; + + auto stage_and_publish = [&] { + sealed->TestStageLoadFieldDataThenPublish( + pk, + replacement_column, + replacement_dataset.raw_->num_rows(), + DataType::INT64, + schema, + runtime, + staged.get(), + current, + final_delta, + [&] { + ASSERT_NE(runtime->pk_index_slot, nullptr); + EXPECT_NE(runtime->pk_index_slot, + current->runtime->pk_index_slot); + + auto still_published = + sealed->TestGetPublishedStateSnapshot(); + EXPECT_EQ(still_published, current); + EXPECT_TRUE(sealed->Contain(old_pk)); + EXPECT_FALSE(sealed->Contain(replacement_pk)); + if (fail_before_publish) { + throw std::runtime_error("abort staged PK publish"); + } + }); + }; + + if (fail_before_publish) { + EXPECT_THROW(stage_and_publish(), std::runtime_error); + } else { + EXPECT_NO_THROW(stage_and_publish()); + } + }; + + stage_replacement(true); + EXPECT_EQ(sealed->TestGetPublishedStateSnapshot(), current); + EXPECT_TRUE(sealed->Contain(old_pk)); + EXPECT_FALSE(sealed->Contain(replacement_pk)); + + stage_replacement(false); + auto published = sealed->TestGetPublishedStateSnapshot(); + ASSERT_NE(published, current); + EXPECT_NE(published->runtime->pk_index_slot, + current->runtime->pk_index_slot); + EXPECT_FALSE(sealed->Contain(old_pk)); + EXPECT_TRUE(sealed->Contain(replacement_pk)); + + EXPECT_EQ(current->runtime->pk_index_slot, old_pk_index_slot); +} + TEST(SealedSegmentCowState, ClearPublishedStateDropsRuntimeSnapshot) { auto schema = std::make_shared(); auto pk = schema->AddDebugField("pk", DataType::INT64); @@ -5056,6 +5172,42 @@ TEST(SealedSegmentCowState, ExternalSyntheticFieldsUseStagedLoadInfoRows) { EXPECT_EQ(published->load_info->GetNumOfRows(), 3); } +TEST(SealedSegmentCowState, ExternalZeroRowsClearsStagedPkState) { + auto schema = std::make_shared(); + auto pk = schema->AddDebugField("pk", DataType::INT64); + schema->set_primary_field_id(pk); + schema->set_external_source("s3://bucket/data"); + schema->set_external_spec(R"({"format":"parquet"})"); + + auto segment = CreateSealedSegment(schema); + auto* sealed = dynamic_cast(segment.get()); + ASSERT_NE(sealed, nullptr); + + auto runtime = sealed->TestCloneMutableRuntimeResourceState(); + proto::segcore::SegmentLoadInfo nonempty_proto; + nonempty_proto.set_segmentid(1005); + nonempty_proto.set_num_of_rows(7); + SegmentLoadInfo nonempty_load_info(nonempty_proto, schema); + sealed->TestSynthesizeExternalSystemFields( + nonempty_load_info, schema, runtime.get()); + ASSERT_TRUE(runtime->fields.count(pk) > 0); + ASSERT_NE(runtime->virtual_pk2offset, nullptr); + + proto::segcore::SegmentLoadInfo empty_proto; + empty_proto.set_segmentid(1005); + empty_proto.set_num_of_rows(0); + SegmentLoadInfo empty_load_info(empty_proto, schema); + sealed->TestSynthesizeExternalSystemFields( + empty_load_info, schema, runtime.get()); + + EXPECT_EQ(runtime->fields.count(pk), 0); + EXPECT_EQ(runtime->pk_index_slot, nullptr); + ASSERT_NE(runtime->virtual_pk2offset, nullptr); + EXPECT_FALSE(runtime->virtual_pk2offset->contain(PkType{int64_t{0}})); + ASSERT_NE(runtime->timestamps, nullptr); + EXPECT_TRUE(runtime->timestamps->empty()); +} + TEST(SealedSegmentCowState, PublishRuntimeStateRecomputesExternalSystemFieldReady) { auto schema = std::make_shared();