mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-07 17:48:29 +08:00
d3784c6515
issue: #44212 Implement search/query storage usage statistics in go side(result reduce), only record storage usage in vector search C++ path. Need to be implemented in query c++ path in next prs. --------- Signed-off-by: chasingegg <chao.gao@zilliz.com> Signed-off-by: marcelo.chen <marcelo.chen@zilliz.com> Co-authored-by: marcelo.chen <marcelo.chen@zilliz.com>
272 lines
12 KiB
C++
272 lines
12 KiB
C++
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
|
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
|
// or implied. See the License for the specific language governing permissions and limitations under the License
|
|
|
|
#include "common/BitsetView.h"
|
|
#include "common/QueryInfo.h"
|
|
#include "common/Tracer.h"
|
|
#include "common/Types.h"
|
|
#include "SearchOnGrowing.h"
|
|
#include <cstddef>
|
|
#include "knowhere/comp/index_param.h"
|
|
#include "knowhere/config.h"
|
|
#include "log/Log.h"
|
|
#include "query/CachedSearchIterator.h"
|
|
#include "query/SearchBruteForce.h"
|
|
#include "query/SearchOnIndex.h"
|
|
#include "exec/operator/Utils.h"
|
|
|
|
namespace milvus::query {
|
|
|
|
void
|
|
FloatSegmentIndexSearch(const segcore::SegmentGrowingImpl& segment,
|
|
const SearchInfo& info,
|
|
const void* query_data,
|
|
int64_t num_queries,
|
|
const BitsetView& bitset,
|
|
milvus::OpContext* op_context,
|
|
SearchResult& search_result) {
|
|
auto& schema = segment.get_schema();
|
|
auto& indexing_record = segment.get_indexing_record();
|
|
auto& record = segment.get_insert_record();
|
|
|
|
auto vecfield_id = info.field_id_;
|
|
auto& field = schema[vecfield_id];
|
|
auto is_sparse = field.get_data_type() == DataType::VECTOR_SPARSE_U32_F32;
|
|
// TODO(SPARSE): see todo in PlanImpl.h::PlaceHolder.
|
|
auto dim = is_sparse ? 0 : field.get_dim();
|
|
|
|
AssertInfo(IsVectorDataType(field.get_data_type()),
|
|
"[FloatSearch]Field data type isn't VECTOR_FLOAT, "
|
|
"VECTOR_FLOAT16, VECTOR_BFLOAT16 or VECTOR_SPARSE_U32_F32");
|
|
dataset::SearchDataset search_dataset{info.metric_type_,
|
|
num_queries,
|
|
info.topk_,
|
|
info.round_decimal_,
|
|
dim,
|
|
query_data};
|
|
if (indexing_record.is_in(vecfield_id)) {
|
|
const auto& field_indexing =
|
|
indexing_record.get_vec_field_indexing(vecfield_id);
|
|
|
|
auto indexing = field_indexing.get_segment_indexing();
|
|
SearchInfo search_conf = field_indexing.get_search_params(info);
|
|
auto vec_index = dynamic_cast<index::VectorIndex*>(indexing.get());
|
|
SearchOnIndex(search_dataset,
|
|
*vec_index,
|
|
search_conf,
|
|
bitset,
|
|
op_context,
|
|
search_result,
|
|
is_sparse);
|
|
}
|
|
}
|
|
|
|
void
|
|
SearchOnGrowing(const segcore::SegmentGrowingImpl& segment,
|
|
const SearchInfo& info,
|
|
const void* query_data,
|
|
const size_t* query_lims,
|
|
int64_t num_queries,
|
|
Timestamp timestamp,
|
|
const BitsetView& bitset,
|
|
milvus::OpContext* op_context,
|
|
SearchResult& search_result) {
|
|
auto& schema = segment.get_schema();
|
|
auto& record = segment.get_insert_record();
|
|
auto active_count =
|
|
std::min(int64_t(bitset.size()), segment.get_active_count(timestamp));
|
|
|
|
// step 1.1: get meta
|
|
// step 1.2: get which vector field to search
|
|
auto vecfield_id = info.field_id_;
|
|
auto& field = schema[vecfield_id];
|
|
CheckBruteForceSearchParam(field, info);
|
|
|
|
auto data_type = field.get_data_type();
|
|
auto element_type = field.get_element_type();
|
|
AssertInfo(IsVectorDataType(data_type),
|
|
"[SearchOnGrowing]Data type isn't vector type");
|
|
|
|
auto topk = info.topk_;
|
|
auto metric_type = info.metric_type_;
|
|
auto round_decimal = info.round_decimal_;
|
|
|
|
// step 2: small indexing search
|
|
if (segment.get_indexing_record().SyncDataWithIndex(field.get_id())) {
|
|
AssertInfo(
|
|
data_type != DataType::VECTOR_ARRAY,
|
|
"vector array(embedding list) is not supported for growing segment "
|
|
"indexing search");
|
|
|
|
FloatSegmentIndexSearch(segment,
|
|
info,
|
|
query_data,
|
|
num_queries,
|
|
bitset,
|
|
op_context,
|
|
search_result);
|
|
} else {
|
|
std::shared_lock<std::shared_mutex> read_chunk_mutex(
|
|
segment.get_chunk_mutex());
|
|
// check SyncDataWithIndex() again, in case the vector chunks has been removed.
|
|
if (segment.get_indexing_record().SyncDataWithIndex(field.get_id())) {
|
|
AssertInfo(data_type != DataType::VECTOR_ARRAY,
|
|
"vector array(embedding list) is not supported for "
|
|
"growing segment indexing search");
|
|
|
|
return FloatSegmentIndexSearch(segment,
|
|
info,
|
|
query_data,
|
|
num_queries,
|
|
bitset,
|
|
op_context,
|
|
search_result);
|
|
}
|
|
SubSearchResult final_qr(num_queries, topk, metric_type, round_decimal);
|
|
// TODO(SPARSE): see todo in PlanImpl.h::PlaceHolder.
|
|
auto dim = field.get_data_type() == DataType::VECTOR_SPARSE_U32_F32
|
|
? 0
|
|
: field.get_dim();
|
|
dataset::SearchDataset search_dataset{metric_type,
|
|
num_queries,
|
|
topk,
|
|
round_decimal,
|
|
dim,
|
|
query_data,
|
|
query_lims};
|
|
int32_t current_chunk_id = 0;
|
|
|
|
// get K1 and B from index for bm25 brute force
|
|
std::map<std::string, std::string> index_info;
|
|
if (metric_type == knowhere::metric::BM25) {
|
|
index_info = segment.get_indexing_record()
|
|
.get_field_index_meta(vecfield_id)
|
|
.GetIndexParams();
|
|
}
|
|
|
|
// step 3: brute force search where small indexing is unavailable
|
|
auto vec_ptr = record.get_data_base(vecfield_id);
|
|
|
|
if (info.iterator_v2_info_.has_value()) {
|
|
AssertInfo(data_type != DataType::VECTOR_ARRAY,
|
|
"vector array(embedding list) is not supported for "
|
|
"vector iterator");
|
|
|
|
CachedSearchIterator cached_iter(search_dataset,
|
|
vec_ptr,
|
|
active_count,
|
|
info,
|
|
index_info,
|
|
bitset,
|
|
data_type);
|
|
cached_iter.NextBatch(info, search_result);
|
|
return;
|
|
}
|
|
|
|
auto vec_size_per_chunk = vec_ptr->get_size_per_chunk();
|
|
auto max_chunk = upper_div(active_count, vec_size_per_chunk);
|
|
|
|
for (int chunk_id = current_chunk_id; chunk_id < max_chunk;
|
|
++chunk_id) {
|
|
auto chunk_data = vec_ptr->get_chunk_data(chunk_id);
|
|
|
|
auto element_begin = chunk_id * vec_size_per_chunk;
|
|
auto element_end =
|
|
std::min(active_count, (chunk_id + 1) * vec_size_per_chunk);
|
|
auto size_per_chunk = element_end - element_begin;
|
|
|
|
query::dataset::RawDataset sub_data;
|
|
std::unique_ptr<uint8_t[]> buf = nullptr;
|
|
std::vector<size_t> offsets;
|
|
if (data_type != DataType::VECTOR_ARRAY) {
|
|
sub_data = query::dataset::RawDataset{
|
|
element_begin, dim, size_per_chunk, chunk_data};
|
|
} else {
|
|
// TODO(SpadeA): For VectorArray(Embedding List), data is
|
|
// discreted stored in FixedVector which means we will copy the
|
|
// data to a contiguous memory buffer. This is inefficient and
|
|
// will be optimized in the future.
|
|
auto vec_ptr = reinterpret_cast<const VectorArray*>(chunk_data);
|
|
auto size = 0;
|
|
for (int i = 0; i < size_per_chunk; ++i) {
|
|
size += vec_ptr[i].byte_size();
|
|
}
|
|
|
|
buf = std::make_unique<uint8_t[]>(size);
|
|
offsets.reserve(size_per_chunk + 1);
|
|
offsets.push_back(0);
|
|
|
|
auto offset = 0;
|
|
auto ptr = buf.get();
|
|
for (int i = 0; i < size_per_chunk; ++i) {
|
|
memcpy(ptr, vec_ptr[i].data(), vec_ptr[i].byte_size());
|
|
ptr += vec_ptr[i].byte_size();
|
|
|
|
offset += vec_ptr[i].length();
|
|
offsets.push_back(offset);
|
|
}
|
|
sub_data = query::dataset::RawDataset{element_begin,
|
|
dim,
|
|
size_per_chunk,
|
|
buf.get(),
|
|
offsets.data()};
|
|
}
|
|
|
|
if (data_type == DataType::VECTOR_ARRAY) {
|
|
AssertInfo(
|
|
query_lims != nullptr,
|
|
"query_lims is nullptr, but data_type is vector array");
|
|
}
|
|
|
|
if (milvus::exec::UseVectorIterator(info)) {
|
|
AssertInfo(data_type != DataType::VECTOR_ARRAY,
|
|
"vector array(embedding list) is not supported for "
|
|
"vector iterator");
|
|
|
|
auto sub_qr =
|
|
PackBruteForceSearchIteratorsIntoSubResult(search_dataset,
|
|
sub_data,
|
|
info,
|
|
index_info,
|
|
bitset,
|
|
data_type);
|
|
final_qr.merge(sub_qr);
|
|
} else {
|
|
auto sub_qr = BruteForceSearch(search_dataset,
|
|
sub_data,
|
|
info,
|
|
index_info,
|
|
bitset,
|
|
data_type,
|
|
element_type,
|
|
op_context);
|
|
final_qr.merge(sub_qr);
|
|
}
|
|
}
|
|
if (milvus::exec::UseVectorIterator(info)) {
|
|
std::vector<int64_t> chunk_rows(max_chunk, 0);
|
|
for (int i = 1; i < max_chunk; ++i) {
|
|
chunk_rows[i] = i * vec_size_per_chunk;
|
|
}
|
|
search_result.AssembleChunkVectorIterators(
|
|
num_queries, max_chunk, chunk_rows, final_qr.chunk_iterators());
|
|
} else {
|
|
search_result.distances_ = std::move(final_qr.mutable_distances());
|
|
search_result.seg_offsets_ =
|
|
std::move(final_qr.mutable_seg_offsets());
|
|
}
|
|
search_result.unity_topK_ = topk;
|
|
search_result.total_nq_ = num_queries;
|
|
}
|
|
}
|
|
|
|
} // namespace milvus::query
|