xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: combined indexer should use key instead of index (#87) 

Co-authored-by: yinzefeng.yzf <yinzefeng.yzf@alibaba-inc.com>
This commit is contained in:
lichen2015 2026-02-12 10:38:05 +08:00 committed by GitHub
parent dc088a4446
commit e7ad7cc31e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 128 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/core/algorithm/flat/flat_streamer.cc
View File

@ -376,7 +376,7 @@ int FlatStreamer<BATCH_SIZE>::search_bf_by_p_keys_impl(
if (!filter.is_valid() || !filter(key)) {
dist_t dist = 0;
IndexStorage::MemoryBlock block;
entity_.get_vector_by_key(key, block);
if (entity_.get_vector_by_key(key, block) != 0) continue;
entity_.row_major_distance(query, block.data(), 1, &dist);
heap->emplace(key, dist);
}
@ -418,7 +418,7 @@ int FlatStreamer<BATCH_SIZE>::group_by_search_impl(
if (!bf_context->filter().is_valid() || !bf_context->filter()(key)) {
dist_t dist = 0;
IndexStorage::MemoryBlock block;
entity_.get_vector_by_key(key, block);
if (entity_.get_vector_by_key(key, block) != 0) continue;
entity_.row_major_distance(query, block.data(), 1, &dist);
std::string group_id = group_by(key);
@ -466,7 +466,7 @@ int FlatStreamer<BATCH_SIZE>::group_by_search_p_keys_impl(
if (!bf_context->filter().is_valid() || !bf_context->filter()(key)) {
dist_t dist = 0;
IndexStorage::MemoryBlock block;
entity_.get_vector_by_key(key, block);
if (entity_.get_vector_by_key(key, block) != 0) continue;
entity_.row_major_distance(query, block.data(), 1, &dist);
std::string group_id = group_by(key);

2
src/core/algorithm/flat/flat_streamer_context.h
View File

@ -122,7 +122,7 @@ class FlatStreamerContext : public IndexStreamer::Context {
owner_->entity().get_vector_by_key(key, block);
results_[idx].emplace_back(key, score, key, block);
} else {
results_[idx].emplace_back(key, score);
results_[idx].emplace_back(key, score, key);
}
}
}

88
src/db/index/column/vector_column/combined_vector_column_indexer.cc
View File

@ -40,22 +40,53 @@ CombinedVectorColumnIndexer::CombinedVectorColumnIndexer(
}
}
int block_offset = 0;
for (size_t i = 0; i < indexers_.size(); ++i) {
auto &block_meta = blocks_[i];
block_offsets_.push_back(block_offset);
block_offset += block_meta.doc_count_;
}
min_doc_id_ = segment_meta.min_doc_id();
}
Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
const vector_column_params::VectorData &vector_data,
const vector_column_params::QueryParams &query_params) {
core::IndexDocumentList doc_list;
std::vector<std::string> reverted_vector_list;
std::vector<std::string> reverted_sparse_values_list;
int block_offset = 0;
// query_params.bf_pks is segment level, here we need to convert it to block
// level
std::vector<std::vector<uint64_t>> block_bf_pks(indexers_.size());
if (!query_params.bf_pks.empty()) {
// dispatcher pks to corresponding block_bf_pks
for (auto &pk : query_params.bf_pks[0]) {
for (size_t i = 0; i < block_offsets_.size(); ++i) {
if (pk >= block_offsets_[i] &&
pk < block_offsets_[i] + blocks_[i].doc_count_) {
block_bf_pks[i].push_back(
static_cast<uint64_t>(pk - block_offsets_[i]));
break;
}
}
}
}
auto q_params = query_params.query_params;
for (size_t i = 0; i < indexers_.size(); ++i) {
auto &block_meta = blocks_[i];
if (!query_params.bf_pks.empty() && block_bf_pks[i].empty()) {
LOG_DEBUG(
"query_params has bf_pks, but block_bf_pks[%zu] is empty, just skip "
"this indexer",
i);
continue;
}
zvec::Result<zvec::IndexResults::Ptr> result{nullptr};
float scale_factor{};
bool need_refine{false};
if (q_params && q_params->is_using_refiner()) {
if (normal_indexers_.size() != indexers_.size()) {
return tl::make_unexpected(Status::InvalidArgument(
@ -63,7 +94,6 @@ Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
"] not match indexers size[", indexers_.size(), "]"));
}
// query_params of HNSW doesn't have scale_factor
float scale_factor{};
if (q_params->type() == IndexType::FLAT) {
scale_factor = std::dynamic_pointer_cast<FlatQueryParams>(q_params)
->scale_factor();
@ -71,29 +101,34 @@ Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
scale_factor =
std::dynamic_pointer_cast<IVFQueryParams>(q_params)->scale_factor();
}
vector_column_params::QueryParams modified_query_params{
query_params.data_type,
query_params.dimension,
query_params.topk,
query_params.filter,
query_params.fetch_vector,
query_params.query_params,
query_params.group_by
? std::make_unique<vector_column_params::GroupByParams>(
query_params.group_by->group_topk,
query_params.group_by->group_count,
query_params.group_by->group_by)
: nullptr,
query_params.bf_pks,
std::shared_ptr<vector_column_params::RefinerParam>(
new vector_column_params::RefinerParam{scale_factor,
normal_indexers_[i]}),
query_params.extra_params};
result = indexers_[i]->Search(vector_data, modified_query_params);
} else {
result = indexers_[i]->Search(vector_data, query_params);
need_refine = true;
}
vector_column_params::QueryParams modified_query_params{
query_params.data_type,
query_params.dimension,
query_params.topk,
query_params.filter,
query_params.fetch_vector,
query_params.query_params,
query_params.group_by
? std::make_unique<vector_column_params::GroupByParams>(
query_params.group_by->group_topk,
query_params.group_by->group_count,
query_params.group_by->group_by)
: nullptr,
{},
need_refine ? std::shared_ptr<vector_column_params::RefinerParam>(
new vector_column_params::RefinerParam{
scale_factor, normal_indexers_[i]})
: nullptr,
query_params.extra_params};
if (!query_params.bf_pks.empty()) {
modified_query_params.bf_pks.emplace_back(block_bf_pks[i]);
}
result = indexers_[i]->Search(vector_data, modified_query_params);
if (!result) {
return tl::make_unexpected(result.error());
}
@ -105,10 +140,9 @@ Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
const auto &sub_docs = vector_index_results->docs();
for (size_t j = 0; j < sub_docs.size(); ++j) {
auto doc = sub_docs[j];
doc.set_index(block_offset + sub_docs[j].index());
doc.set_key(block_offsets_[i] + sub_docs[j].key());
doc_list.emplace_back(std::move(doc));
}
block_offset += block_meta.doc_count_;
auto &&temp_vector_list = vector_index_results->reverted_vector_list();
reverted_vector_list.insert(

1
src/db/index/column/vector_column/combined_vector_column_indexer.h
View File

@ -52,6 +52,7 @@ class CombinedVectorColumnIndexer {
std::vector<VectorColumnIndexer::Ptr> indexers_;
std::vector<VectorColumnIndexer::Ptr> normal_indexers_;
std::vector<BlockMeta> blocks_;
std::vector<uint32_t> block_offsets_;
MetricType metric_type_{MetricType::UNDEFINED};
bool is_quantized_{false};
uint64_t min_doc_id_{0};

4
tests/core/algorithm/flat/flat_streamer_test.cc
View File

@ -847,8 +847,8 @@ TEST_F(FlatStreamerTest, TestMaxIndexSize) {
writeCnt1 * 128 * 4 + writeCnt1 * 8 + writeCnt1 * 28 / 32;
LOG_INFO("increment1: %lu, expect_size: %lu", increment1, expect_size);
ASSERT_GT(expect_size, increment1 * 0.8f);
ASSERT_LT(expect_size, increment1 * 1.2f);
ASSERT_GT(expect_size, increment1 * 0.75f);
ASSERT_LT(expect_size, increment1 * 1.25f);
streamer->flush(0UL);
streamer.reset();

60
tests/db/index/segment/segment_test.cc
View File

@ -1170,6 +1170,66 @@ TEST_P(SegmentTest, CombinedVectorColumnIndexerWithQuantVectorIndex) {
ASSERT_EQ(count, 10);
}
TEST_P(SegmentTest, CombinedVectorColumnIndexerQueryWithPks) {
options.max_buffer_size_ = 10 * 1024;
auto tmp_schema = test::TestHelper::CreateSchemaWithVectorIndex(
false, "demo", std::make_shared<HnswIndexParams>(MetricType::IP));
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path, *tmp_schema, 0, 0, id_map, delete_store, version_manager,
options, 0, 0);
ASSERT_TRUE(segment != nullptr);
uint64_t MAX_DOC = 1000;
test::TestHelper::SegmentInsertDoc(segment, *schema, 0, MAX_DOC);
auto combined_indexer = segment->get_combined_vector_indexer("dense_fp32");
ASSERT_TRUE(combined_indexer != nullptr);
Doc verify_doc = test::TestHelper::CreateDoc(999, *schema);
std::vector<std::vector<uint64_t>> bf_pks = {
{10, 20, 30, 40, 50, 60, 70, 80, 90, 999}};
// query
auto dense_fp32_field = schema->get_field("dense_fp32");
auto query_vector = verify_doc.get<std::vector<float>>("dense_fp32").value();
auto query = vector_column_params::VectorData{
vector_column_params::DenseVector{.data = query_vector.data()}};
auto query_params = vector_column_params::QueryParams{
.data_type = dense_fp32_field->data_type(),
.dimension = dense_fp32_field->dimension(),
.topk = 10,
.filter = nullptr,
.fetch_vector = false,
.query_params = std::make_shared<zvec::QueryParams>(IndexType::HNSW),
.group_by = nullptr,
.bf_pks = bf_pks,
.refiner_param = nullptr,
.extra_params = {}};
auto results = combined_indexer->Search(query, query_params);
ASSERT_TRUE(results.has_value());
auto vector_results =
dynamic_cast<VectorIndexResults *>(results.value().get());
ASSERT_TRUE(vector_results);
ASSERT_EQ(vector_results->count(), 10);
int count = 0;
std::vector<uint64_t> result_doc_ids;
auto iter = vector_results->create_iterator();
while (iter->valid()) {
count++;
result_doc_ids.push_back(iter->doc_id());
iter->next();
}
ASSERT_EQ(count, 10);
// need reverse result_doc_ids
std::reverse(result_doc_ids.begin(), result_doc_ids.end());
ASSERT_EQ(result_doc_ids, bf_pks[0]);
}
TEST_P(SegmentTest, ConcurrentInsertOperations) {
auto segment = test::TestHelper::CreateSegmentWithDoc(