diff --git a/include/svs/concurrent/dynamic_index.h b/include/svs/concurrent/dynamic_index.h index 8c260b1ca..76921d318 100644 --- a/include/svs/concurrent/dynamic_index.h +++ b/include/svs/concurrent/dynamic_index.h @@ -1746,11 +1746,41 @@ class MutableVamanaIndex { const std::filesystem::path& graph_directory, const std::filesystem::path& data_directory ) { - // Post-consolidation, all entries should be "valid". - // Therefore, we don't need to save the slot metadata. - consolidate(); - compact(); + repair_if_needed(); + save_impl(config_directory, graph_directory, data_directory); + } + + void save(std::ostream& os) { + repair_if_needed(); + save_impl(os); + } + + private: + bool needs_repair() const { + for (size_t i = 0, imax = status_.size(); i < imax; ++i) { + if (std::atomic_ref(const_cast(status_[i])) + .load(std::memory_order_acquire) != SlotMetadata::Valid) { + return true; + } + } + return false; + } + + /// Skipped when every slot is Valid to avoid a buffer roundtrip + void repair_if_needed() { + if (needs_repair()) { + consolidate(); + compact(); + } + } + /// Requires every slot Valid; slot metadata is not saved, so any other state + /// makes the loader reject the index. + void save_impl( + const std::filesystem::path& config_directory, + const std::filesystem::path& graph_directory, + const std::filesystem::path& data_directory + ) const { // Save auxiliary data structures. lib::save_to_disk( lib::SaveOverride([&](const lib::SaveContext& ctx) { @@ -1774,12 +1804,9 @@ class MutableVamanaIndex { lib::save_to_disk(graph_, graph_directory); } - void save(std::ostream& os) { - // Post-consolidation, all entries should be "valid". - // Therefore, we don't need to save the slot metadata. - consolidate(); - compact(); - + /// Requires every slot Valid; slot metadata is not saved, so any other state + /// makes the loader reject the index. + void save_impl(std::ostream& os) const { lib::begin_serialization(os); auto save_table = lib::SaveTable( "vamana_dynamic_auxiliary_parameters", @@ -1799,6 +1826,7 @@ class MutableVamanaIndex { lib::save_to_stream(graph_, os); } + public: ///// ///// Calibrate ///// diff --git a/include/svs/concurrent/multi.h b/include/svs/concurrent/multi.h index 5e00fc56f..8284ca1a3 100644 --- a/include/svs/concurrent/multi.h +++ b/include/svs/concurrent/multi.h @@ -779,11 +779,31 @@ class MultiMutableVamanaIndex { const std::filesystem::path& graph_directory, const std::filesystem::path& data_directory ) { - // Post-consolidation, all entries should be "valid". - // Therefore, we don't need to save the slot metadata. - consolidate(); - compact(); + repair_if_needed(); + save_impl(config_directory, graph_directory, data_directory); + } + + void save(std::ostream& os) { + repair_if_needed(); + save_impl(os); + } + + private: + /// Skipped when the parent index needs no repair to avoid a buffer roundtrip + void repair_if_needed() { + if (index_->needs_repair()) { + consolidate(); + compact(); + } + } + /// Requires every slot to be Valid: slot metadata is not saved, so a Deleted or + /// Empty slot makes the loader reject the index. + void save_impl( + const std::filesystem::path& config_directory, + const std::filesystem::path& graph_directory, + const std::filesystem::path& data_directory + ) const { auto labels = get_labels(); size_t num_labels = labels.size(); @@ -815,10 +835,9 @@ class MultiMutableVamanaIndex { lib::save_to_disk(index_->graph_, graph_directory); } - void save(std::ostream& os) { - consolidate(); - compact(); - + /// Requires every slot to be Valid: slot metadata is not saved, so a Deleted or + /// Empty slot makes the loader reject the index. + void save_impl(std::ostream& os) const { auto labels = get_labels(); size_t num_labels = labels.size(); diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 9336f82e1..e0d7c4c35 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -1130,11 +1130,37 @@ class MutableVamanaIndex { const std::filesystem::path& graph_directory, const std::filesystem::path& data_directory ) { - // Post-consolidation, all entries should be "valid". - // Therefore, we don't need to save the slot metadata. - consolidate(); - compact(); + repair_if_needed(); + save_impl(config_directory, graph_directory, data_directory); + } + + void save(std::ostream& os) { + repair_if_needed(); + save_impl(os); + } + + private: + bool needs_repair() const { + return std::any_of(status_.begin(), status_.end(), [](SlotMetadata status) { + return status != SlotMetadata::Valid; + }); + } + + /// Skipped when every slot is Valid to avoid a buffer roundtrip + void repair_if_needed() { + if (needs_repair()) { + consolidate(); + compact(); + } + } + /// Requires every slot to be Valid: slot metadata is not saved, so a Deleted or + /// Empty slot makes the loader reject the index. + void save_impl( + const std::filesystem::path& config_directory, + const std::filesystem::path& graph_directory, + const std::filesystem::path& data_directory + ) const { // Save auxiliary data structures. lib::save_to_disk( lib::SaveOverride([&](const lib::SaveContext& ctx) { @@ -1158,12 +1184,9 @@ class MutableVamanaIndex { lib::save_to_disk(graph_, graph_directory); } - void save(std::ostream& os) { - // Post-consolidation, all entries should be "valid". - // Therefore, we don't need to save the slot metadata. - consolidate(); - compact(); - + /// Requires every slot to be Valid: slot metadata is not saved, so a Deleted or + /// Empty slot makes the loader reject the index. + void save_impl(std::ostream& os) const { lib::begin_serialization(os); auto save_table = lib::SaveTable( "vamana_dynamic_auxiliary_parameters", @@ -1183,6 +1206,7 @@ class MutableVamanaIndex { lib::save_to_stream(graph_, os); } + public: ///// ///// Calibrate ///// diff --git a/include/svs/index/vamana/multi.h b/include/svs/index/vamana/multi.h index 277771fb9..3408971bd 100644 --- a/include/svs/index/vamana/multi.h +++ b/include/svs/index/vamana/multi.h @@ -637,11 +637,31 @@ class MultiMutableVamanaIndex { const std::filesystem::path& graph_directory, const std::filesystem::path& data_directory ) { - // Post-consolidation, all entries should be "valid". - // Therefore, we don't need to save the slot metadata. - consolidate(); - compact(); + repair_if_needed(); + save_impl(config_directory, graph_directory, data_directory); + } + + void save(std::ostream& os) { + repair_if_needed(); + save_impl(os); + } + + private: + /// Skipped when the parent index needs no repair to avoid a buffer roundtrip + void repair_if_needed() { + if (index_->needs_repair()) { + consolidate(); + compact(); + } + } + /// Requires every slot to be Valid: slot metadata is not saved, so a Deleted or + /// Empty slot makes the loader reject the index. + void save_impl( + const std::filesystem::path& config_directory, + const std::filesystem::path& graph_directory, + const std::filesystem::path& data_directory + ) const { auto labels = get_labels(); size_t num_labels = labels.size(); @@ -673,10 +693,9 @@ class MultiMutableVamanaIndex { lib::save_to_disk(index_->graph_, graph_directory); } - void save(std::ostream& os) { - consolidate(); - compact(); - + /// Requires every slot to be Valid: slot metadata is not saved, so a Deleted or + /// Empty slot makes the loader reject the index. + void save_impl(std::ostream& os) const { auto labels = get_labels(); size_t num_labels = labels.size(); diff --git a/tests/svs/index/vamana/dynamic_index_2.cpp b/tests/svs/index/vamana/dynamic_index_2.cpp index e590ae547..d35dd58c6 100644 --- a/tests/svs/index/vamana/dynamic_index_2.cpp +++ b/tests/svs/index/vamana/dynamic_index_2.cpp @@ -22,6 +22,7 @@ #include "svs/lib/float16.h" #include "svs/lib/preprocessor.h" #include "svs/lib/timing.h" +#include "svs/orchestrators/dynamic_vamana.h" #include "svs/misc/dynamic_helper.h" @@ -675,3 +676,79 @@ CATCH_TEST_CASE("Dynamic Vamana Index Default Parameters", "[parameter][vamana]" ); } } + +CATCH_TEST_CASE( + "Dynamic Vamana Index Stream Save Round Trip", + "[graph_index][dynamic_index][saveload]" +) { + const size_t num_threads = 2; + using Distance = svs::distance::DistanceL2; + using Data_t = svs::data::BlockedData; + + auto data = test_dataset::data_blocked_f32(); + std::vector indices(data.size()); + std::iota(indices.begin(), indices.end(), 0); + + svs::index::vamana::VamanaBuildParameters parameters{1.2, 64, 10, 20, 10, true}; + auto index = svs::index::vamana::MutableVamanaIndex( + parameters, std::move(data), indices, Distance(), num_threads + ); + + const size_t num_neighbors = 10; + auto queries = test_dataset::queries(); + auto search_params = svs::index::vamana::VamanaSearchParameters{}; + search_params.buffer_config_ = svs::index::vamana::SearchBufferConfig{num_neighbors}; + + CATCH_SECTION("No deletions: stream save reloads identical results") { + auto results = svs::QueryResult(queries.size(), num_neighbors); + index.search(results.view(), queries.cview(), search_params); + + std::stringstream stream; + index.save(stream); + auto loaded = + svs::DynamicVamana::assemble(stream, Distance(), num_threads); + + CATCH_REQUIRE(loaded.size() == index.size()); + + auto loaded_results = svs::QueryResult(queries.size(), num_neighbors); + loaded.search(loaded_results.view(), queries.cview(), search_params); + for (size_t q = 0; q < queries.size(); ++q) { + for (size_t i = 0; i < num_neighbors; ++i) { + CATCH_REQUIRE(loaded_results.index(q, i) == results.index(q, i)); + } + } + + std::vector sample_ids(indices.begin(), indices.begin() + 5); + auto original_vectors = + svs::data::SimpleData(sample_ids.size(), index.dimensions()); + auto reloaded_vectors = + svs::data::SimpleData(sample_ids.size(), index.dimensions()); + index.reconstruct_at( + original_vectors.view(), std::span(sample_ids) + ); + loaded.reconstruct_at(reloaded_vectors.view(), sample_ids); + for (size_t i = 0; i < sample_ids.size(); ++i) { + auto a = original_vectors.get_datum(i); + auto b = reloaded_vectors.get_datum(i); + CATCH_REQUIRE(std::equal(a.begin(), a.end(), b.begin())); + } + } + + // The loader rejects Deleted or Empty slots, so this reload throws if + // repair_if_needed() wrongly skips the repair. + CATCH_SECTION("Deletions: stream save repairs and drops deleted ids") { + std::vector ids_to_delete(indices.begin(), indices.begin() + 3); + index.delete_entries(ids_to_delete); + size_t live_count = index.size(); + + std::stringstream stream; + index.save(stream); + auto loaded = + svs::DynamicVamana::assemble(stream, Distance(), num_threads); + + CATCH_REQUIRE(loaded.size() == live_count); + for (auto id : ids_to_delete) { + CATCH_REQUIRE(!loaded.has_id(id)); + } + } +}