diff --git a/src/algorithm/inner_index_interface.h b/src/algorithm/inner_index_interface.h index 38bd7dc1ce..77d182336f 100644 --- a/src/algorithm/inner_index_interface.h +++ b/src/algorithm/inner_index_interface.h @@ -698,7 +698,7 @@ class InnerIndexInterface { protected: std::atomic total_count_{0}; - std::atomic current_memory_usage_{0}; + mutable std::atomic current_memory_usage_{0}; mutable std::shared_mutex memory_usage_mutex_{}; bool has_raw_vector_{false}; diff --git a/src/algorithm/sindi/sindi.cpp b/src/algorithm/sindi/sindi.cpp index 0c3396b7ff..dc057a7a17 100644 --- a/src/algorithm/sindi/sindi.cpp +++ b/src/algorithm/sindi/sindi.cpp @@ -425,12 +425,7 @@ SINDI::Add(const DatasetPtr& base) { throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, "SINDI DMQ rerank does not support incremental Add"); } - const auto previous_window_count = mutable_term_datacell_->GetWindowCount(); - auto failed_ids = this->add(base, true); - if (mutable_term_datacell_->GetWindowCount() != previous_window_count) { - this->cal_memory_usage(); - } - return failed_ids; + return this->add(base, true); } std::vector @@ -455,9 +450,9 @@ SINDI::add(const DatasetPtr& base, bool sort_affected_windows) { Vector pruned_ids(allocator_); Vector pruned_vals(allocator_); Vector remapped_ids(allocator_); - const auto first_affected_window = cur_element_count_ / window_size_; - // This remains -1 when every input vector is rejected, so post-insert loops are no-ops. - int64_t last_affected_window = -1; + const auto first_affected_window = + static_cast(cur_element_count_.load(std::memory_order_relaxed) / window_size_); + std::optional last_affected_window; std::vector rerank_vectors; if (use_reorder_) { rerank_vectors.reserve(data_num); @@ -507,16 +502,32 @@ SINDI::add(const DatasetPtr& base, bool sort_affected_windows) { if (use_reorder_) { rerank_vectors.push_back(sparse_vectors[i]); } - last_affected_window = cur_element_count_ / window_size_; + last_affected_window = static_cast( + cur_element_count_.load(std::memory_order_relaxed) / window_size_); cur_element_count_++; } if (not rerank_vectors.empty()) { rerank_flat_->BatchInsertVector(rerank_vectors.data(), static_cast(rerank_vectors.size())); } - if (sort_affected_windows) { - for (int64_t window = first_affected_window; window <= last_affected_window; ++window) { - mutable_term_datacell_->SortByValue(static_cast(window)); + if (sort_affected_windows && last_affected_window.has_value()) { + bool posting_state_changed = false; + const auto current_count = + static_cast(cur_element_count_.load(std::memory_order_relaxed)); + for (uint32_t window = first_affected_window;; ++window) { + const auto window_end = + (static_cast(window) + 1) * static_cast(window_size_); + if (window_end <= current_count) { + posting_state_changed |= mutable_term_datacell_->NormalizeDirtyPostings(window); + } else { + posting_state_changed |= mutable_term_datacell_->FinalizeInsertBatch(window); + } + if (window == last_affected_window.value()) { + break; + } + } + if (posting_state_changed) { + this->cal_memory_usage(); } } return failed_ids; @@ -1144,7 +1155,7 @@ SINDI::UseTermListsHeapInsert(const SINDISearchParameter& search_param, } void -SINDI::cal_memory_usage() { +SINDI::cal_memory_usage() const { auto memory = sizeof(SINDI); if (term_datacell_ != nullptr) { memory += term_datacell_->GetMemoryUsage(); @@ -1163,9 +1174,17 @@ SINDI::cal_memory_usage() { this->current_memory_usage_.store(static_cast(memory)); } +void +SINDI::normalize_dirty_postings_for_serialization() const { + if (mutable_term_datacell_ != nullptr && mutable_term_datacell_->NormalizeDirtyPostings()) { + this->cal_memory_usage(); + } +} + void SINDI::Serialize(StreamWriter& writer) const { - std::shared_lock rlock(this->global_mutex_); + std::scoped_lock wlock(this->global_mutex_); + this->normalize_dirty_postings_for_serialization(); if (cur_element_count_ == 0) { const auto cur_element_count = cur_element_count_.load(); @@ -1303,9 +1322,10 @@ SINDI::serialize_windows(StreamWriter& writer) const { void SINDI::serialize_streaming_body(StreamWriter& writer) const { - std::shared_lock rlock(this->global_mutex_); + std::scoped_lock wlock(this->global_mutex_); CHECK_ARGUMENT(not immutable_enabled_, "immutable SINDI runtime does not support SerializeStreaming"); + this->normalize_dirty_postings_for_serialization(); auto windows_tag = static_cast(StreamSerializationTag::SINDI_WINDOWS); auto label_tag = static_cast(StreamSerializationTag::LABEL_TABLE); diff --git a/src/algorithm/sindi/sindi.h b/src/algorithm/sindi/sindi.h index 619b0dc00e..3dd585fc71 100644 --- a/src/algorithm/sindi/sindi.h +++ b/src/algorithm/sindi/sindi.h @@ -269,7 +269,12 @@ class SINDI : public InnerIndexInterface { /// Recalculate and cache the memory-usage counter. void - cal_memory_usage(); + cal_memory_usage() const; + + /// Canonicalize mutable posting runs before writing the sorted-posting format. + /// The caller must hold global_mutex_ exclusively. + void + normalize_dirty_postings_for_serialization() const; /** * @brief Compact a sparse vector's dim-ids into the remapped space diff --git a/src/algorithm/sindi/sindi_test.cpp b/src/algorithm/sindi/sindi_test.cpp index 160a808758..c373ae2f2a 100644 --- a/src/algorithm/sindi/sindi_test.cpp +++ b/src/algorithm/sindi/sindi_test.cpp @@ -500,7 +500,7 @@ TEST_CASE("SINDI term prune keeps highest stored values after build", "[ut][SIND } } -TEST_CASE("SINDI sorts incremental partial windows", "[ut][SINDI]") { +TEST_CASE("SINDI defers incremental partial window normalization", "[ut][SINDI]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); IndexCommonParam common_param; common_param.allocator_ = allocator; @@ -536,7 +536,7 @@ TEST_CASE("SINDI sorts incremental partial windows", "[ut][SINDI]") { auto appended = Dataset::Make(); appended->NumElements(1)->SparseVectors(&appended_vector)->Ids(&appended_label)->Owner(false); REQUIRE(index.Add(appended).empty()); - REQUIRE(SINDITestAccess::MutableTermIsSorted(index, 0, term)); + REQUIRE_FALSE(SINDITestAccess::MutableTermIsSorted(index, 0, term)); float query_value = 1.0F; SparseVector query_vector{1, &term, &query_value}; @@ -545,6 +545,11 @@ TEST_CASE("SINDI sorts incremental partial windows", "[ut][SINDI]") { const auto search_parameters = R"({"sindi": {"n_candidate": 1, "term_retain_threshold": 1}})"; REQUIRE(index.KnnSearch(query, 1, search_parameters, nullptr)->GetIds()[0] == appended_label); + std::stringstream stream; + IOStreamWriter writer(stream); + REQUIRE_NOTHROW(index.Serialize(writer)); + REQUIRE(SINDITestAccess::MutableTermIsSorted(index, 0, term)); + appended_value = 3.0F; appended_label = 13; REQUIRE(index.Add(appended).empty()); diff --git a/src/algorithm/sindi_v2/sindi_v2.cpp b/src/algorithm/sindi_v2/sindi_v2.cpp index fb6005ec8b..5c3674a4e9 100644 --- a/src/algorithm/sindi_v2/sindi_v2.cpp +++ b/src/algorithm/sindi_v2/sindi_v2.cpp @@ -514,8 +514,8 @@ SINDIV2::Add(const DatasetPtr& base) { Vector pruned_ids(allocator_); Vector pruned_vals(allocator_); Vector remapped_ids(allocator_); - const auto first_affected_window = cur_element_count_ / window_size_; - int64_t last_affected_window = -1; + const auto first_affected_window = static_cast(cur_element_count_ / window_size_); + std::optional last_affected_window; std::vector dmq_rerank_vectors; if (use_reorder_ && rerank_type_ == SPARSE_RERANK_TYPE_DMQ8) { dmq_rerank_vectors.reserve(data_num); @@ -581,7 +581,7 @@ SINDIV2::Add(const DatasetPtr& base) { {sparse_vectors + i, static_cast(cur_element_count_), {}}); } - last_affected_window = cur_element_count_ / window_size_; + last_affected_window = static_cast(cur_element_count_ / window_size_); cur_element_count_++; } @@ -593,8 +593,20 @@ SINDIV2::Add(const DatasetPtr& base) { write_rerank_flat_with_layout(rerank_flat_, rerank_layout_records, rerank_layout_); } - for (int64_t window = first_affected_window; window <= last_affected_window; ++window) { - mutable_term_datacell->SortByValue(static_cast(window)); + if (last_affected_window.has_value()) { + const auto current_count = static_cast(cur_element_count_); + for (uint32_t window = first_affected_window;; ++window) { + const auto window_end = + (static_cast(window) + 1) * static_cast(window_size_); + if (window_end <= current_count) { + mutable_term_datacell->NormalizeDirtyPostings(window); + } else { + mutable_term_datacell->FinalizeInsertBatch(window); + } + if (window == last_affected_window.value()) { + break; + } + } } this->cal_memory_usage(); return failed_ids; @@ -1063,7 +1075,7 @@ SINDIV2::UseTermListsHeapInsert(const SINDIV2SearchParameter& search_param) cons } void -SINDIV2::cal_memory_usage() { +SINDIV2::cal_memory_usage() const { auto memory = sizeof(SINDIV2); if (term_datacell_ != nullptr) { memory += term_datacell_->GetMemoryUsage(); @@ -1106,9 +1118,8 @@ SINDIV2::Serialize(StreamWriter& writer) const { if (term_datacell_ != nullptr && std::dynamic_pointer_cast(term_datacell_) != nullptr) { - const auto mutable_datacell = this->get_mutable_term_datacell(); - for (uint32_t window = 0; window < mutable_datacell->GetWindowCount(); ++window) { - mutable_datacell->SortByValue(window); + if (this->get_mutable_term_datacell()->NormalizeDirtyPostings()) { + this->cal_memory_usage(); } } diff --git a/src/algorithm/sindi_v2/sindi_v2.h b/src/algorithm/sindi_v2/sindi_v2.h index c7ff1384e8..6b7f4c6e0b 100644 --- a/src/algorithm/sindi_v2/sindi_v2.h +++ b/src/algorithm/sindi_v2/sindi_v2.h @@ -161,7 +161,7 @@ class SINDIV2 : public InnerIndexInterface { get_min_max_window_id(const FilterPtr& filter) const; void - cal_memory_usage(); + cal_memory_usage() const; SparseVector sort_and_prune_sparse_vector_for_build(const SparseVector& input, diff --git a/src/algorithm/sindi_v2/sindi_v2_test.cpp b/src/algorithm/sindi_v2/sindi_v2_test.cpp index 5df98f2b37..667cf9ec5d 100644 --- a/src/algorithm/sindi_v2/sindi_v2_test.cpp +++ b/src/algorithm/sindi_v2/sindi_v2_test.cpp @@ -267,7 +267,7 @@ TEST_CASE("SINDIV2 term prune keeps highest stored values after build", "[ut][SI } } -TEST_CASE("SINDIV2 sorts incremental partial windows", "[ut][SINDIV2]") { +TEST_CASE("SINDIV2 defers incremental partial window normalization", "[ut][SINDIV2]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); IndexCommonParam common_param; common_param.allocator_ = allocator; @@ -305,7 +305,7 @@ TEST_CASE("SINDIV2 sorts incremental partial windows", "[ut][SINDIV2]") { auto appended = Dataset::Make(); appended->NumElements(1)->SparseVectors(&appended_vector)->Ids(&appended_label)->Owner(false); REQUIRE(index.Add(appended).empty()); - REQUIRE(SINDIV2TestAccess::MutableTermIsSorted(index, 0, term)); + REQUIRE_FALSE(SINDIV2TestAccess::MutableTermIsSorted(index, 0, term)); float query_value = 1.0F; SparseVector query_vector{1, &term, &query_value}; diff --git a/src/datacell/mutable_sindi_term_datacell.cpp b/src/datacell/mutable_sindi_term_datacell.cpp index 420a6f402f..72b34dfa86 100644 --- a/src/datacell/mutable_sindi_term_datacell.cpp +++ b/src/datacell/mutable_sindi_term_datacell.cpp @@ -29,7 +29,6 @@ namespace vsag { void MutableSindiTermDataCell::Finalize() { for (auto& window : windows_) { - this->SortByValue(window); this->Compact(window); } } @@ -210,6 +209,74 @@ MutableSindiTermDataCell::GetTermDictCount() const { return term_dict_count; } +MutableSindiTermDataCell::PostingRunSelection +MutableSindiTermDataCell::SelectPostingRuns(const MutableSINDIWindow& window, + uint32_t term, + uint32_t retained_count) const { + const auto posting_count = window.term_sizes_[term]; + retained_count = std::min(retained_count, posting_count); + if (window.postings_sorted_) { + return {retained_count, posting_count, 0}; + } + const auto dirty = window.dirty_posting_prefixes_.find(term); + if (dirty == window.dirty_posting_prefixes_.end()) { + return {retained_count, posting_count, 0}; + } + + const auto prefix_count = dirty->second; + const auto suffix_count = posting_count - prefix_count; + if (retained_count == 0) { + return {0, prefix_count, 0}; + } + if (retained_count == posting_count) { + return {prefix_count, prefix_count, suffix_count}; + } + + uint32_t lower = retained_count > prefix_count ? retained_count - prefix_count : 0; + uint32_t upper = std::min(retained_count, suffix_count); + while (lower <= upper) { + const auto take_suffix = lower + (upper - lower) / 2; + const auto take_prefix = retained_count - take_suffix; + if (take_prefix > 0 && take_suffix < suffix_count && + PostingEntryBefore(window, term, prefix_count + take_suffix, take_prefix - 1)) { + lower = take_suffix + 1; + continue; + } + if (take_suffix > 0 && take_prefix < prefix_count && + PostingEntryBefore(window, term, take_prefix, prefix_count + take_suffix - 1)) { + upper = take_suffix - 1; + continue; + } + return {take_prefix, prefix_count, take_suffix}; + } + + CHECK_ARGUMENT(false, "failed to partition sorted posting runs"); +} + +void +MutableSindiTermDataCell::ScanPostingRange(uint32_t term_iterator, + const SparseTermComputerPtr& computer, + const uint16_t* term_ids, + const uint8_t* term_data, + uint32_t term_count, + float* dists, + SparseEvaluationTracker* evaluation_tracker) const { + if (term_count == 0) { + return; + } + if (evaluation_tracker != nullptr) { + evaluation_tracker->Mark(term_ids, term_count); + } + if (sparse_value_quant_type_ == SparseValueQuantizationType::SQ8) { + computer->ScanForAccumulateSQ8(term_iterator, term_ids, term_data, term_count, dists); + } else if (sparse_value_quant_type_ == SparseValueQuantizationType::FP16) { + computer->ScanForAccumulateFP16Bytes(term_iterator, term_ids, term_data, term_count, dists); + } else { + computer->ScanForAccumulateFloatBytes( + term_iterator, term_ids, term_data, term_count, dists); + } +} + void MutableSindiTermDataCell::QueryWindow(float* dists, uint32_t window_id, @@ -237,26 +304,26 @@ MutableSindiTermDataCell::QueryWindow(float* dists, const auto posting_count = window.term_sizes_[term]; const auto term_size = computer->GetTermScanCount(posting_count); - query_context.evaluation_tracker.Mark(window.term_ids_[term]->data(), term_size); - - if (sparse_value_quant_type_ == SparseValueQuantizationType::SQ8) { - computer->ScanForAccumulateSQ8(it, - window.term_ids_[term]->data(), - window.term_datas_[term]->data(), - term_size, - dists); - } else if (sparse_value_quant_type_ == SparseValueQuantizationType::FP16) { - computer->ScanForAccumulateFP16Bytes(it, - window.term_ids_[term]->data(), - window.term_datas_[term]->data(), - term_size, - dists); - } else { - computer->ScanForAccumulateFloatBytes(it, - window.term_ids_[term]->data(), - window.term_datas_[term]->data(), - term_size, - dists); + const auto selected = this->SelectPostingRuns(window, term, term_size); + const auto code_size = this->GetTermValueCodeSize(); + const auto* term_ids = window.term_ids_[term]->data(); + const auto* term_data = window.term_datas_[term]->data(); + this->ScanPostingRange(it, + computer, + term_ids, + term_data, + selected.prefix_count, + dists, + &query_context.evaluation_tracker); + if (selected.suffix_count > 0) { + this->ScanPostingRange( + it, + computer, + term_ids + selected.suffix_offset, + term_data + static_cast(selected.suffix_offset) * code_size, + selected.suffix_count, + dists, + &query_context.evaluation_tracker); } } computer->ResetTerm(); @@ -500,61 +567,110 @@ MutableSindiTermDataCell::InsertHeapByTermLists(const MutableSINDIWindow& window continue; } - uint32_t i = 0; const auto posting_count = window.term_sizes_[term]; const auto term_size = computer->GetTermScanCount(posting_count); - auto& one_term_ids = *window.term_ids_[term]; - if constexpr (mode == InnerSearchMode::KNN_SEARCH) { - if (heap.size() < n_candidate) { - for (; i < term_size; i++) { - id = one_term_ids[i]; - const bool heap_filled = fill_heap_initial(id, - dists[id], - cur_heap_top, - heap, - offset_id, - n_candidate, - filter, - param.distance_threshold, - param.enable_reorder); - if constexpr (type == InnerSearchType::WITH_FILTER_LIMIT) { - if (filter_callback_remaining != nullptr and - *filter_callback_remaining == 0) { - computer->ResetTerm(); - return true; + if (window.postings_sorted_) { + uint32_t i = 0; + const auto& term_ids = *window.term_ids_[term]; + if constexpr (mode == InnerSearchMode::KNN_SEARCH) { + if (heap.size() < n_candidate) { + for (; i < term_size; ++i) { + id = term_ids[i]; + const bool heap_filled = fill_heap_initial(id, + dists[id], + cur_heap_top, + heap, + offset_id, + n_candidate, + filter, + param.distance_threshold, + param.enable_reorder); + if constexpr (type == InnerSearchType::WITH_FILTER_LIMIT) { + if (filter_callback_remaining != nullptr && + *filter_callback_remaining == 0) { + computer->ResetTerm(); + return true; + } + } + if (heap_filled) { + ++i; + break; } } - if (heap_filled) { - i++; - break; + } + } + for (; i < term_size; ++i) { + id = term_ids[i]; + if constexpr (mode == InnerSearchMode::RANGE_SEARCH) { + if (candidate_tracker != nullptr && !candidate_tracker->MarkOne(id)) { + continue; + } + } + insert_candidate_into_heap(id, + dists[id], + cur_heap_top, + heap, + offset_id, + n_candidate, + radius, + filter, + param.distance_threshold, + param.enable_reorder); + if constexpr (type == InnerSearchType::WITH_FILTER_LIMIT) { + if (filter_callback_remaining != nullptr && *filter_callback_remaining == 0) { + computer->ResetTerm(); + return true; } } } + continue; } - - for (; i < term_size; i++) { - id = one_term_ids[i]; - if constexpr (mode == InnerSearchMode::RANGE_SEARCH) { - if (candidate_tracker != nullptr && not candidate_tracker->MarkOne(id)) { - continue; + const auto selected = this->SelectPostingRuns(window, term, term_size); + bool heap_initialized = heap.size() >= n_candidate; + const auto completed = + this->ForEachSelectedPosting(window, term, selected, [&](uint16_t selected_id) { + id = selected_id; + if constexpr (mode == InnerSearchMode::RANGE_SEARCH) { + if (candidate_tracker != nullptr && !candidate_tracker->MarkOne(id)) { + return true; + } } - } - insert_candidate_into_heap(id, - dists[id], - cur_heap_top, - heap, - offset_id, - n_candidate, - radius, - filter, - param.distance_threshold, - param.enable_reorder); - if constexpr (type == InnerSearchType::WITH_FILTER_LIMIT) { - if (filter_callback_remaining != nullptr and *filter_callback_remaining == 0) { - computer->ResetTerm(); - return true; + if constexpr (mode == InnerSearchMode::KNN_SEARCH) { + if (!heap_initialized) { + heap_initialized = fill_heap_initial(id, + dists[id], + cur_heap_top, + heap, + offset_id, + n_candidate, + filter, + param.distance_threshold, + param.enable_reorder); + if constexpr (type == InnerSearchType::WITH_FILTER_LIMIT) { + return filter_callback_remaining == nullptr || + *filter_callback_remaining != 0; + } + return true; + } } - } + insert_candidate_into_heap(id, + dists[id], + cur_heap_top, + heap, + offset_id, + n_candidate, + radius, + filter, + param.distance_threshold, + param.enable_reorder); + if constexpr (type == InnerSearchType::WITH_FILTER_LIMIT) { + return filter_callback_remaining == nullptr || *filter_callback_remaining != 0; + } + return true; + }); + if (!completed) { + computer->ResetTerm(); + return true; } } computer->ResetTerm(); @@ -655,6 +771,7 @@ MutableSindiTermDataCell::InsertVector(const SparseVector& sparse_base, uint32_t window->term_datas_[term] = std::make_unique>(allocator_); } + const auto old_term_size = window->term_sizes_[term]; window->term_ids_[term]->push_back(window_local_id); auto& data_vec = *window->term_datas_[term]; @@ -664,6 +781,11 @@ MutableSindiTermDataCell::InsertVector(const SparseVector& sparse_base, uint32_t val, sparse_value_quant_type_, quantization_params_.get(), data_vec.data() + old_size); window->term_sizes_[term] += 1; + window->dirty_posting_prefixes_.try_emplace(term, old_term_size); + if (window->pending_posting_flags_[term] == 0) { + window->pending_posting_flags_[term] = 1; + window->pending_posting_terms_.push_back(term); + } } if (sparse_base.len_ > 0) { window->postings_sorted_ = false; @@ -678,7 +800,7 @@ MutableSindiTermDataCell::SortByValue(uint32_t window_id) { } void -MutableSindiTermDataCell::SortByValue(MutableSINDIWindow& window) const { +MutableSindiTermDataCell::SortByValue(MutableSINDIWindow& window) { if (window.postings_sorted_) { return; } @@ -703,7 +825,219 @@ MutableSindiTermDataCell::SortByValue(MutableSINDIWindow& window) const { sorted_ids, sorted_data); } + window.dirty_posting_prefixes_.clear(); + window.dirty_posting_prefixes_.rehash(0); + for (uint32_t term : window.pending_posting_terms_) { + window.pending_posting_flags_[term] = 0; + } + window.pending_posting_terms_.clear(); + window.postings_sorted_ = true; + this->ReleaseNormalizationScratch(); +} + +bool +MutableSindiTermDataCell::PostingEntryBefore(const MutableSINDIWindow& window, + uint32_t term, + uint32_t left, + uint32_t right) const { + const auto& ids = *window.term_ids_[term]; + const auto& data = *window.term_datas_[term]; + const auto compare = [&ids, &data, left, right](auto code_type) { + using CodeType = decltype(code_type); + CodeType left_code{}; + CodeType right_code{}; + std::memcpy(&left_code, + data.data() + static_cast(left) * sizeof(CodeType), + sizeof(CodeType)); + std::memcpy(&right_code, + data.data() + static_cast(right) * sizeof(CodeType), + sizeof(CodeType)); + if (left_code != right_code) { + return left_code > right_code; + } + return ids[left] < ids[right]; + }; + + switch (sparse_value_quant_type_) { + case SparseValueQuantizationType::SQ8: + return compare(uint8_t{}); + case SparseValueQuantizationType::FP16: + return compare(uint16_t{}); + case SparseValueQuantizationType::FP32: + return compare(float{}); + default: + CHECK_ARGUMENT(false, "unknown sparse value quantization type"); + } +} + +void +MutableSindiTermDataCell::NormalizePosting(MutableSINDIWindow& window, uint32_t term) { + const auto dirty = window.dirty_posting_prefixes_.find(term); + if (dirty == window.dirty_posting_prefixes_.end()) { + return; + } + + const auto prefix_count = dirty->second; + const auto posting_count = window.term_sizes_[term]; + const auto code_size = this->GetTermValueCodeSize(); + auto& ids = *window.term_ids_[term]; + auto& data = *window.term_datas_[term]; + normalization_ids_scratch_.resize(posting_count); + normalization_data_scratch_.resize(static_cast(posting_count) * code_size); + + uint32_t prefix = 0; + uint32_t suffix = prefix_count; + uint32_t output = 0; + while (prefix < prefix_count && suffix < posting_count) { + const auto source = PostingEntryBefore(window, term, prefix, suffix) ? prefix++ : suffix++; + normalization_ids_scratch_[output] = ids[source]; + std::memcpy(normalization_data_scratch_.data() + static_cast(output) * code_size, + data.data() + static_cast(source) * code_size, + code_size); + ++output; + } + while (prefix < prefix_count) { + normalization_ids_scratch_[output] = ids[prefix]; + std::memcpy(normalization_data_scratch_.data() + static_cast(output) * code_size, + data.data() + static_cast(prefix) * code_size, + code_size); + ++prefix; + ++output; + } + while (suffix < posting_count) { + normalization_ids_scratch_[output] = ids[suffix]; + std::memcpy(normalization_data_scratch_.data() + static_cast(output) * code_size, + data.data() + static_cast(suffix) * code_size, + code_size); + ++suffix; + ++output; + } + + std::copy(normalization_ids_scratch_.begin(), normalization_ids_scratch_.end(), ids.begin()); + std::copy(normalization_data_scratch_.begin(), normalization_data_scratch_.end(), data.begin()); + window.dirty_posting_prefixes_.erase(term); +} + +bool +MutableSindiTermDataCell::FinalizeInsertBatch(uint32_t window_id) { + CHECK_ARGUMENT(window_id < windows_.size(), "mutable SINDI window id out of range"); + return this->FinalizeInsertBatch(windows_[window_id], true); +} + +bool +MutableSindiTermDataCell::FinalizeInsertBatch(MutableSINDIWindow& window, bool release_scratch) { + if (window.dirty_posting_prefixes_.empty()) { + return false; + } + + struct posting_sort_scratch { + explicit posting_sort_scratch(Allocator* allocator) + : order(allocator), sorted_ids(allocator), sorted_data(allocator) { + } + + Vector order; + Vector sorted_ids; + Vector sorted_data; + }; + std::optional scratch; + bool posting_state_changed = false; + const auto code_size = this->GetTermValueCodeSize(); + for (uint32_t term : window.pending_posting_terms_) { + window.pending_posting_flags_[term] = 0; + const auto dirty = window.dirty_posting_prefixes_.find(term); + if (dirty == window.dirty_posting_prefixes_.end()) { + continue; + } + if (!scratch.has_value()) { + scratch.emplace(allocator_); + } + auto& sort_scratch = scratch.value(); + posting_state_changed = true; + const auto prefix_count = dirty->second; + const auto suffix_count = window.term_sizes_[term] - prefix_count; + auto& ids = *window.term_ids_[term]; + auto& data = *window.term_datas_[term]; + sindi_datacell_utils::SortPostingListByValue( + ids.data() + prefix_count, + data.data() + static_cast(prefix_count) * code_size, + suffix_count, + sparse_value_quant_type_, + sort_scratch.order, + sort_scratch.sorted_ids, + sort_scratch.sorted_data); + + if (prefix_count == 0) { + window.dirty_posting_prefixes_.erase(dirty); + continue; + } + const auto normalize_threshold = + std::clamp(prefix_count, MIN_DIRTY_POSTING_SIZE, MAX_DIRTY_POSTING_SIZE); + if (suffix_count >= normalize_threshold) { + this->NormalizePosting(window, term); + } + } + window.pending_posting_terms_.clear(); + if (window.dirty_posting_prefixes_.empty()) { + window.dirty_posting_prefixes_.rehash(0); + window.postings_sorted_ = true; + } else { + window.postings_sorted_ = false; + } + if (release_scratch) { + this->ReleaseNormalizationScratch(); + } + return posting_state_changed; +} + +void +MutableSindiTermDataCell::ReleaseNormalizationScratch() { + Vector empty_ids(allocator_); + Vector empty_data(allocator_); + normalization_ids_scratch_.swap(empty_ids); + normalization_data_scratch_.swap(empty_data); +} + +bool +MutableSindiTermDataCell::NormalizeDirtyPostings(uint32_t window_id) { + CHECK_ARGUMENT(window_id < windows_.size(), "mutable SINDI window id out of range"); + return this->NormalizeDirtyPostings(windows_[window_id], true); +} + +bool +MutableSindiTermDataCell::NormalizeDirtyPostings() { + bool posting_state_changed = false; + for (auto& window : windows_) { + posting_state_changed |= this->NormalizeDirtyPostings(window, false); + } + this->ReleaseNormalizationScratch(); + return posting_state_changed; +} + +bool +MutableSindiTermDataCell::NormalizeDirtyPostings(MutableSINDIWindow& window, bool release_scratch) { + const auto finalized = this->FinalizeInsertBatch(window, false); + if (window.dirty_posting_prefixes_.empty()) { + if (release_scratch) { + this->ReleaseNormalizationScratch(); + } + return finalized; + } + + Vector dirty_terms(allocator_); + dirty_terms.reserve(window.dirty_posting_prefixes_.size()); + for (const auto& item : window.dirty_posting_prefixes_) { + dirty_terms.push_back(item.first); + } + for (uint32_t term : dirty_terms) { + this->NormalizePosting(window, term); + } + window.dirty_posting_prefixes_.clear(); + window.dirty_posting_prefixes_.rehash(0); window.postings_sorted_ = true; + if (release_scratch) { + this->ReleaseNormalizationScratch(); + } + return true; } void @@ -729,14 +1063,19 @@ MutableSindiTermDataCell::ResizeTermList(MutableSINDIWindow& window, Vector>> new_ids(new_capacity, allocator_); Vector>> new_datas(new_capacity, allocator_); Vector new_sizes(new_capacity, 0, allocator_); + Vector new_pending_flags(new_capacity, 0, allocator_); std::move(window.term_ids_.begin(), window.term_ids_.end(), new_ids.begin()); std::move(window.term_datas_.begin(), window.term_datas_.end(), new_datas.begin()); std::copy(window.term_sizes_.begin(), window.term_sizes_.end(), new_sizes.begin()); + std::copy(window.pending_posting_flags_.begin(), + window.pending_posting_flags_.end(), + new_pending_flags.begin()); window.term_ids_.swap(new_ids); window.term_datas_.swap(new_datas); window.term_sizes_.swap(new_sizes); + window.pending_posting_flags_.swap(new_pending_flags); window.term_capacity_ = new_capacity; } @@ -747,6 +1086,7 @@ MutableSindiTermDataCell::Compact() { void MutableSindiTermDataCell::Compact(MutableSINDIWindow& window) { + this->NormalizeDirtyPostings(window, true); uint32_t compact_term_capacity = 0; const uint64_t compactable_capacity = std::min(std::min(static_cast(window.term_capacity_), @@ -762,6 +1102,7 @@ MutableSindiTermDataCell::Compact(MutableSINDIWindow& window) { Vector>> compact_ids(compact_term_capacity, allocator_); Vector>> compact_datas(compact_term_capacity, allocator_); Vector compact_sizes(compact_term_capacity, 0, allocator_); + Vector compact_pending_flags(compact_term_capacity, 0, allocator_); for (uint32_t i = 0; i < compact_term_capacity; ++i) { compact_sizes[i] = window.term_sizes_[i]; if (window.term_sizes_[i] != 0) { @@ -777,6 +1118,7 @@ MutableSindiTermDataCell::Compact(MutableSINDIWindow& window) { window.term_ids_.swap(compact_ids); window.term_datas_.swap(compact_datas); window.term_sizes_.swap(compact_sizes); + window.pending_posting_flags_.swap(compact_pending_flags); window.term_capacity_ = compact_term_capacity; } @@ -855,6 +1197,8 @@ MutableSindiTermDataCell::GetMemoryUsage() const { memory += MutableSindiTermDataCell::GetWindowMemoryUsage(window); } memory += sizeof(QuantizationParams); + memory += normalization_ids_scratch_.capacity() * sizeof(uint16_t); + memory += normalization_data_scratch_.capacity() * sizeof(uint8_t); return static_cast(memory); } @@ -876,6 +1220,10 @@ MutableSindiTermDataCell::GetWindowMemoryUsage(const MutableSINDIWindow& window) } } memory += window.term_sizes_.capacity() * sizeof(uint32_t); + memory += window.dirty_posting_prefixes_.bucket_count() * sizeof(void*); + memory += window.dirty_posting_prefixes_.size() * sizeof(std::pair); + memory += window.pending_posting_terms_.capacity() * sizeof(uint32_t); + memory += window.pending_posting_flags_.capacity() * sizeof(uint8_t); return memory; } diff --git a/src/datacell/mutable_sindi_term_datacell.h b/src/datacell/mutable_sindi_term_datacell.h index c970d27160..c14818b0db 100644 --- a/src/datacell/mutable_sindi_term_datacell.h +++ b/src/datacell/mutable_sindi_term_datacell.h @@ -20,6 +20,7 @@ #include "algorithm/sindi/sindi_parameter.h" #include "datacell/sindi_datacell_utils.h" #include "datacell/sindi_search_term_datacell.h" +#include "hash_types.h" #include "impl/searcher/basic_searcher.h" #include "quantization/sparse_quantization/sparse_term_computer.h" #include "storage/stream_reader.h" @@ -32,13 +33,21 @@ namespace vsag { struct MutableSINDIWindow { explicit MutableSINDIWindow(Allocator* allocator = nullptr) - : term_ids_(allocator), term_datas_(allocator), term_sizes_(allocator) { + : term_ids_(allocator), + term_datas_(allocator), + term_sizes_(allocator), + dirty_posting_prefixes_(allocator), + pending_posting_terms_(allocator), + pending_posting_flags_(allocator) { } uint32_t term_capacity_{0}; Vector>> term_ids_; Vector>> term_datas_; Vector term_sizes_; + UnorderedMap dirty_posting_prefixes_; + Vector pending_posting_terms_; + Vector pending_posting_flags_; bool postings_sorted_{true}; }; @@ -60,7 +69,9 @@ class MutableSindiTermDataCell : public SindiSearchTermDataCell { sparse_value_quant_type_(sparse_value_quant_type), quantization_params_(std::move(quantization_params)), window_size_(window_size), - windows_(allocator) { + windows_(allocator), + normalization_ids_scratch_(allocator), + normalization_data_scratch_(allocator) { } void @@ -166,6 +177,15 @@ class MutableSindiTermDataCell : public SindiSearchTermDataCell { void SortByValue(uint32_t window_id); + bool + FinalizeInsertBatch(uint32_t window_id); + + bool + NormalizeDirtyPostings(uint32_t window_id); + + bool + NormalizeDirtyPostings(); + void ResizeWindowCount(uint32_t window_count); @@ -206,6 +226,16 @@ class MutableSindiTermDataCell : public SindiSearchTermDataCell { GetTermDictCount() const override; private: + struct PostingRunSelection { + uint32_t prefix_count{0}; + uint32_t suffix_offset{0}; + uint32_t suffix_count{0}; + }; + + static constexpr uint32_t MIN_DIRTY_POSTING_SIZE = 32; + + static constexpr uint32_t MAX_DIRTY_POSTING_SIZE = 256; + [[nodiscard]] uint32_t GetTermValueCodeSize() const; @@ -219,7 +249,70 @@ class MutableSindiTermDataCell : public SindiSearchTermDataCell { ResizeTermList(MutableSINDIWindow& window, InnerIdType new_term_capacity) const; void - SortByValue(MutableSINDIWindow& window) const; + SortByValue(MutableSINDIWindow& window); + + [[nodiscard]] PostingRunSelection + SelectPostingRuns(const MutableSINDIWindow& window, + uint32_t term, + uint32_t retained_count) const; + + [[nodiscard]] bool + PostingEntryBefore(const MutableSINDIWindow& window, + uint32_t term, + uint32_t left, + uint32_t right) const; + + void + NormalizePosting(MutableSINDIWindow& window, uint32_t term); + + bool + FinalizeInsertBatch(MutableSINDIWindow& window, bool release_scratch); + + bool + NormalizeDirtyPostings(MutableSINDIWindow& window, bool release_scratch); + + void + ReleaseNormalizationScratch(); + + void + ScanPostingRange(uint32_t term_iterator, + const SparseTermComputerPtr& computer, + const uint16_t* term_ids, + const uint8_t* term_data, + uint32_t term_count, + float* dists, + SparseEvaluationTracker* evaluation_tracker) const; + + template + bool + ForEachSelectedPosting(const MutableSINDIWindow& window, + uint32_t term, + const PostingRunSelection& selected, + Callback&& callback) const { + uint32_t prefix = 0; + uint32_t suffix = selected.suffix_offset; + const auto prefix_end = selected.prefix_count; + const auto suffix_end = selected.suffix_offset + selected.suffix_count; + const auto& ids = *window.term_ids_[term]; + while (prefix < prefix_end && suffix < suffix_end) { + const auto posting = + PostingEntryBefore(window, term, prefix, suffix) ? prefix++ : suffix++; + if (!callback(ids[posting])) { + return false; + } + } + while (prefix < prefix_end) { + if (!callback(ids[prefix++])) { + return false; + } + } + while (suffix < suffix_end) { + if (!callback(ids[suffix++])) { + return false; + } + } + return true; + } void Compact(MutableSINDIWindow& window); @@ -294,5 +387,9 @@ class MutableSindiTermDataCell : public SindiSearchTermDataCell { uint32_t window_size_{0}; Vector windows_; + + Vector normalization_ids_scratch_; + + Vector normalization_data_scratch_; }; } // namespace vsag diff --git a/src/datacell/mutable_sindi_term_datacell_test.cpp b/src/datacell/mutable_sindi_term_datacell_test.cpp index dd55a43f38..17b768e1b3 100644 --- a/src/datacell/mutable_sindi_term_datacell_test.cpp +++ b/src/datacell/mutable_sindi_term_datacell_test.cpp @@ -27,6 +27,20 @@ using namespace vsag; namespace { +class ModuloFilter final : public Filter { +public: + explicit ModuloFilter(uint32_t divisor) : divisor_(divisor) { + } + + bool + CheckValid(int64_t id) const override { + return id % divisor_ == 0; + } + +private: + uint32_t divisor_; +}; + uint64_t QueryFirstWindow(const MutableSindiTermDataCellPtr& data_cell, float* dists, @@ -37,6 +51,57 @@ QueryFirstWindow(const MutableSindiTermDataCellPtr& data_cell, return query_context.evaluation_tracker.Count(); } +template +void +RequireHeapEquivalence(const MutableSindiTermDataCellPtr& data_cell, + const SparseTermComputerPtr& computer, + uint32_t posting_count, + const FilterPtr& filter, + Allocator* allocator) { + std::vector term_list_distances(posting_count, 0.0F); + std::vector full_scan_distances(posting_count, 0.0F); + QueryFirstWindow(data_cell, term_list_distances.data(), computer, allocator); + QueryFirstWindow(data_cell, full_scan_distances.data(), computer, allocator); + + InnerSearchParam search_param; + search_param.ef = 20; + search_param.radius = 0.99F; + search_param.is_inner_id_allowed = filter; + MaxHeap term_list_heap(allocator); + MaxHeap full_scan_heap(allocator); + REQUIRE_FALSE(data_cell->InsertHeapByTermLists( + term_list_distances.data(), computer, term_list_heap, search_param, 0)); + REQUIRE_FALSE(data_cell->InsertHeapByDists( + full_scan_distances.data(), full_scan_distances.size(), full_scan_heap, search_param, 0)); + + std::vector> term_list_results; + std::vector> full_scan_results; + while (!term_list_heap.empty()) { + term_list_results.push_back(term_list_heap.top()); + term_list_heap.pop(); + } + while (!full_scan_heap.empty()) { + full_scan_results.push_back(full_scan_heap.top()); + full_scan_heap.pop(); + } + const auto by_distance_and_id = [](const auto& left, const auto& right) { + return left.first == right.first ? left.second < right.second : left.first < right.first; + }; + std::sort(term_list_results.begin(), term_list_results.end(), by_distance_and_id); + std::sort(full_scan_results.begin(), full_scan_results.end(), by_distance_and_id); + REQUIRE(term_list_results.size() == full_scan_results.size()); + for (uint32_t i = 0; i < term_list_results.size(); ++i) { + REQUIRE(term_list_results[i].first == full_scan_results[i].first); + const auto distance = term_list_results[i].first; + const bool unique_distance = + (i == 0 || term_list_results[i - 1].first != distance) && + (i + 1 == term_list_results.size() || term_list_results[i + 1].first != distance); + if (unique_distance) { + REQUIRE(term_list_results[i].second == full_scan_results[i].second); + } + } +} + } // namespace TEST_CASE("MutableSindiTermDataCell uses caller document id coordinates", @@ -125,7 +190,8 @@ TEST_CASE("MutableSindiTermDataCell sorts postings by stored value", } } -TEST_CASE("MutableSindiTermDataCell prunes sorted postings", "[ut][MutableSindiTermDataCell]") { +TEST_CASE("MutableSindiTermDataCell prunes incrementally appended postings", + "[ut][MutableSindiTermDataCell]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); auto quantization_params = std::make_shared(); auto data_cell = std::make_shared( @@ -144,9 +210,6 @@ TEST_CASE("MutableSindiTermDataCell prunes sorted postings", "[ut][MutableSindiT REQUIRE(data_cell->GetWindow(0).term_sizes_[term] == 4); REQUIRE_FALSE(data_cell->GetWindow(0).postings_sorted_); - data_cell->SortByValue(0); - REQUIRE(data_cell->GetWindow(0).postings_sorted_); - float query_value = 1.0F; SparseVector query{1, &term, &query_value}; SINDISearchParameter search_parameter; @@ -160,6 +223,130 @@ TEST_CASE("MutableSindiTermDataCell prunes sorted postings", "[ut][MutableSindiT REQUIRE(sorted_distances[2] == 0.0F); } +TEST_CASE("MutableSindiTermDataCell two-run search matches eager normalization", + "[ut][MutableSindiTermDataCell]") { + const auto quantization = GENERATE(SparseValueQuantizationType::FP32, + SparseValueQuantizationType::FP16, + SparseValueQuantizationType::SQ8); + DYNAMIC_SECTION("quantization=" << static_cast(quantization)) { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + auto quantization_params = std::make_shared(); + quantization_params->min_val = 0.0F; + quantization_params->max_val = 1024.0F; + quantization_params->diff = 1024.0F; + uint32_t term = 9; + constexpr uint32_t prefix_count = 300; + constexpr uint32_t posting_count = 340; + auto split = std::make_shared( + 16, 512, allocator.get(), quantization, quantization_params); + auto eager = std::make_shared( + 16, 512, allocator.get(), quantization, quantization_params); + + std::array values{}; + for (uint32_t id = 0; id < posting_count; ++id) { + values[id] = static_cast((id * 43U) % 401U + 25U); + if (id == prefix_count) { + values[id] = 1000.0F; + } + SparseVector vector{1, &term, values.data() + id}; + split->InsertVector(vector, id); + eager->InsertVector(vector, id); + if (id + 1 == prefix_count) { + split->SortByValue(0); + } + } + REQUIRE(split->FinalizeInsertBatch(0)); + eager->SortByValue(0); + REQUIRE(split->GetWindow(0).dirty_posting_prefixes_.at(term) == prefix_count); + + const std::array retained_counts{1, 31, 32, 255, 256, posting_count}; + for (const auto retained_count : retained_counts) { + CAPTURE(quantization, retained_count); + float query_value = 1.0F; + SparseVector query{1, &term, &query_value}; + SINDISearchParameter search_parameter; + search_parameter.term_retain_threshold = retained_count; + auto split_computer = + std::make_shared(query, search_parameter, allocator.get()); + auto eager_computer = + std::make_shared(query, search_parameter, allocator.get()); + std::array split_distances{}; + std::array eager_distances{}; + REQUIRE( + QueryFirstWindow(split, split_distances.data(), split_computer, allocator.get()) == + retained_count); + REQUIRE( + QueryFirstWindow(eager, eager_distances.data(), eager_computer, allocator.get()) == + retained_count); + REQUIRE(split_distances == eager_distances); + } + + float query_value = 1.0F; + SparseVector query{1, &term, &query_value}; + SINDISearchParameter search_parameter; + search_parameter.term_retain_threshold = 256; + auto computer = + std::make_shared(query, search_parameter, allocator.get()); + RequireHeapEquivalence( + split, computer, posting_count, nullptr, allocator.get()); + RequireHeapEquivalence( + split, computer, posting_count, nullptr, allocator.get()); + auto filter = std::make_shared(3); + RequireHeapEquivalence( + split, computer, posting_count, filter, allocator.get()); + RequireHeapEquivalence( + split, computer, posting_count, filter, allocator.get()); + + REQUIRE(split->NormalizeDirtyPostings(0)); + REQUIRE_FALSE(split->NormalizeDirtyPostings(0)); + REQUIRE(*split->GetWindow(0).term_ids_[term] == *eager->GetWindow(0).term_ids_[term]); + REQUIRE(*split->GetWindow(0).term_datas_[term] == *eager->GetWindow(0).term_datas_[term]); + REQUIRE(split->normalization_ids_scratch_.capacity() == 0); + REQUIRE(split->normalization_data_scratch_.capacity() == 0); + } +} + +TEST_CASE("MutableSindiTermDataCell bounds dirty posting runs", "[ut][MutableSindiTermDataCell]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + auto data_cell = std::make_shared( + 16, 1024, allocator.get(), SparseValueQuantizationType::FP32, nullptr); + uint32_t term = 7; + std::array values{}; + for (uint32_t id = 0; id < 300; ++id) { + values[id] = static_cast((id * 37U) % 113U + 1U); + SparseVector vector{1, &term, values.data() + id}; + data_cell->InsertVector(vector, id); + } + REQUIRE(data_cell->FinalizeInsertBatch(0)); + REQUIRE(data_cell->GetWindow(0).dirty_posting_prefixes_.empty()); + + for (uint32_t id = 300; id < 555; ++id) { + values[id] = static_cast((id * 37U) % 113U + 1U); + SparseVector vector{1, &term, values.data() + id}; + data_cell->InsertVector(vector, id); + REQUIRE(data_cell->FinalizeInsertBatch(0)); + } + REQUIRE(data_cell->GetWindow(0).dirty_posting_prefixes_.at(term) == 300); + + values[555] = 1000.0F; + SparseVector vector{1, &term, values.data() + 555}; + data_cell->InsertVector(vector, 555); + REQUIRE(data_cell->FinalizeInsertBatch(0)); + REQUIRE(data_cell->GetWindow(0).dirty_posting_prefixes_.empty()); + REQUIRE(data_cell->GetWindow(0).postings_sorted_); + + values[100] = 2000.0F; + SparseVector appended{1, &term, values.data() + 100}; + data_cell->InsertVector(appended, 556); + REQUIRE(data_cell->FinalizeInsertBatch(0)); + REQUIRE_FALSE(data_cell->GetWindow(0).dirty_posting_prefixes_.empty()); + data_cell->Compact(); + REQUIRE(data_cell->GetWindow(0).dirty_posting_prefixes_.empty()); + REQUIRE(data_cell->GetWindow(0).postings_sorted_); + REQUIRE(data_cell->normalization_ids_scratch_.capacity() == 0); + REQUIRE(data_cell->normalization_data_scratch_.capacity() == 0); +} + TEST_CASE("MutableSindiTermDataCell normalizes legacy posting order on deserialize", "[ut][MutableSindiTermDataCell]") { auto allocator = SafeAllocator::FactoryDefaultAllocator();