From 305601ad250ead9d68f7e15921d691c3f82bab44 Mon Sep 17 00:00:00 2001 From: Cai Yudong Date: Thu, 25 Aug 2022 18:16:55 +0800 Subject: [PATCH] Add API SearchOnSealedChunk() to search on sealed segment without index (#18830) Signed-off-by: yudong.cai Signed-off-by: yudong.cai --- internal/core/src/query/SearchOnSealed.cpp | 42 +++++++++++++++---- internal/core/src/query/SearchOnSealed.h | 14 ++++++- .../core/src/segcore/SegmentGrowingImpl.cpp | 4 +- .../core/src/segcore/SegmentSealedImpl.cpp | 28 ++++--------- 4 files changed, 57 insertions(+), 31 deletions(-) diff --git a/internal/core/src/query/SearchOnSealed.cpp b/internal/core/src/query/SearchOnSealed.cpp index 0e4d4b6eee..4bb193c13e 100644 --- a/internal/core/src/query/SearchOnSealed.cpp +++ b/internal/core/src/query/SearchOnSealed.cpp @@ -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_; @@ -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 diff --git a/internal/core/src/query/SearchOnSealed.h b/internal/core/src/query/SearchOnSealed.h index f813ca0de5..756fc987c7 100644 --- a/internal/core/src/query/SearchOnSealed.h +++ b/internal/core/src/query/SearchOnSealed.h @@ -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 diff --git a/internal/core/src/segcore/SegmentGrowingImpl.cpp b/internal/core/src/segcore/SegmentGrowingImpl.cpp index 5c7e1e5a9e..94aaa6eb08 100644 --- a/internal/core/src/segcore/SegmentGrowingImpl.cpp +++ b/internal/core/src/segcore/SegmentGrowingImpl.cpp @@ -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); } diff --git a/internal/core/src/segcore/SegmentSealedImpl.cpp b/internal/core/src/segcore/SegmentSealedImpl.cpp index bbd77c408e..a294ed3547 100644 --- a/internal/core/src/segcore/SegmentSealedImpl.cpp +++ b/internal/core/src/segcore/SegmentSealedImpl.cpp @@ -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