Skip to content

Commit

Permalink
enhance: optimize search reduce perf(#32507)
Browse files Browse the repository at this point in the history
Signed-off-by: MrPresent-Han <chun.han@gmail.com>
  • Loading branch information
MrPresent-Han committed Jul 12, 2024
1 parent d8e68cb commit 23fea06
Showing 1 changed file with 17 additions and 17 deletions.
34 changes: 17 additions & 17 deletions internal/core/src/segcore/reduce/Reduce.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -91,17 +91,21 @@ ReduceHelper::FilterInvalidSearchResult(SearchResult* search_result) {
auto& offsets = search_result->seg_offsets_;
auto& distances = search_result->distances_;

int segment_row_count = segment->get_row_count();
//1. for sealed segment, segment_row_count will not change as delete records will take effect as bitset
//2. for growing segment, segment_row_count is the minimum position acknowledged, which will only increase after
//the time at which the search operation is executed, so it's safe here to keep this value inside stack
for (auto i = 0; i < nq; ++i) {
for (auto j = 0; j < topK; ++j) {
auto index = i * topK + j;
if (offsets[index] != INVALID_SEG_OFFSET) {
AssertInfo(0 <= offsets[index] &&
offsets[index] < segment->get_row_count(),
offsets[index] < segment_row_count,
fmt::format("invalid offset {}, segment {} with "
"rows num {}, data or index corruption",
offsets[index],
segment->get_segment_id(),
segment->get_row_count()));
segment_row_count));
real_topks[i]++;
offsets[valid_index] = offsets[index];
distances[valid_index] = distances[index];
Expand Down Expand Up @@ -136,23 +140,22 @@ ReduceHelper::FillPrimaryKey() {
search_results_[valid_index++] = search_result;
}
}
search_results_.resize(valid_index);
num_segments_ = search_results_.size();
if(valid_index < search_results_.size()) {
search_results_.resize(valid_index);
num_segments_ = search_results_.size();
}
}

void
ReduceHelper::RefreshSearchResults() {
tracer::AutoSpan span(
"ReduceHelper::RefreshSearchResults", trace_ctx_, false);
std::vector<int64_t> real_topks(total_nq_, 0);
for (int i = 0; i < num_segments_; i++) {
std::vector<int64_t> real_topks(total_nq_, 0);
auto search_result = search_results_[i];
if (search_result->result_offsets_.size() != 0) {
RefreshSingleSearchResult(search_result, i, real_topks);
}
std::partial_sum(real_topks.begin(),
real_topks.end(),
search_result->topk_per_nq_prefix_sum_.begin() + 1);
}
}

Expand All @@ -164,23 +167,20 @@ ReduceHelper::RefreshSingleSearchResult(SearchResult* search_result,
for (int j = 0; j < total_nq_; j++) {
size += final_search_records_[seg_res_idx][j].size();
}
std::vector<milvus::PkType> primary_keys(size);
std::vector<float> distances(size);
std::vector<int64_t> seg_offsets(size);

uint32_t index = 0;
for (int j = 0; j < total_nq_; j++) {
for (auto offset : final_search_records_[seg_res_idx][j]) {
primary_keys[index] = search_result->primary_keys_[offset];
distances[index] = search_result->distances_[offset];
seg_offsets[index] = search_result->seg_offsets_[offset];
search_result->primary_keys_[index] = search_result->primary_keys_[offset];
search_result->distances_[index] = search_result->distances_[offset];
search_result->seg_offsets_[index] = search_result->seg_offsets_[offset];
index++;
real_topks[j]++;
}
}
search_result->primary_keys_.swap(primary_keys);
search_result->distances_.swap(distances);
search_result->seg_offsets_.swap(seg_offsets);
search_result->primary_keys_.resize(index);
search_result->distances_.resize(index);
search_result->seg_offsets_.resize(index);
}

void
Expand Down

0 comments on commit 23fea06

Please sign in to comment.