MDD-53: Add build_with_compression API for higher quality graph construction

bindings/python/src/vamana.cpp

@@ -123,6 +123,58 @@ void register_vamana_build_from_file(Dispatcher& dispatcher) {
 
 using VamanaBuildTypes = std::variant<UnspecializedVectorDataLoader>;
 
+template <typename Q, typename T, size_t N>
+svs::Vamana build_with_compression_uncompressed(
+    const svs::index::vamana::VamanaBuildParameters& parameters,
+    svs::VectorDataLoader<T, N, RebindAllocator<T>> 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<Q>(
+        parameters, std::move(data), compression_params, distance_type, num_threads
+    );
+}
+
+template <typename Dispatcher>
+void register_vamana_build_with_compression_from_file(Dispatcher& dispatcher) {
+    for_standard_specializations(
+        [&dispatcher]<typename Q, typename T, size_t N, EnableBuild B>() {
+            if constexpr (enable_build_from_file<B>) {
+                auto method = &build_with_compression_uncompressed<Q, T, N>;
+                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_<svs::CompressionType>(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

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 <cstdint>
+
+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

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 <algorithm>
 #include <concepts>
@@ -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<uint32_t>>
+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<I>(entry_point), logging::Level::Trace, logger);
+    builder.construct(
+        verified_parameters.alpha,
+        lib::narrow<I>(entry_point),
+        logging::Level::Trace,
+        logger
+    );
+
+    if (compression_params.enabled()) {
+        if (compression_params.type == CompressionType::ScalarInt8) {
+            auto compressed_data = quantization::scalar::SQDataset<std::int8_t>::compress(
+                uncompressed_data, threadpool
+            );
+
+            return VamanaIndex{
+                std::move(graph),
+                std::move(compressed_data),
+                lib::narrow<I>(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<I>(entry_point),
+        std::move(distance),
+        std::move(threadpool),
+        logger};
+}
+
 } // namespace svs::index::vamana

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<VamanaInterface> {
         }
     }
 
+    template <
+        manager::QueryTypeDefinition QueryTypes,
+        typename DataLoader,
+        typename Distance,
+        typename ThreadPoolProto = size_t,
+        typename Allocator = HugepageAllocator<uint32_t>>
+    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<std::decay_t<Distance>, DistanceType>) {
+            auto dispatcher = DistanceDispatcher(distance);
+            return dispatcher([&](auto distance_function) {
+                return make_vamana<manager::as_typelist<QueryTypes>>(
+                    index::vamana::auto_build_with_compression(
+                        parameters,
+                        std::forward<DataLoader>(data_loader),
+                        compression_params,
+                        std::move(distance_function),
+                        std::move(threadpool),
+                        graph_allocator
+                    )
+                );
+            });
+        } else {
+            return make_vamana<manager::as_typelist<QueryTypes>>(
+                index::vamana::auto_build_with_compression(
+                    parameters,
+                    std::forward<DataLoader>(data_loader),
+                    compression_params,
+                    distance,
+                    std::move(threadpool),
+                    graph_allocator
+                )
+            );
+        }
+    }
+
     ///
     /// @brief Return a new iterator (an instance of `svs::VamanaIterator`) for the query.
     ///

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 <filesystem>
+
+namespace {
+
+template <typename Distance> void test_build_with_compression(const Distance& distance) {
+    const double epsilon = 0.01;
+    const auto queries = svs::data::SimpleData<float>::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<Distance>, 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<float>(
+        expected_result.build_parameters_.value(),
+        svs::VectorDataLoader<float, 128>(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<Distance>);
+    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());
+}