Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/algorithm/pyramid/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

set (PYRAMID_SRCS
pyramid.cpp
pyramid_index_node.cpp
pyramid_zparameters.cpp
)

Expand Down
309 changes: 160 additions & 149 deletions src/algorithm/pyramid/pyramid.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Copy link
Copy Markdown
Collaborator Author

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 in pyramid.cpp. The header is already included via pyramid.h (which includes it at line 32). Consider removing this redundant include to keep the include list minimal.

Copy link
Copy Markdown
Collaborator Author

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" does not appear to be used anywhere in this file. Consider removing it to keep includes minimal.

#include "quantization/rabitq_quantization/rabitq_quantizer_parameter.h"
#include "query_context.h"
#include "storage/empty_index_binary_set.h"
Expand Down Expand Up @@ -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);
Expand All @@ -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);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] create_knn_search_param does not set search_param.enable_reorder or search_param.enable_rabitq_one_bit_search, unlike KnnSearch which sets both (lines 266-271). When SearchWithRequest calls create_knn_search_param for KNN searches, reorder-related behavior (e.g. distance threshold filtering in search_node) will differ from the KnnSearch path because enable_reorder defaults to false. Similarly, enable_rabitq_one_bit_search will default to false, skipping the RaBitQ lower-bound optimization.

Consider adding these fields to create_knn_search_param:

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));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] Avoid signed overflow in the ef bound

Both operands of AMPLIFICATION_FACTOR * k are signed int64_t, so any positive k above INT64_MAX / 100 overflows before conversion to uint64_t, invoking undefined behavior instead of returning a controlled argument error. Validate the multiplication bound or perform a checked unsigned calculation.

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;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The SearchWithRequest KNN path (via create_knn_search_param) does not set search_param.enable_reorder or search_param.enable_rabitq_one_bit_search, unlike the existing KnnSearch method which sets both. This means SearchWithRequest will bypass reorder and rabitq optimizations even when the index was configured with use_reorder_=true.

Consider adding these fields to create_knn_search_param or factoring the shared parameter setup into a common helper to avoid drift between the two search paths.

search_param.parallel_search_thread_count = parsed_param.parallel_search_thread_count;
if (this->support_duplicate_) {
search_param.consider_duplicate = true;
}

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[critical] create_knn_search_param is missing enable_reorder and enable_rabitq_one_bit_search settings that are present in KnnSearch.

In KnnSearch (pyramid.cpp:266-271), these fields are set:

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 create_knn_search_param (pyramid.cpp:153-184) does not set either of them. This means KNN searches via SearchWithRequest will silently skip reorder, which is a correctness regression compared to KnnSearch.

Suggested fix: add the same two lines to create_knn_search_param:

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_;

}

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[critical] create_knn_search_param does not set search_param.enable_reorder, unlike the equivalent logic in KnnSearch (line 266). This means reorder-based distance refinement will be silently skipped for all KNN searches performed through SearchWithRequest, leading to lower-quality (quantized) distances in results compared to KnnSearch.

Add after search_param.search_mode = KNN_SEARCH;:

search_param.enable_reorder = use_reorder_;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[critical] create_knn_search_param is missing search_param.enable_rabitq_one_bit_search and search_param.distance_threshold that KnnSearch sets (lines 269-271). This means:

  1. RaBitQ one-bit search optimization will never be used via SearchWithRequest
  2. Distance threshold filtering (used for early termination) will not work

KnnSearch sets:

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 create_knn_search_param as well.

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);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] create_knn_search_param does not set search_param.enable_reorder or search_param.enable_rabitq_one_bit_search, unlike the equivalent code paths in KnnSearch (line 266-272) and RangeSearch (line 339-342). This means SearchWithRequest will silently skip the reorder phase and the RaBitQ one-bit search optimization, producing different (potentially lower-quality) results than KnnSearch with the same parameters.

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; statement.

return search_param;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] create_knn_search_param does not set search_param.enable_reorder, while KnnSearch (line 266) explicitly sets search_param.enable_reorder = use_reorder_. The default value of enable_reorder in InnerSearchParam is true, so this works correctly when use_reorder_ is true, but would be inconsistent if use_reorder_ were false.

Currently this has no behavioral impact because SearchWithRequest does not populate search_param.distance_threshold (the only code path that checks enable_reorder is gated by distance_threshold.has_value() in search_node line 1544). However, if request.threshold_ support is added in the future, this could cause search_node to skip the distance-threshold check even when reorder is disabled, potentially admitting out-of-range results into the candidate heap.

Consider adding search_param.enable_reorder = use_reorder_; to create_knn_search_param for consistency with KnnSearch.

}

std::vector<int64_t>
Expand All @@ -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);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] The label_table_->ResetRemap + InsertRemap calls added to build_by_odescent execute unconditionally for all Pyramid builds, not just when reasoning is enabled. This adds O(N) overhead to every build. If the remap table is only needed for the reasoning feature, consider making this conditional (e.g. gated behind a config flag) to avoid the overhead for non-reasoning use cases.

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()) {
Expand Down Expand Up @@ -505,6 +388,126 @@ Pyramid::RangeSearch(const DatasetPtr& query,
return result;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] create_knn_search_param does not set enable_rabitq_one_bit_search, while the existing KnnSearch method does (line 269-271). This means searches via SearchWithRequest (KNN mode) will always use the default false value, which differs from KnnSearch behavior where it is set from parsed_param.has_rabitq_one_bit_search. This could lead to inconsistent search quality when rabitq quantization is used.

Suggested fix: add the same enable_rabitq_one_bit_search initialization to create_knn_search_param:

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_;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The 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 ctx leaves search_impl's outer heap and result buffers, plus Pyramid's FLAT scratch vectors and heap, allocated from allocator_. A caller supplying a per-search allocator therefore still consumes the index allocator for material request allocations, defeating isolation or bounded-pool accounting. Use the selected allocator consistently throughout search and result packing.

}
Comment thread
LHT129 marked this conversation as resolved.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] ctx.rabitq_error_rate is not set in SearchWithRequest, unlike KnnSearch and RangeSearch.

Both KnnSearch (pyramid.cpp:251) and RangeSearch (pyramid.cpp:331) set:

ctx.rabitq_error_rate = parsed_param.rabitq_error_rate;

This is missing from SearchWithRequest. If rabitq quantization is used, the error rate setting may affect distance computation accuracy during search.


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_);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve Pyramid's RaBitQ search parameters

Unlike both KnnSearch and RangeSearch, this path never assigns ctx.rabitq_error_rate, never sets search_param.enable_rabitq_one_bit_search, and always passes null lower-bound candidates to reorder. A request enabling one-bit RaBitQ—or a split-code index where it is the default—therefore silently takes a different full-code traversal/reorder path and ignores the requested error rate. Propagate the same RaBitQ setup used by the legacy search paths.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] SearchWithRequest does not set ctx.rabitq_error_rate, unlike both KnnSearch (line 251) and RangeSearch (line 331) which set ctx.rabitq_error_rate = parsed_param.rabitq_error_rate;. This means RaBitQ quantization error rate control will not be applied when searching through SearchWithRequest, potentially affecting recall when RaBitQ is in use.

Add after parsed_param is obtained:

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_);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Honor the KNN distance threshold

This path never copies request.threshold_ into the search parameters or filters the packed result afterward. For example, a KNN request with topk_=2, threshold 0.5, and candidates at distances 0 and 100 returns both candidates, unlike legacy KnnSearch and other request implementations. Preserve the inclusive threshold contract when constructing and packing the request search.

} else {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[critical] SearchWithRequest KNN path is missing ctx.rabitq_error_rate setting. KnnSearch (line 264) and RangeSearch both set:

ctx.rabitq_error_rate = parsed_param.rabitq_error_rate;

Without this, RaBitQ quantization error rate control will not work when searching via SearchWithRequest, leading to incorrect distance lower bounds and potentially missed results.

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;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The range-search parameter setup in SearchWithRequest (lines 418-427) largely duplicates the same logic in RangeSearch (lines 335-345). Consider extracting a create_range_search_param helper method (similar to create_knn_search_param) to avoid the duplication. This would ensure future changes to range-search parameter logic only need to be made in one place.

Suggested helper signature:

InnerSearchParam
create_range_search_param(const PyramidSearchParameters& parsed_param,
                          float radius,
                          int64_t limited_size,
                          const FilterPtr& filter) const;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The SearchWithRequest range search path does not set search_param.enable_reorder or search_param.enable_rabitq_one_bit_search, unlike the existing RangeSearch method. This creates a behavioral difference between the two range search code paths.

search_param.radius = request.radius_ * RADIUS_EPSILON;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[critical] SearchWithRequest range search path is missing several settings that RangeSearch applies:

  1. search_param.enable_reorder = use_reorder_ — without this, precise distance reordering is disabled for range searches via SearchWithRequest
  2. search_param.enable_rabitq_one_bit_search — RaBitQ one-bit search optimization is never activated
  3. ctx.rabitq_error_rate = parsed_param.rabitq_error_rate — RaBitQ error rate control is missing

RangeSearch sets all three of these. They should be added here for consistency.

search_param.search_mode = RANGE_SEARCH;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The range search branch in SearchWithRequest (lines 415-432) also does not set search_param.enable_reorder or search_param.enable_rabitq_one_bit_search, unlike RangeSearch (lines 339-341). This means reorder and RaBitQ lower-bound optimizations will be silently disabled for range searches through SearchWithRequest.

Add after search_param.search_mode = RANGE_SEARCH;:

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_;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The range search path in SearchWithRequest also does not set search_param.enable_reorder or search_param.enable_rabitq_one_bit_search, unlike RangeSearch (lines 339-342). Same issue as create_knn_search_param — this will silently skip the reorder phase and RaBitQ one-bit search optimization for range queries via SearchWithRequest.

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.search_mode = RANGE_SEARCH;.

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()

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The 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:
search_param.enable_reorder = use_reorder_;
search_param.enable_rabitq_one_bit_search = ...

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;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[critical] The range search branch in SearchWithRequest (pyramid.cpp:416-432) is missing enable_reorder and enable_rabitq_one_bit_search settings that are present in RangeSearch.

In RangeSearch (pyramid.cpp:339-341), these fields are set:

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: enable_reorder is also missing from RangeSearch, but that is a pre-existing issue.)

However, the range search path in SearchWithRequest does not set enable_reorder or enable_rabitq_one_bit_search at all. This means range searches via SearchWithRequest will not use reorder or rabitq one-bit search optimizations.

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_);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The 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 request.filter_; it never reads enable_bitset_filter_ or bitset_filter_, and it applies filter_ even when enable_filter_ is false. Consequently a KNN or range request whose bitset excludes the nearest label can still return that label, silently bypassing a requested exclusion. Combine only the enabled filter fields, as the SearchRequest contract requires.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The 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) {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] The search_func lambda in SearchWithRequest (line 433) captures search_param by reference ([&]), but search_param is a local variable on the stack. This is fine for the current single-threaded call path since search_func is only used synchronously within search_impl. However, if search_impl or search_hierarchy ever dispatches search_func to a thread pool (as search_hierarchy already does for parallel path search at line 1464), the reference to search_param would dangle after SearchWithRequest returns. The same pattern exists in KnnSearch and RangeSearch, so this is a pre-existing concern, not specific to this PR.

return this->search_node(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The 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 RecordVisit or RecordFilterReject for it. If a filter rejects an expected entry-point label while other nodes remain eligible, the report diagnoses that target as not_reachable rather than filter_rejected. Record the entry-point events before graph expansion.

node, vl, search_param, query, base_codes_, ctx, parsed_param.subindex_ef_search);
Comment thread
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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The range search branch in SearchWithRequest (lines 415-431) does not set enable_rabitq_one_bit_search, while the existing RangeSearch method does (line 339-341). This means range searches via SearchWithRequest will always use the default false value for enable_rabitq_one_bit_search, which differs from RangeSearch behavior.

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 =

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[critical] The cur_element_count_mutex_ is held during precise_flatten->Query(...), which is a potentially expensive distance computation. This lock is also acquired by Add operations (via cur_element_count_mutex_), so holding it during distance computation will block concurrent insertions unnecessarily.

The lock should only protect the label-to-inner-id mapping. The distance computation should be moved outside the lock scope.

Suggested fix: collect expected_inner_ids inside the lock, release the lock, then compute true distances:

{
    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_;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The 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 base_codes_, which may use SQ, PQ, or RaBitQ and is the same approximate representation used during search. Even after fixing initialization order, the recorded "true" distance then matches the quantized distance, so a miss caused by quantization is reported as unknown. Use a lossless representation or report this diagnosis as unavailable for such configurations.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The 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 base_codes_, which may be SQ, PQ, or RaBitQ and is the same approximate representation used by search. Comparing that value with the search distance cannot identify a miss caused by quantization and can report unknown instead of quantization_error. Use a lossless reference representation, or explicitly make this diagnosis unavailable when none is stored.

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]);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Initialize target traces before setting distances

SetTrueDistance only updates an entry already present in expected_traces_, but InitializeExpectedTargets is not called until line 291. Every distance computed here is therefore discarded, leaving true_distance at zero and making quantization_error impossible to diagnose. Initialize the targets before setting these distances.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Initialize expected traces before setting distances

ReasoningContext::SetTrueDistance only updates IDs already present in expected_traces_, but InitializeExpectedTargets is called later at line 475. Every distance computed by this prepass is discarded, leaving missed targets with true_distance == 0 and preventing the promised quantization diagnosis. Initialize the targets before these calls.

}
}
}

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);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] SearchWithRequest does not collect rabitq lower-bound candidates, unlike KnnSearch and RangeSearch.

Both KnnSearch (pyramid.cpp:282-304) and RangeSearch (pyramid.cpp:354-376) collect rabitq_lower_bound_candidates when enable_rabitq_one_bit_search is active and pass them to search_impl for use during reorder. This optimization is missing from SearchWithRequest (line 485), which always passes nullptr as the last argument to search_impl.

If enable_rabitq_one_bit_search is added to create_knn_search_param (per the comment above), this optimization should also be wired through to maintain parity with the existing search paths.

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();

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Use the result dimension as the hit count

search_impl packs one query as Dim(hit_count)->NumElements(1), and MakeEmptyDataset() is Dim(0)->NumElements(1) with null IDs. Therefore this value is always 1: an empty reasoning search dereferences result_ids[0], while a nonempty top-k search marks only its first hit and reports expected labels at later ranks as missed. Use GetDim() or retain the result inner IDs before packing, and handle the empty case.

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,
Expand Down Expand Up @@ -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();
Expand All @@ -1533,6 +1538,9 @@ Pyramid::search_node(const IndexNode* node,
codes->Query(dists.data(), computer, ids_ptr, id_count, &ctx);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The reasoning hooks in search_node pass 0 as the hop parameter to RecordVisit and RecordEviction. This means the reasoning report cannot distinguish at which search hop a candidate was visited or evicted, limiting the diagnostic value (e.g. it cannot report "visited at hop 3" vs "visited at hop 10").

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) {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The loop variable int i should be uint64_t to match the type of id_count (which is uint64_t from Vector::size()). This avoids a signed/unsigned comparison warning and potential truncation for very large node sizes.

if (ctx.reasoning_ctx != nullptr) {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The 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
Expand All @@ -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();
}
}
Expand Down
Loading
Loading