-
Notifications
You must be signed in to change notification settings - Fork 104
feat(pyramid): integrate reasoning mechanism for Pyramid index #2333
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,6 +16,7 @@ | |
|
|
||
| set (PYRAMID_SRCS | ||
| pyramid.cpp | ||
| pyramid_index_node.cpp | ||
| pyramid_zparameters.cpp | ||
| ) | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -26,6 +26,8 @@ | |
| #include "impl/heap/standard_heap.h" | ||
| #include "impl/odescent/odescent_graph_builder.h" | ||
| #include "impl/pruning_strategy.h" | ||
| #include "impl/reasoning/search_reasoning.h" | ||
| #include "io/memory_io/memory_io_parameter.h" | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [note] The include |
||
| #include "quantization/rabitq_quantization/rabitq_quantizer_parameter.h" | ||
| #include "query_context.h" | ||
| #include "storage/empty_index_binary_set.h" | ||
|
|
@@ -133,17 +135,6 @@ split(const std::string& str, char delimiter) { | |
| return vec; | ||
| } | ||
|
|
||
| static inline uint64_t | ||
| get_suitable_max_degree(int64_t data_num) { | ||
| if (data_num < 100'000) { | ||
| return 24; | ||
| } | ||
| if (data_num < 1000'000) { | ||
| return 32; | ||
| } | ||
| return 64; | ||
| } | ||
|
|
||
| static inline uint64_t | ||
| get_suitable_ef_search(int64_t topk, int64_t data_num, uint64_t subindex_ef_search = 50) { | ||
| auto topk_float = static_cast<float>(topk); | ||
|
|
@@ -159,149 +150,37 @@ get_suitable_ef_search(int64_t topk, int64_t data_num, uint64_t subindex_ef_sear | |
| return std::max(static_cast<uint64_t>(4.0F * topk_float), subindex_ef_search * 8); | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] Consider adding these fields to search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_; |
||
| } | ||
|
|
||
| IndexNode::IndexNode(Allocator* allocator, | ||
| GraphInterfaceParamPtr graph_param, | ||
| uint32_t index_min_size) | ||
| : ids_(allocator), | ||
| children_(allocator), | ||
| allocator_(allocator), | ||
| graph_param_(std::move(graph_param)), | ||
| index_min_size_(index_min_size) { | ||
| } | ||
|
|
||
| void | ||
| IndexNode::Build(ODescent& odescent) { | ||
| std::unique_lock lock(mutex_); | ||
| // Build an index when the level corresponding to the current node requires indexing | ||
| if (not ids_.empty()) { | ||
| Init(); | ||
| } | ||
| if (status_ == Status::GRAPH) { | ||
| entry_point_ = ids_[0]; | ||
| odescent.SetMaxDegree(static_cast<int32_t>(graph_param_->max_degree_)); | ||
| odescent.Build(ids_); | ||
| odescent.SaveGraph(graph_); | ||
| Vector<InnerIdType>(allocator_).swap(ids_); | ||
| } | ||
| for (const auto& item : children_) { | ||
| item.second->Build(odescent); | ||
| } | ||
| } | ||
|
|
||
| void | ||
| IndexNode::AddChild(const std::string& key) { | ||
| // AddChild is not thread-safe; ensure thread safety in calls to it. | ||
| children_[key] = std::make_unique<IndexNode>(allocator_, graph_param_, index_min_size_); | ||
| children_[key]->level_ = level_ + 1; | ||
| } | ||
|
|
||
| IndexNode* | ||
| IndexNode::GetChild(const std::string& key, bool need_init) { | ||
| std::unique_lock lock(mutex_); | ||
| auto result = children_.find(key); | ||
| if (result != children_.end()) { | ||
| return result->second.get(); | ||
| } | ||
| if (not need_init) { | ||
| return nullptr; | ||
| } | ||
| AddChild(key); | ||
| return children_[key].get(); | ||
| } | ||
|
|
||
| void | ||
| IndexNode::Deserialize(StreamReader& reader) { | ||
| // deserialize `entry_point_` | ||
| StreamReader::ReadObj(reader, entry_point_); | ||
| // deserialize `level_` | ||
| StreamReader::ReadObj(reader, level_); | ||
| // deserialize `status_` | ||
| StreamReader::ReadObj(reader, status_); | ||
| if (status_ == Status::GRAPH) { | ||
| graph_ = std::make_shared<SparseGraphDataCell>( | ||
| std::dynamic_pointer_cast<SparseGraphDatacellParameter>(graph_param_), allocator_); | ||
| graph_->Deserialize(reader); | ||
| } else if (status_ == Status::FLAT) { | ||
| StreamReader::ReadVector(reader, ids_); | ||
| } | ||
| // deserialize `children` | ||
| uint64_t children_size = 0; | ||
| StreamReader::ReadObj(reader, children_size); | ||
| for (uint64_t i = 0; i < children_size; ++i) { | ||
| std::string key = StreamReader::ReadString(reader); | ||
| AddChild(key); | ||
| children_[key]->Deserialize(reader); | ||
| } | ||
| } | ||
| InnerSearchParam | ||
| Pyramid::create_knn_search_param(const PyramidSearchParameters& parsed_param, | ||
| int64_t k, | ||
| const FilterPtr& filter) const { | ||
| CHECK_ARGUMENT(k > 0, fmt::format("k({}) must be greater than 0", k)); | ||
| CHECK_ARGUMENT(parsed_param.hierarchy_op == PyramidSearchParameters::HierarchyOp::SINGLE, | ||
| "multi-hierarchy search (union/intersection) is not yet implemented"); | ||
| auto ef_search_threshold = | ||
| std::max<uint64_t>(AMPLIFICATION_FACTOR * k, static_cast<uint64_t>(1000)); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P3] Avoid signed overflow in the ef bound Both operands of |
||
| CHECK_ARGUMENT( // NOLINT | ||
| (1 <= parsed_param.ef_search) and (parsed_param.ef_search <= ef_search_threshold), | ||
| fmt::format( | ||
| "ef_search({}) must be in range [1, {}]", parsed_param.ef_search, ef_search_threshold)); | ||
|
|
||
| void | ||
| IndexNode::Serialize(StreamWriter& writer) const { | ||
| // serialize `entry_point_` | ||
| StreamWriter::WriteObj(writer, entry_point_); | ||
| // serialize `level_` | ||
| StreamWriter::WriteObj(writer, level_); | ||
| // serialize `status_` | ||
| StreamWriter::WriteObj(writer, status_); | ||
| if (status_ == Status::GRAPH) { | ||
| graph_->Serialize(writer); | ||
| } else if (status_ == Status::FLAT) { | ||
| StreamWriter::WriteVector(writer, ids_); | ||
| } | ||
| // serialize `children` | ||
| uint64_t children_size = children_.size(); | ||
| StreamWriter::WriteObj(writer, children_size); | ||
| for (const auto& item : children_) { | ||
| // calculate size of `key` | ||
| StreamWriter::WriteString(writer, item.first); | ||
| // calculate size of `content` | ||
| item.second->Serialize(writer); | ||
| } | ||
| } | ||
| void | ||
| IndexNode::Init() { | ||
| if (status_ == Status::NO_INDEX) { | ||
| if (ids_.size() >= index_min_size_) { | ||
| if (not ids_.empty() and level_ != 0) { | ||
| auto new_max_degree = get_suitable_max_degree(static_cast<int64_t>(ids_.size())); | ||
| if (new_max_degree < graph_param_->max_degree_) { | ||
| auto new_graph_param = std::make_shared<SparseGraphDatacellParameter>(); | ||
| new_graph_param->FromJson(graph_param_->ToJson()); | ||
| new_graph_param->max_degree_ = | ||
| get_suitable_max_degree(static_cast<int64_t>(ids_.size())); | ||
| graph_param_ = new_graph_param; | ||
| } | ||
| } | ||
| graph_ = std::make_shared<SparseGraphDataCell>( | ||
| std::dynamic_pointer_cast<SparseGraphDatacellParameter>(graph_param_), allocator_); | ||
| status_ = Status::GRAPH; | ||
| } else { | ||
| status_ = Status::FLAT; | ||
| } | ||
| InnerSearchParam search_param; | ||
| search_param.ef = std::max<uint64_t>(parsed_param.ef_search, static_cast<uint64_t>(k)); | ||
| search_param.radius = std::numeric_limits<float>::max(); | ||
| search_param.topk = k; | ||
| search_param.search_mode = KNN_SEARCH; | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The Consider adding these fields to |
||
| search_param.parallel_search_thread_count = parsed_param.parallel_search_thread_count; | ||
| if (this->support_duplicate_) { | ||
| search_param.consider_duplicate = true; | ||
| } | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical] In search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_;But Suggested fix: add the same two lines to search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_; |
||
| } | ||
|
|
||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical] Add after search_param.enable_reorder = use_reorder_;
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical]
search_param.distance_threshold = threshold;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_;These should be added to |
||
| void | ||
| IndexNode::Search(const SearchFunc& search_func, | ||
| const VisitedListPtr& vl, | ||
| const DistHeapPtr& search_result, | ||
| uint64_t ef_search) const { | ||
| bool has_index = false; | ||
| { | ||
| std::shared_lock lock(mutex_); | ||
| has_index = status_ != IndexNode::Status::NO_INDEX; | ||
| } | ||
| if (has_index) { | ||
| auto self_search_result = search_func(this, vl); | ||
| search_result->Merge(*self_search_result); | ||
| while (search_result->Size() > ef_search) { | ||
| search_result->Pop(); | ||
| } | ||
| return; | ||
| if (parsed_param.enable_time_record) { | ||
| search_param.time_cost = std::make_shared<Timer>(); | ||
| search_param.time_cost->SetThreshold(parsed_param.timeout_ms); | ||
| } | ||
|
|
||
| for (const auto& [key, node] : children_) { | ||
| node->Search(search_func, vl, search_result, ef_search); | ||
| } | ||
| search_param.is_inner_id_allowed = this->create_search_filter(filter); | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] Consider adding: search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_;before the |
||
| return search_param; | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [note] Currently this has no behavioral impact because Consider adding |
||
| } | ||
|
|
||
| std::vector<int64_t> | ||
|
|
@@ -312,6 +191,10 @@ Pyramid::build_by_odescent(const DatasetPtr& base) { | |
|
|
||
| resize(data_num); | ||
| std::memcpy(label_table_->label_table_.data(), data_ids, sizeof(LabelType) * data_num); | ||
| label_table_->ResetRemap(data_num); | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [note] The |
||
| for (InnerIdType id = 0; id < static_cast<InnerIdType>(data_num); ++id) { | ||
| label_table_->InsertRemap(data_ids[id], id); | ||
| } | ||
|
|
||
| base_codes_->BatchInsertVector(data_vectors, data_num); | ||
| if (has_precise_reorder()) { | ||
|
|
@@ -505,6 +388,126 @@ Pyramid::RangeSearch(const DatasetPtr& query, | |
| return result; | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] Suggested fix: add the same search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_; |
||
| } | ||
|
|
||
| DatasetPtr | ||
| Pyramid::SearchWithRequest(const SearchRequest& request) const { | ||
| SearchStatistics stats; | ||
| QueryContext ctx{.alloc = this->allocator_, .stats = &stats}; | ||
| if (request.search_allocator_ != nullptr) { | ||
| ctx.alloc = request.search_allocator_; | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Propagate the request allocator through core search allocations Assigning the allocator only to |
||
| } | ||
|
LHT129 marked this conversation as resolved.
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] Both ctx.rabitq_error_rate = parsed_param.rabitq_error_rate;This is missing from |
||
|
|
||
| const auto& query = request.query_; | ||
| CHECK_ARGUMENT(query != nullptr, "query dataset is required"); | ||
| CHECK_ARGUMENT(query->GetFloat32Vectors() != nullptr, "query vectors is required"); | ||
|
|
||
| const bool is_knn = request.mode_ == SearchMode::KNN_SEARCH; | ||
| if (is_knn) { | ||
| this->validate_knn_args(query, request.topk_); | ||
| } else { | ||
| CHECK_ARGUMENT(request.mode_ == SearchMode::RANGE_SEARCH, "unsupported search mode"); | ||
| this->validate_range_args(query, request.radius_, request.limited_size_); | ||
| } | ||
|
|
||
| auto parsed_param = PyramidSearchParameters::FromJson(request.params_str_); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Preserve Pyramid's RaBitQ search parameters Unlike both
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] Add after ctx.rabitq_error_rate = parsed_param.rabitq_error_rate; |
||
| InnerSearchParam search_param; | ||
| if (is_knn) { | ||
| search_param = this->create_knn_search_param(parsed_param, request.topk_, request.filter_); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Honor the KNN distance threshold This path never copies |
||
| } else { | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical] ctx.rabitq_error_rate = parsed_param.rabitq_error_rate;Without this, RaBitQ quantization error rate control will not work when searching via |
||
| CHECK_ARGUMENT(parsed_param.hierarchy_op == PyramidSearchParameters::HierarchyOp::SINGLE, | ||
| "multi-hierarchy search (union/intersection) is not yet implemented"); | ||
| search_param.ef = parsed_param.ef_search; | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The range-search parameter setup in Suggested helper signature:
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The |
||
| search_param.radius = request.radius_ * RADIUS_EPSILON; | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical]
|
||
| search_param.search_mode = RANGE_SEARCH; | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The range search branch in Add after search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_;
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The range search path in Consider adding: search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_;after |
||
| search_param.parallel_search_thread_count = parsed_param.parallel_search_thread_count; | ||
| search_param.topk = request.limited_size_ == -1 ? std::numeric_limits<int64_t>::max() | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical] The range search path in SearchWithRequest is missing search_param.enable_reorder and search_param.enable_rabitq_one_bit_search, which are set in the original RangeSearch method. This means reorder and RaBitQ one-bit search will not work when range search is invoked through SearchWithRequest. Compare with RangeSearch: Suggested fix: add these two lines to the range search branch in SearchWithRequest (after search_param.search_mode = RANGE_SEARCH). |
||
| : request.limited_size_; | ||
| if (this->support_duplicate_) { | ||
| search_param.consider_duplicate = true; | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical] The range search branch in In search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_;(Note: However, the range search path in Suggested fix: add these fields in the range search branch (around line 425): search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_; |
||
| } | ||
| if (parsed_param.enable_time_record) { | ||
| search_param.time_cost = std::make_shared<Timer>(); | ||
| search_param.time_cost->SetThreshold(parsed_param.timeout_ms); | ||
| } | ||
| search_param.is_inner_id_allowed = this->create_search_filter(request.filter_); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P1] Compose all enabled request filters The request path constructs its filter solely from
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The range search path in SearchWithRequest does not set search_param.enable_rabitq_one_bit_search, unlike RangeSearch which sets it. This means range searches via SearchWithRequest will not benefit from the rabitq one-bit search optimization. Consider adding: search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search ? parsed_param.rabitq_one_bit_search : default_rabitq_one_bit_search_; |
||
| } | ||
| SearchFunc search_func = [&](const IndexNode* node, const VisitedListPtr& vl) { | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [note] The |
||
| return this->search_node( | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Instrument the GRAPH entry point The GRAPH branch delegates reasoning events to BasicSearcher's flatten overload, but that overload scores and filters its entry point without calling |
||
| node, vl, search_param, query, base_codes_, ctx, parsed_param.subindex_ef_search); | ||
|
LHT129 marked this conversation as resolved.
|
||
| }; | ||
|
|
||
| // Setup reasoning context if expected labels are provided. | ||
| std::shared_ptr<ReasoningContext> reasoning_ctx; | ||
| if (is_knn && not request.expected_labels_.empty()) { | ||
| reasoning_ctx = std::make_shared<ReasoningContext>(ctx.alloc); | ||
| reasoning_ctx->SetSearchParams( | ||
| request.topk_, "Pyramid", use_reorder_, request.filter_ != nullptr); | ||
|
|
||
| UnorderedMap<int64_t, InnerIdType> label_to_inner_id(ctx.alloc); | ||
| Vector<InnerIdType> expected_inner_ids(ctx.alloc); | ||
| { | ||
| // Add holds this lock while mutating labels and vector storage. | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The range search branch in Suggested fix: add the same initialization in the range search branch: search_param.enable_rabitq_one_bit_search = parsed_param.has_rabitq_one_bit_search
? parsed_param.rabitq_one_bit_search
: default_rabitq_one_bit_search_; |
||
| std::lock_guard lock(cur_element_count_mutex_); | ||
| for (const auto& label : request.expected_labels_) { | ||
| // `true` = return_even_removed: include removed labels so reasoning can diagnose them. | ||
| auto [success, inner_id] = label_table_->TryGetIdByLabel(label, true); | ||
| if (success) { | ||
| label_to_inner_id[label] = inner_id; | ||
| } | ||
| } | ||
| expected_inner_ids.reserve(label_to_inner_id.size()); | ||
| for (const auto& [label, inner_id] : label_to_inner_id) { | ||
| expected_inner_ids.push_back(inner_id); | ||
| } | ||
| if (not expected_inner_ids.empty()) { | ||
| auto precise_flatten = | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [critical] The The lock should only protect the label-to-inner-id mapping. The distance computation should be moved outside the lock scope. Suggested fix: collect {
std::lock_guard lock(cur_element_count_mutex_);
for (const auto& label : request.expected_labels_) {
auto [success, inner_id] = label_table_->TryGetIdByLabel(label, true);
if (success) {
label_to_inner_id[label] = inner_id;
}
}
expected_inner_ids.reserve(label_to_inner_id.size());
for (const auto& [label, inner_id] : label_to_inner_id) {
expected_inner_ids.push_back(inner_id);
}
} // release lock here
if (not expected_inner_ids.empty()) {
auto precise_flatten = ...;
auto computer = precise_flatten->FactoryComputer(query->GetFloat32Vectors());
Vector<float> true_dists(expected_inner_ids.size(), ctx.alloc);
precise_flatten->Query(true_dists.data(), computer, ...);
for (size_t i = 0; i < expected_inner_ids.size(); ++i) {
reasoning_ctx->SetTrueDistance(expected_inner_ids[i], true_dists[i]);
}
} |
||
| this->precise_codes_ != nullptr ? this->precise_codes_ : this->base_codes_; | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Do not use lossy search codes as true distances When reorder is disabled, this selects
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Do not use lossy base codes as true distances With reorder disabled, this selects |
||
| auto computer = precise_flatten->FactoryComputer(query->GetFloat32Vectors()); | ||
| Vector<float> true_dists(expected_inner_ids.size(), ctx.alloc); | ||
| precise_flatten->Query(true_dists.data(), | ||
| computer, | ||
| expected_inner_ids.data(), | ||
| static_cast<InnerIdType>(expected_inner_ids.size()), | ||
| &ctx); | ||
| for (size_t i = 0; i < expected_inner_ids.size(); ++i) { | ||
| reasoning_ctx->SetTrueDistance(expected_inner_ids[i], true_dists[i]); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Initialize target traces before setting distances
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P2] Initialize expected traces before setting distances
|
||
| } | ||
| } | ||
| } | ||
|
|
||
| Vector<int64_t> expected_labels_vec( | ||
| request.expected_labels_.begin(), request.expected_labels_.end(), ctx.alloc); | ||
| reasoning_ctx->InitializeExpectedTargets(expected_labels_vec, label_to_inner_id); | ||
| ctx.reasoning_ctx = reasoning_ctx.get(); | ||
| } | ||
|
|
||
| std::string hierarchy_name = | ||
| parsed_param.hierarchies.empty() ? "" : parsed_param.hierarchies[0]; | ||
| auto result = this->search_impl(query, search_func, search_param, ctx, hierarchy_name, nullptr); | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] Both If |
||
| result->Statistics(stats.Dump()); | ||
|
|
||
| if (reasoning_ctx) { | ||
| Vector<InnerIdType> result_inner_ids(ctx.alloc); | ||
| const auto* result_ids = result->GetIds(); | ||
| const auto num_results = result->GetNumElements(); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P1] Use the result dimension as the hit count
|
||
| result_inner_ids.reserve(static_cast<size_t>(num_results)); | ||
| { | ||
| std::lock_guard lock(cur_element_count_mutex_); | ||
| for (int64_t i = 0; i < num_results; ++i) { | ||
| // `true` = return_even_removed: match the same behavior used for expected_labels. | ||
| auto [success, inner_id] = label_table_->TryGetIdByLabel(result_ids[i], true); | ||
| if (success) { | ||
| result_inner_ids.push_back(inner_id); | ||
| } | ||
| } | ||
| } | ||
| reasoning_ctx->MarkResult(result_inner_ids); | ||
| reasoning_ctx->DiagnoseExpectedTargets(); | ||
| result->Reasoning(reasoning_ctx->GenerateReport()); | ||
| } | ||
|
|
||
| return result; | ||
| } | ||
|
|
||
| DatasetPtr | ||
| Pyramid::search_impl(const DatasetPtr& query, | ||
| const SearchFunc& search_func, | ||
|
|
@@ -1522,6 +1525,8 @@ Pyramid::search_node(const IndexNode* node, | |
| for (uint64_t i = 0; i < id_count; ++i) { | ||
| if (inner_filter->CheckValid(ids_ptr[i])) { | ||
| valid_ids.push_back(ids_ptr[i]); | ||
| } else if (ctx.reasoning_ctx != nullptr) { | ||
| ctx.reasoning_ctx->RecordFilterReject(ids_ptr[i]); | ||
| } | ||
| } | ||
| ids_ptr = valid_ids.data(); | ||
|
|
@@ -1533,6 +1538,9 @@ Pyramid::search_node(const IndexNode* node, | |
| codes->Query(dists.data(), computer, ids_ptr, id_count, &ctx); | ||
|
|
||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The reasoning hooks in If hop tracking is feasible to add in a follow-up, consider threading the current hop count through the search recursion so the reasoning report can provide richer diagnostics. |
||
| for (int i = 0; i < id_count; ++i) { | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] The loop variable |
||
| if (ctx.reasoning_ctx != nullptr) { | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [note] The reasoning tracking (RecordVisit, RecordFilterReject, RecordEviction) is only added in the FLAT branch of search_node. The GRAPH branch (line 1558) delegates to searcher_->Search() which may not propagate reasoning context to the underlying graph search. This means expected labels that reside in GRAPH nodes will not be tracked during the graph traversal phase, potentially producing incomplete reasoning reports for those targets. |
||
| ctx.reasoning_ctx->RecordVisit(ids_ptr[i], dists[i], 0); | ||
| } | ||
| if (search_param.distance_threshold.has_value() and | ||
| (not std::isfinite(dists[i]) || | ||
| (not search_param.enable_reorder and | ||
|
|
@@ -1541,6 +1549,9 @@ Pyramid::search_node(const IndexNode* node, | |
| } | ||
| results->Push(dists[i], ids_ptr[i]); | ||
| if (results->Size() > search_param.ef) { | ||
| if (ctx.reasoning_ctx != nullptr) { | ||
| ctx.reasoning_ctx->RecordEviction(results->Top().second, 0); | ||
| } | ||
| results->Pop(); | ||
| } | ||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[note] The
#include "io/memory_io/memory_io_parameter.h"added here does not appear to be used anywhere inpyramid.cpp. The header is already included viapyramid.h(which includes it at line 32). Consider removing this redundant include to keep the include list minimal.