furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 02:05:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

enhance: move sealed segment PK state into runtime snapshot (#51395)

Browse Source 
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 <congqi.xia@zilliz.com>
This commit is contained in:
congqixia
2026-07-16 14:22:39 +08:00
committed by GitHub
parent 4e6bea5789
commit 066d19df1e
6 changed files with 418 additions and 238 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp
+211 -179
View File
@@ -333,8 +333,10 @@ ChunkedSegmentSealedImpl::ReadTimestamp(
}
PinWrapper<const storagev2translator::PkIndexCell*>
ChunkedSegmentSealedImpl::PinPkIndex(milvus::OpContext* op_ctx) const {
auto slot = *pk_index_slot_.rlock();
ChunkedSegmentSealedImpl::PinPkIndex(
const std::shared_ptr<const RuntimeResourceState>& runtime,
milvus::OpContext* op_ctx) const {
auto slot = runtime != nullptr ? runtime->pk_index_slot : nullptr;
if (!slot) {
return PinWrapper<const storagev2translator::PkIndexCell*>(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<CacheSlot<storagev2translator::PkIndexCell>>
ChunkedSegmentSealedImpl::BuildPkIndexSlot(
const std::shared_ptr<ChunkedColumnInterface>& 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<ChunkedColumnInterface>& 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<storagev2translator::PkIndexCell>> translator =
std::make_unique<storagev2translator::PkIndexTranslator>(
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<const VirtualPKOffsetMap>(id_, 0);
}
auto timestamps = std::make_shared<TimestampData>();
timestamps->InitFromOwnedData({});
runtime->timestamps = std::move(timestamps);
runtime->timestamp_index = std::make_shared<const TimestampIndex>();
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<const VirtualPKOffsetMap>(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<TimestampData>();
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<const PublishedSegmentState>& 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<std::vector<OffsetMap::OffsetType>, 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<QueryIteratorCursor>& 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<int64_t> 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<PkType>& pks,
const Timestamp* timestamps,
const std::function<void(const SegOffset offset,
const Timestamp ts)>& callback) {
if (commit_ts_ != 0) {
nullptr,
[this](const std::vector<PkType>& pks,
const Timestamp* timestamps,
const std::function<void(const SegOffset offset,
const Timestamp ts)>& 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<int64_t>(pks[j]);
auto it = std::lower_bound(
src, src + chunk_row_num, target);
@@ -4485,6 +4514,9 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl(
auto string_chunk =
static_cast<StringChunk*>(pw.get());
for (size_t j = 0; j < pks.size(); ++j) {
if (!delete_is_after_insert(j)) {
continue;
}
auto& target =
std::get<std::string>(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<const SegmentLoadInfo>(
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<Timestamp, PkType>& 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<CacheSlot<storagev2translator::PkIndexCell>> 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 ||
internal/core/src/segcore/ChunkedSegmentSealedImpl.h
+48 -31
View File
@@ -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<const TimestampIndex> timestamp_index;
std::shared_ptr<CacheSlot<storagev2translator::TimestampIndexCell>>
timestamp_index_slot;
std::shared_ptr<CacheSlot<storagev2translator::PkIndexCell>>
pk_index_slot;
std::shared_ptr<const OffsetMap> virtual_pk2offset;
};
struct PublishedSegmentState {
@@ -423,11 +422,6 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
return stats_.mem_size.load() + deleted_record_.mem_size();
}
InsertRecord<true>&
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<const PublishedSegmentState>& snapshot) const;
void
pk_binary_range(milvus::OpContext* op_ctx,
@@ -2147,7 +2143,8 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
std::optional<Timestamp> effective_commit_ts) const;
PinWrapper<const storagev2translator::PkIndexCell*>
PinPkIndex(milvus::OpContext* op_ctx) const;
PinPkIndex(const std::shared_ptr<const RuntimeResourceState>& 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<ChunkedColumnInterface>& column,
DataType data_type,
bool is_replace);
void
init_storage_v2_pk_index(
FieldId field_id,
const std::shared_ptr<ChunkedColumnInterface>& column,
DataType data_type);
std::shared_ptr<CacheSlot<storagev2translator::PkIndexCell>>
BuildPkIndexSlot(const std::shared_ptr<ChunkedColumnInterface>& column,
DataType data_type,
bool eager,
milvus::OpContext* op_ctx) const;
private:
std::unique_ptr<SpanBase>
@@ -2190,14 +2180,9 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
folly::Synchronized<std::unordered_map<FieldId, index::CacheIndexBasePtr>>
scalar_indexings_;
// inserted fields data and row_ids, timestamps
InsertRecord<true> insert_record_;
folly::Synchronized<
std::shared_ptr<CacheSlot<storagev2translator::TimestampIndexCell>>>
timestamp_index_slot_;
folly::Synchronized<
std::shared_ptr<CacheSlot<storagev2translator::PkIndexCell>>>
pk_index_slot_;
// deleted pks
mutable DeletedRecord<true> deleted_record_;
@@ -2371,6 +2356,38 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
committer.Publish(current, final_delta);
}
template <typename Verifier>
void
TestStageLoadFieldDataThenPublish(
FieldId field_id,
const std::shared_ptr<ChunkedColumnInterface>& column,
size_t num_rows,
DataType data_type,
const SchemaPtr& schema_snapshot,
std::shared_ptr<RuntimeResourceState> runtime,
PublishedSegmentState* staged_state,
const std::shared_ptr<const PublishedSegmentState>& 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 <typename Verifier>
void
TestStageLoadIndexGenerateInterimThenPublish(
internal/core/src/segcore/DeletedRecord.h
+4 -20
View File
@@ -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<Timestamp(int64_t row_id)> func) {
get_insert_timestamp_func_ = std::move(func);
}
std::vector<std::pair<Timestamp, BitsetType>>
get_snapshots() const {
std::shared_lock<std::shared_mutex> 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<Timestamp(int64_t row_id)> 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
internal/core/src/segcore/DeletedRecordTest.cpp
+1 -2
View File
@@ -61,9 +61,8 @@ TEST(DeleteMVCC, common_case) {
auto pks = dataset.get_col<int64_t>(pk);
auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);
ASSERT_EQ(c, segment->get_real_count());
auto& insert_record = segment->get_insert_record();
DeletedRecord<true> delete_record(
&insert_record,
nullptr,
[&pks](const std::vector<PkType>& delete_pks,
const Timestamp* timestamps,
std::function<void(const SegOffset offset, const Timestamp ts)>
internal/core/src/segcore/SegmentSealed.h
-4
View File
@@ -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<true>&
get_insert_record() = 0;
virtual PinWrapper<index::NgramInvertedIndex*>
GetNgramIndex(milvus::OpContext* op_ctx,
FieldId field_id) const override = 0;
internal/core/unittest/test_sealed.cpp
+154 -2
View File
@@ -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<Schema>();
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<ChunkedSegmentSealedImpl*>(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<ChunkedSegmentSealedImpl*>(replacement_segment.get());
ASSERT_NE(replacement_sealed, nullptr);
auto old_pks = old_dataset.get_col<int64_t>(pk);
auto replacement_pks = replacement_dataset.get_col<int64_t>(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<Schema>();
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<Schema>();
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<ChunkedSegmentSealedImpl*>(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<Schema>();