Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add new API for searching on sealed segment without index #18830

Merged
merged 1 commit into from
Aug 25, 2022
Merged

Add new API for searching on sealed segment without index #18830

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
42 changes: 35 additions & 7 deletions internal/core/src/query/SearchOnSealed.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,18 +16,20 @@
#include "knowhere/index/vector_index/ConfAdapterMgr.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
#include "query/SearchBruteForce.h"
#include "query/SearchOnSealed.h"
#include "query/helper.h"

namespace milvus::query {

void
SearchOnSealed(const Schema& schema,
const segcore::SealedIndexingRecord& record,
const SearchInfo& search_info,
const void* query_data,
int64_t num_queries,
const BitsetView& bitset,
SearchResult& result) {
SearchOnSealedIndex(const Schema& schema,
const segcore::SealedIndexingRecord& record,
const SearchInfo& search_info,
const void* query_data,
int64_t num_queries,
const BitsetView& bitset,
SearchResult& result) {
auto topk = search_info.topk_;
auto round_decimal = search_info.round_decimal_;

Expand Down Expand Up @@ -77,4 +79,30 @@ SearchOnSealed(const Schema& schema,
std::copy_n(ids, total_num, result.seg_offsets_.data());
std::copy_n(distances, total_num, result.distances_.data());
}

void
SearchOnSealed(const Schema& schema,
const segcore::InsertRecord& record,
const SearchInfo& search_info,
const void* query_data,
int64_t num_queries,
int64_t row_count,
const BitsetView& bitset,
SearchResult& result) {
auto field_id = search_info.field_id_;
auto& field = schema[field_id];

query::dataset::SearchDataset dataset{search_info.metric_type_, num_queries, search_info.topk_,
search_info.round_decimal_, field.get_dim(), query_data};
auto vec_data = record.get_field_data_base(field_id);
AssertInfo(vec_data->num_chunk() == 1, "num chunk not equal to 1 for sealed segment");
auto chunk_data = vec_data->get_chunk_data(0);
auto sub_qr = query::BruteForceSearch(dataset, chunk_data, row_count, bitset);

result.distances_ = std::move(sub_qr.mutable_distances());
result.seg_offsets_ = std::move(sub_qr.mutable_seg_offsets());
result.unity_topK_ = dataset.topk;
result.total_nq_ = dataset.num_queries;
}

} // namespace milvus::query
14 changes: 12 additions & 2 deletions internal/core/src/query/SearchOnSealed.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,23 @@

namespace milvus::query {

void
SearchOnSealedIndex(const Schema& schema,
const segcore::SealedIndexingRecord& record,
const SearchInfo& search_info,
const void* query_data,
int64_t num_queries,
const BitsetView& view,
SearchResult& result);

void
SearchOnSealed(const Schema& schema,
const segcore::SealedIndexingRecord& record,
const segcore::InsertRecord& record,
const SearchInfo& search_info,
const void* query_data,
int64_t num_queries,
const BitsetView& view,
int64_t row_count,
const BitsetView& bitset,
SearchResult& result);

} // namespace milvus::query
4 changes: 2 additions & 2 deletions internal/core/src/segcore/SegmentGrowingImpl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,8 +184,8 @@ SegmentGrowingImpl::vector_search(query::SearchInfo& search_info,
SearchResult& output) const {
auto& sealed_indexing = this->get_sealed_indexing_record();
if (sealed_indexing.is_ready(search_info.field_id_)) {
query::SearchOnSealed(this->get_schema(), sealed_indexing, search_info, query_data, query_count, bitset,
output);
query::SearchOnSealedIndex(this->get_schema(), sealed_indexing, search_info, query_data, query_count, bitset,
output);
} else {
query::SearchOnGrowing(*this, search_info, query_data, query_count, timestamp, bitset, output);
}
Expand Down
28 changes: 8 additions & 20 deletions internal/core/src/segcore/SegmentSealedImpl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -363,27 +363,15 @@ SegmentSealedImpl::vector_search(query::SearchInfo& search_info,
if (get_bit(index_ready_bitset_, field_id)) {
AssertInfo(vector_indexings_.is_ready(field_id),
"vector indexes isn't ready for field " + std::to_string(field_id.get()));
query::SearchOnSealed(*schema_, vector_indexings_, search_info, query_data, query_count, bitset, output);
return;
} else if (!get_bit(field_data_ready_bitset_, field_id)) {
PanicInfo("Field Data is not loaded");
query::SearchOnSealedIndex(*schema_, vector_indexings_, search_info, query_data, query_count, bitset, output);
} else {
AssertInfo(get_bit(field_data_ready_bitset_, field_id),
"Field Data is not loaded: " + std::to_string(field_id.get()));
AssertInfo(row_count_opt_.has_value(), "Can't get row count value");
auto row_count = row_count_opt_.value();
query::SearchOnSealed(*schema_, insert_record_, search_info, query_data, query_count, row_count, bitset,
output);
}

query::dataset::SearchDataset dataset{search_info.metric_type_, query_count, search_info.topk_,
search_info.round_decimal_, field_meta.get_dim(), query_data};
AssertInfo(get_bit(field_data_ready_bitset_, field_id),
"Can't get bitset element at " + std::to_string(field_id.get()));
AssertInfo(row_count_opt_.has_value(), "Can't get row count value");
auto row_count = row_count_opt_.value();
auto vec_data = insert_record_.get_field_data_base(field_id);
AssertInfo(vec_data->num_chunk() == 1, "num chunk not equal to 1 for sealed segment");
auto chunk_data = vec_data->get_chunk_data(0);
auto sub_qr = query::BruteForceSearch(dataset, chunk_data, row_count, bitset);

output.distances_ = std::move(sub_qr.mutable_distances());
output.seg_offsets_ = std::move(sub_qr.mutable_seg_offsets());
output.unity_topK_ = dataset.topk;
output.total_nq_ = dataset.num_queries;
}

void
Expand Down