From dfe678ab7df2fe421cebbeb14a53c31a775e0c61 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 28 Oct 2025 20:11:45 +0000 Subject: [PATCH] MDD-53: Add build_with_compression API for higher quality graph construction - Add CompressionType enum and CompressionParameters struct - Add auto_build_with_compression() that builds with uncompressed data then compresses - Add Vamana::build_with_compression() orchestrator method - Add Python bindings with dispatcher for runtime type selection - Add integration tests for build with compression - Builds graphs on uncompressed data for better quality, stores compressed for efficiency Co-Authored-By: milind@cognition.ai --- bindings/python/src/vamana.cpp | 82 +++++++++++++++++ include/svs/core/compression.h | 37 ++++++++ include/svs/index/vamana/index.h | 77 ++++++++++++++++ include/svs/orchestrators/vamana.h | 44 ++++++++++ .../integration/vamana/compression_build.cpp | 87 +++++++++++++++++++ 5 files changed, 327 insertions(+) create mode 100644 include/svs/core/compression.h create mode 100644 tests/integration/vamana/compression_build.cpp diff --git a/bindings/python/src/vamana.cpp b/bindings/python/src/vamana.cpp index 6b4c1c75..16c08939 100644 --- a/bindings/python/src/vamana.cpp +++ b/bindings/python/src/vamana.cpp @@ -123,6 +123,58 @@ void register_vamana_build_from_file(Dispatcher& dispatcher) { using VamanaBuildTypes = std::variant; +template +svs::Vamana build_with_compression_uncompressed( + const svs::index::vamana::VamanaBuildParameters& parameters, + svs::VectorDataLoader> data, + svs::CompressionType compression_type, + svs::DistanceType distance_type, + size_t num_threads +) { + auto compression_params = svs::CompressionParameters(compression_type); + return svs::Vamana::build_with_compression( + parameters, std::move(data), compression_params, distance_type, num_threads + ); +} + +template +void register_vamana_build_with_compression_from_file(Dispatcher& dispatcher) { + for_standard_specializations( + [&dispatcher]() { + if constexpr (enable_build_from_file) { + auto method = &build_with_compression_uncompressed; + dispatcher.register_target(svs::lib::dispatcher_build_docs, method); + } + } + ); +} + +using BuildWithCompressionDispatcher = svs::lib::Dispatcher< + svs::Vamana, + const svs::index::vamana::VamanaBuildParameters&, + VamanaBuildTypes, + svs::CompressionType, + svs::DistanceType, + size_t>; + +BuildWithCompressionDispatcher build_with_compression_dispatcher() { + auto dispatcher = BuildWithCompressionDispatcher{}; + register_vamana_build_with_compression_from_file(dispatcher); + return dispatcher; +} + +svs::Vamana build_with_compression_from_file( + const svs::index::vamana::VamanaBuildParameters& parameters, + VamanaBuildTypes data_source, + svs::CompressionType compression_type, + svs::DistanceType distance_type, + size_t num_threads +) { + return build_with_compression_dispatcher().invoke( + parameters, std::move(data_source), compression_type, distance_type, num_threads + ); +} + ///// ///// Build from Array ///// @@ -402,6 +454,36 @@ Method {}: py::arg("num_threads") = 1, fmt::format(docstring_proto, dynamic).c_str() ); + + py::enum_(m, "CompressionType") + .value("None", svs::CompressionType::None) + .value("ScalarInt8", svs::CompressionType::ScalarInt8); + + vamana.def_static( + "build_with_compression", + &detail::build_with_compression_from_file, + py::arg("build_parameters"), + py::arg("data_loader"), + py::arg("compression_type"), + py::arg("distance_type"), + py::arg("num_threads") = 1, + R"( +Build a Vamana index with optional compression. + +Builds the graph using uncompressed data for higher quality, then +compresses the data for storage in the final index. + +Args: + build_parameters: Build parameters for graph construction. + data_loader: Path to uncompressed vector data. + compression_type: Type of compression to apply (CompressionType.ScalarInt8 or CompressionType.None). + distance_type: Distance metric to use. + num_threads: Number of threads for parallel operations. + +Returns: + Built Vamana index with compressed data. + )" + ); } } // namespace detail diff --git a/include/svs/core/compression.h b/include/svs/core/compression.h new file mode 100644 index 00000000..cbdacc99 --- /dev/null +++ b/include/svs/core/compression.h @@ -0,0 +1,37 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace svs { + +enum class CompressionType { + None, + ScalarInt8, +}; + +struct CompressionParameters { + CompressionType type; + + explicit CompressionParameters(CompressionType type = CompressionType::None) + : type{type} {} + + bool enabled() const { return type != CompressionType::None; } +}; + +} // namespace svs diff --git a/include/svs/index/vamana/index.h b/include/svs/index/vamana/index.h index b7c13664..38bd472c 100644 --- a/include/svs/index/vamana/index.h +++ b/include/svs/index/vamana/index.h @@ -17,7 +17,9 @@ #pragma once // svs +#include "svs/core/compression.h" #include "svs/core/data.h" + #include "svs/core/graph.h" #include "svs/core/loading.h" #include "svs/core/medioid.h" @@ -34,6 +36,8 @@ #include "svs/lib/saveload.h" #include "svs/lib/threads.h" +#include "svs/quantization/scalar/scalar.h" + // stl #include #include @@ -1058,4 +1062,77 @@ void verify_and_set_default_index_parameters( throw std::invalid_argument("prune_to must be <= graph_max_degree"); } } + +template < + typename DataProto, + typename Distance, + typename ThreadPoolProto, + typename Allocator = HugepageAllocator> +auto auto_build_with_compression( + const VamanaBuildParameters& parameters, + DataProto data_proto, + const CompressionParameters& compression_params, + Distance distance, + ThreadPoolProto threadpool_proto, + const Allocator& graph_allocator = {}, + svs::logging::logger_ptr logger = svs::logging::get() +) { + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + + auto uncompressed_data = svs::detail::dispatch_load(std::move(data_proto), threadpool); + + auto entry_point = extensions::compute_entry_point(uncompressed_data, threadpool); + + auto verified_parameters = parameters; + verify_and_set_default_index_parameters(verified_parameters, distance); + auto graph = default_graph( + uncompressed_data.size(), verified_parameters.graph_max_degree, graph_allocator + ); + using I = typename decltype(graph)::index_type; + + auto builder = VamanaBuilder( + graph, + uncompressed_data, + distance, + verified_parameters, + threadpool, + extensions::estimate_prefetch_parameters(uncompressed_data), + logger + ); + + builder.construct(1.0F, lib::narrow(entry_point), logging::Level::Trace, logger); + builder.construct( + verified_parameters.alpha, + lib::narrow(entry_point), + logging::Level::Trace, + logger + ); + + if (compression_params.enabled()) { + if (compression_params.type == CompressionType::ScalarInt8) { + auto compressed_data = quantization::scalar::SQDataset::compress( + uncompressed_data, threadpool + ); + + return VamanaIndex{ + std::move(graph), + std::move(compressed_data), + lib::narrow(entry_point), + std::move(distance), + std::move(threadpool), + logger}; + } else { + throw ANNEXCEPTION("Unsupported compression type"); + } + } + + return VamanaIndex{ + std::move(graph), + std::move(uncompressed_data), + lib::narrow(entry_point), + std::move(distance), + std::move(threadpool), + logger}; +} + } // namespace svs::index::vamana diff --git a/include/svs/orchestrators/vamana.h b/include/svs/orchestrators/vamana.h index 6b698c4f..a801ea7b 100644 --- a/include/svs/orchestrators/vamana.h +++ b/include/svs/orchestrators/vamana.h @@ -21,6 +21,7 @@ /// @brief Main API for the Vamana type-erased orchestrator. /// +#include "svs/core/compression.h" #include "svs/core/data.h" #include "svs/core/distance.h" #include "svs/core/graph.h" @@ -567,6 +568,49 @@ class Vamana : public manager::IndexManager { } } + template < + manager::QueryTypeDefinition QueryTypes, + typename DataLoader, + typename Distance, + typename ThreadPoolProto = size_t, + typename Allocator = HugepageAllocator> + static Vamana build_with_compression( + const index::vamana::VamanaBuildParameters& parameters, + DataLoader&& data_loader, + const CompressionParameters& compression_params, + Distance distance, + ThreadPoolProto threadpool_proto = 1, + const Allocator& graph_allocator = {} + ) { + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + if constexpr (std::is_same_v, DistanceType>) { + auto dispatcher = DistanceDispatcher(distance); + return dispatcher([&](auto distance_function) { + return make_vamana>( + index::vamana::auto_build_with_compression( + parameters, + std::forward(data_loader), + compression_params, + std::move(distance_function), + std::move(threadpool), + graph_allocator + ) + ); + }); + } else { + return make_vamana>( + index::vamana::auto_build_with_compression( + parameters, + std::forward(data_loader), + compression_params, + distance, + std::move(threadpool), + graph_allocator + ) + ); + } + } + /// /// @brief Return a new iterator (an instance of `svs::VamanaIterator`) for the query. /// diff --git a/tests/integration/vamana/compression_build.cpp b/tests/integration/vamana/compression_build.cpp new file mode 100644 index 00000000..998b55b2 --- /dev/null +++ b/tests/integration/vamana/compression_build.cpp @@ -0,0 +1,87 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "svs/core/compression.h" +#include "svs/core/recall.h" +#include "svs/lib/timing.h" +#include "svs/orchestrators/vamana.h" + +#include "catch2/catch_test_macros.hpp" +#include "fmt/core.h" +#include "svs-benchmark/benchmark.h" + +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" +#include "tests/utils/vamana_reference.h" + +#include + +namespace { + +template void test_build_with_compression(const Distance& distance) { + const double epsilon = 0.01; + const auto queries = svs::data::SimpleData::load(test_dataset::query_file()); + CATCH_REQUIRE(svs_test::prepare_temp_directory()); + size_t num_threads = 2; + + auto expected_result = test_dataset::vamana::expected_build_results( + svs::distance_type_v, svsbenchmark::Uncompressed(svs::DataType::float32) + ); + + auto compression_params = svs::CompressionParameters(svs::CompressionType::ScalarInt8); + auto tic = svs::lib::now(); + auto index = svs::Vamana::build_with_compression( + expected_result.build_parameters_.value(), + svs::VectorDataLoader(test_dataset::data_svs_file()), + compression_params, + distance, + num_threads + ); + + fmt::print("Build with compression time: {}s\n", svs::lib::time_difference(tic)); + CATCH_REQUIRE(index.get_num_threads() == num_threads); + + auto groundtruth = test_dataset::load_groundtruth(svs::distance_type_v); + for (const auto& expected : expected_result.config_and_recall_) { + auto these_queries = test_dataset::get_test_set(queries, expected.num_queries_); + auto these_groundtruth = + test_dataset::get_test_set(groundtruth, expected.num_queries_); + index.set_search_parameters(expected.search_parameters_); + auto results = index.search(these_queries, expected.num_neighbors_); + double recall = svs::k_recall_at_n( + these_groundtruth, results, expected.num_neighbors_, expected.recall_k_ + ); + + fmt::print( + "Window Size: {}, Expected Recall: {}, Actual Recall: {}\n", + index.get_search_window_size(), + expected.recall_, + recall + ); + + CATCH_REQUIRE(recall > expected.recall_ - epsilon); + } +} + +} // namespace + +CATCH_TEST_CASE( + "Vamana Build With Compression", "[integration][build][vamana][compression]" +) { + test_build_with_compression(svs::DistanceL2()); + test_build_with_compression(svs::DistanceIP()); + test_build_with_compression(svs::DistanceCosineSimilarity()); +}