mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
Spade AandGitHub
08cdd63f68
issue: https://github.com/milvus-io/milvus/issues/42148 Correctness: - JSON/CSV import: reject struct array elements with mismatched field count - Parquet import: error on type assertion failure in scalar and vector paths instead of silent zero-fill/data loss - JSON import: validate per-vector dimension in ArrayOfVector FloatVector path - Move element_level inference from pymilvus into C++ ParsePlaceholderGroup: pymilvus cannot reliably infer if it's elelment_level as which kinds of search on ArrayOfVector are supported are determined by metric type but pymilvus does not have this info. - Fix element-level search returning wrong row IDs on growing segments with multiple chunks by using cumulative element offset instead of row offset as begin_id in brute-force search. https://github.com/milvus-io/milvus/issues/48617 Performance: - Replace proto.Clone with in-place FieldName mutation in reconstructStructFieldData Nested index: - field name is missing in CreateIndexInfo which is needed to determine whether the index should be nested or not Tests: added for all above fixes --------- Signed-off-by: SpadeA <tangchenjie1210@gmail.com>
373 lines
16 KiB
C++
373 lines
16 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 <string.h>
|
|
#include <algorithm>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <shared_mutex>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "SearchOnGrowing.h"
|
|
#include "cachinglayer/CacheSlot.h"
|
|
#include "common/BitsetView.h"
|
|
#include "common/Consts.h"
|
|
#include "common/EasyAssert.h"
|
|
#include "common/FieldMeta.h"
|
|
#include "common/IndexMeta.h"
|
|
#include "common/OffsetMapping.h"
|
|
#include "common/QueryInfo.h"
|
|
#include "common/QueryResult.h"
|
|
#include "common/Schema.h"
|
|
#include "common/Types.h"
|
|
#include "common/Utils.h"
|
|
#include "common/VectorArray.h"
|
|
#include "common/protobuf_utils.h"
|
|
#include "exec/operator/Utils.h"
|
|
#include "index/Index.h"
|
|
#include "index/VectorIndex.h"
|
|
#include "knowhere/comp/index_param.h"
|
|
#include "query/CachedSearchIterator.h"
|
|
#include "query/SearchBruteForce.h"
|
|
#include "query/SearchOnIndex.h"
|
|
#include "query/SubSearchResult.h"
|
|
#include "query/Utils.h"
|
|
#include "query/helper.h"
|
|
#include "segcore/ConcurrentVector.h"
|
|
#include "segcore/FieldIndexing.h"
|
|
#include "segcore/InsertRecord.h"
|
|
#include "segcore/SegmentGrowingImpl.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 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_offsets,
|
|
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();
|
|
|
|
// 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_offsets};
|
|
int32_t current_chunk_id = 0;
|
|
|
|
// get index params for bm25 and minhash brute force
|
|
std::map<std::string, std::string> index_info;
|
|
if (metric_type == knowhere::metric::BM25 ||
|
|
metric_type == knowhere::metric::MHJACCARD) {
|
|
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);
|
|
const auto& offset_mapping = vec_ptr->get_offset_mapping();
|
|
|
|
TargetBitmap transformed_bitset;
|
|
BitsetView search_bitset = bitset;
|
|
if (offset_mapping.IsEnabled()) {
|
|
transformed_bitset = TransformBitset(bitset, offset_mapping);
|
|
search_bitset = BitsetView(transformed_bitset);
|
|
}
|
|
|
|
auto active_count = offset_mapping.IsEnabled()
|
|
? offset_mapping.GetValidCount()
|
|
: std::min(int64_t(bitset.size()),
|
|
segment.get_active_count(timestamp));
|
|
|
|
// Check for nullable vector field with all null values
|
|
if (active_count == 0) {
|
|
// All vectors are null, return empty result
|
|
auto total_num = num_queries * info.topk_;
|
|
search_result.seg_offsets_.resize(total_num, INVALID_SEG_OFFSET);
|
|
search_result.distances_.resize(total_num, 0.0f);
|
|
search_result.total_nq_ = num_queries;
|
|
search_result.unity_topK_ = info.topk_;
|
|
return;
|
|
}
|
|
|
|
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,
|
|
search_bitset,
|
|
data_type);
|
|
cached_iter.NextBatch(info, search_result);
|
|
if (offset_mapping.IsEnabled()) {
|
|
TransformOffset(search_result.seg_offsets_, offset_mapping);
|
|
}
|
|
return;
|
|
}
|
|
|
|
auto vec_size_per_chunk = vec_ptr->get_size_per_chunk();
|
|
auto max_chunk = upper_div(active_count, vec_size_per_chunk);
|
|
|
|
// embedding search embedding on embedding list
|
|
bool embedding_search = false;
|
|
if (data_type == DataType::VECTOR_ARRAY &&
|
|
info.array_offsets_ != nullptr) {
|
|
embedding_search = true;
|
|
}
|
|
|
|
// Track cumulative element offset for embedding search.
|
|
// For embedding_search, begin_id must be the cumulative element
|
|
// count (not row offset), because ArrayOffsets maps global
|
|
// element IDs to row IDs.
|
|
int64_t cumulative_element_offset = 0;
|
|
|
|
for (int chunk_id = current_chunk_id; chunk_id < max_chunk;
|
|
++chunk_id) {
|
|
auto chunk_data = vec_ptr->get_chunk_data(chunk_id);
|
|
|
|
auto row_begin = chunk_id * vec_size_per_chunk;
|
|
auto row_end =
|
|
std::min(active_count, (chunk_id + 1) * vec_size_per_chunk);
|
|
auto size_per_chunk = row_end - row_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{
|
|
row_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);
|
|
|
|
if (embedding_search) {
|
|
auto count = 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();
|
|
count += vec_ptr[i].length();
|
|
}
|
|
sub_data = query::dataset::RawDataset{
|
|
cumulative_element_offset, dim, count, buf.get()};
|
|
cumulative_element_offset += count;
|
|
} else {
|
|
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{row_begin,
|
|
dim,
|
|
size_per_chunk,
|
|
buf.get(),
|
|
offsets.data()};
|
|
}
|
|
}
|
|
|
|
auto vector_type = data_type;
|
|
if (embedding_search) {
|
|
vector_type = element_type;
|
|
}
|
|
|
|
if (milvus::exec::UseVectorIterator(info)) {
|
|
AssertInfo(vector_type != DataType::VECTOR_ARRAY,
|
|
"vector array(embedding list) is not supported for "
|
|
"vector iterator");
|
|
|
|
if (buf != nullptr) {
|
|
search_result.chunk_buffers_.emplace_back(std::move(buf));
|
|
}
|
|
|
|
auto sub_qr =
|
|
PackBruteForceSearchIteratorsIntoSubResult(search_dataset,
|
|
sub_data,
|
|
info,
|
|
index_info,
|
|
search_bitset,
|
|
vector_type);
|
|
final_qr.merge(sub_qr);
|
|
} else {
|
|
auto sub_qr = BruteForceSearch(search_dataset,
|
|
sub_data,
|
|
info,
|
|
index_info,
|
|
search_bitset,
|
|
vector_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;
|
|
}
|
|
bool larger_is_closer = PositivelyRelated(info.metric_type_);
|
|
search_result.AssembleChunkVectorIterators(
|
|
num_queries,
|
|
max_chunk,
|
|
chunk_rows,
|
|
final_qr.chunk_iterators(),
|
|
offset_mapping,
|
|
larger_is_closer);
|
|
} else {
|
|
if (info.array_offsets_ != nullptr) {
|
|
auto [seg_offsets, elem_indicies] =
|
|
final_qr.convert_to_element_offsets(
|
|
info.array_offsets_.get());
|
|
search_result.seg_offsets_ = std::move(seg_offsets);
|
|
search_result.element_indices_ = std::move(elem_indicies);
|
|
search_result.element_level_ = true;
|
|
} else {
|
|
search_result.seg_offsets_ =
|
|
std::move(final_qr.mutable_offsets());
|
|
}
|
|
search_result.distances_ = std::move(final_qr.mutable_distances());
|
|
if (offset_mapping.IsEnabled()) {
|
|
TransformOffset(search_result.seg_offsets_, offset_mapping);
|
|
}
|
|
}
|
|
search_result.unity_topK_ = topk;
|
|
search_result.total_nq_ = num_queries;
|
|
}
|
|
}
|
|
|
|
} // namespace milvus::query
|