From e7ad7cc31e4cd6a6702adc9678dc3f0901c041ee Mon Sep 17 00:00:00 2001 From: lichen2015 Date: Thu, 12 Feb 2026 10:38:05 +0800 Subject: [PATCH] fix: combined indexer should use key instead of index (#87) Co-authored-by: yinzefeng.yzf --- src/core/algorithm/flat/flat_streamer.cc | 6 +- .../algorithm/flat/flat_streamer_context.h | 2 +- .../combined_vector_column_indexer.cc | 88 +++++++++++++------ .../combined_vector_column_indexer.h | 1 + .../core/algorithm/flat/flat_streamer_test.cc | 4 +- tests/db/index/segment/segment_test.cc | 60 +++++++++++++ 6 files changed, 128 insertions(+), 33 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index a721cf5..8969efc 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -376,7 +376,7 @@ int FlatStreamer::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::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::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); diff --git a/src/core/algorithm/flat/flat_streamer_context.h b/src/core/algorithm/flat/flat_streamer_context.h index 24cfd9e..1626880 100644 --- a/src/core/algorithm/flat/flat_streamer_context.h +++ b/src/core/algorithm/flat/flat_streamer_context.h @@ -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); } } } diff --git a/src/db/index/column/vector_column/combined_vector_column_indexer.cc b/src/db/index/column/vector_column/combined_vector_column_indexer.cc index f1385b0..70c71d0 100644 --- a/src/db/index/column/vector_column/combined_vector_column_indexer.cc +++ b/src/db/index/column/vector_column/combined_vector_column_indexer.cc @@ -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 CombinedVectorColumnIndexer::Search( const vector_column_params::VectorData &vector_data, const vector_column_params::QueryParams &query_params) { core::IndexDocumentList doc_list; std::vector reverted_vector_list; std::vector 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> 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(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 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 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(q_params) ->scale_factor(); @@ -71,29 +101,34 @@ Result CombinedVectorColumnIndexer::Search( scale_factor = std::dynamic_pointer_cast(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( - 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( - 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( + query_params.group_by->group_topk, + query_params.group_by->group_count, + query_params.group_by->group_by) + : nullptr, + {}, + need_refine ? std::shared_ptr( + 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 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( diff --git a/src/db/index/column/vector_column/combined_vector_column_indexer.h b/src/db/index/column/vector_column/combined_vector_column_indexer.h index 2e723c1..b0b0589 100644 --- a/src/db/index/column/vector_column/combined_vector_column_indexer.h +++ b/src/db/index/column/vector_column/combined_vector_column_indexer.h @@ -52,6 +52,7 @@ class CombinedVectorColumnIndexer { std::vector indexers_; std::vector normal_indexers_; std::vector blocks_; + std::vector block_offsets_; MetricType metric_type_{MetricType::UNDEFINED}; bool is_quantized_{false}; uint64_t min_doc_id_{0}; diff --git a/tests/core/algorithm/flat/flat_streamer_test.cc b/tests/core/algorithm/flat/flat_streamer_test.cc index 022c106..f03012d 100644 --- a/tests/core/algorithm/flat/flat_streamer_test.cc +++ b/tests/core/algorithm/flat/flat_streamer_test.cc @@ -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(); diff --git a/tests/db/index/segment/segment_test.cc b/tests/db/index/segment/segment_test.cc index 5db3f0b..6ca6fff 100644 --- a/tests/db/index/segment/segment_test.cc +++ b/tests/db/index/segment/segment_test.cc @@ -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(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> 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>("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(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(results.value().get()); + ASSERT_TRUE(vector_results); + ASSERT_EQ(vector_results->count(), 10); + + int count = 0; + std::vector 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(