Add support for compressed vector search to IVF (#184)

- Add support for IVF search with SQData
- Add Python tests for IVF 
- Enable IVF tests in CI

Author: ibhati

.github/workflows/build-linux.yml

@@ -30,11 +30,12 @@ concurrency:
 
 jobs:
   build:
-    name: ${{ matrix.cxx }}, ${{ matrix.build_type }}
+    name: ${{ matrix.cxx }}, ${{ matrix.build_type }}, ivf=${{ matrix.ivf }}
     runs-on: ubuntu-22.04
     strategy:
       matrix:
         build_type: [RelWithDebugInfo]
+        ivf: [OFF, ON]
         cxx: [g++-11, g++-12, clang++-15]
         include:
           - cxx: g++-11
@@ -43,6 +44,9 @@ jobs:
             cc: gcc-12
           - cxx: clang++-15
             cc: clang-15
+        exclude:
+          - cxx: g++-12
+            ivf: ON
 
     steps:
     - uses: actions/checkout@v5
@@ -69,7 +73,8 @@ jobs:
               -DSVS_BUILD_TESTS=YES \
               -DSVS_BUILD_EXAMPLES=YES \
               -DSVS_EXPERIMENTAL_LEANVEC=YES \
-              -DSVS_NO_AVX512=NO
+              -DSVS_NO_AVX512=NO \
+              -DSVS_EXPERIMENTAL_ENABLE_IVF=${{ matrix.ivf }}
 
     - name: Build Tests and Utilities
       working-directory: ${{ runner.temp }}/build

benchmark/include/svs-benchmark/ivf/test.h

@@ -48,9 +48,6 @@ struct IVFTest {
     std::filesystem::path graph_;
     std::filesystem::path queries_f32_;
     size_t queries_in_training_set_;
-    // Backend-specific members
-    std::filesystem::path leanvec_data_matrix_;
-    std::filesystem::path leanvec_query_matrix_;
     // Runtime values
     size_t num_threads_;
 
@@ -62,9 +59,6 @@ struct IVFTest {
         std::filesystem::path graph,
         std::filesystem::path queries_f32,
         size_t queries_in_training_set,
-        // backend-specific members
-        std::filesystem::path leanvec_data_matrix,
-        std::filesystem::path leanvec_query_matrix,
         // Runtime values
         size_t num_threads
     )
@@ -74,8 +68,6 @@ struct IVFTest {
         , graph_{std::move(graph)}
         , queries_f32_{std::move(queries_f32)}
         , queries_in_training_set_{queries_in_training_set}
-        , leanvec_data_matrix_{std::move(leanvec_data_matrix)}
-        , leanvec_query_matrix_{std::move(leanvec_query_matrix)}
         , num_threads_{num_threads} {}
 
     static IVFTest example() {
@@ -86,8 +78,6 @@ struct IVFTest {
             "path/to/graph",                     // graph
             "path/to/queries_f32",               // queries_f32
             10000,                               // queries_in_training_set
-            "path/to/leanvec_data_matrix",       // LeanVec data matrix
-            "path/to/leanvec_query_matrix",      // LeanVec query matrix
             0,                                   // Num Threads (not-saved)
         };
     }
@@ -113,9 +103,7 @@ struct IVFTest {
              SVS_LIST_SAVE_(index_config),
              SVS_LIST_SAVE_(graph),
              SVS_LIST_SAVE_(queries_f32),
-             SVS_LIST_SAVE_(queries_in_training_set),
-             SVS_LIST_SAVE_(leanvec_data_matrix),
-             SVS_LIST_SAVE_(leanvec_query_matrix)}
+             SVS_LIST_SAVE_(queries_in_training_set)}
         );
     }
 
@@ -131,8 +119,6 @@ struct IVFTest {
             svsbenchmark::extract_filename(table, "graph", root),
             svsbenchmark::extract_filename(table, "queries_f32", root),
             SVS_LOAD_MEMBER_AT_(table, queries_in_training_set),
-            svsbenchmark::extract_filename(table, "leanvec_data_matrix", root),
-            svsbenchmark::extract_filename(table, "leanvec_query_matrix", root),
             num_threads};
     }
 };

bindings/python/CMakeLists.txt

@@ -47,7 +47,6 @@ if (SVS_EXPERIMENTAL_ENABLE_IVF)
     )
 endif()
 
-
 set(LIB_NAME "_svs")
 pybind11_add_module(${LIB_NAME} MODULE ${CPP_FILES})
 target_link_libraries(${LIB_NAME} PRIVATE pybind11::module)

bindings/python/include/svs/python/ivf.h

@@ -51,14 +51,9 @@ template <typename F> void for_standard_specializations(F&& f) {
     // Pattern:
     // QueryType, DataType, Dimensionality, Enable Building
     // clang-format off
-    X(float,  svs::BFloat16, 512, EnableBuild::FromFileAndArray);
-
-    XN(float,  float,         512);
-    XN(float,  svs::Float16,  512);
-
     X(float,  svs::BFloat16, Dynamic, EnableBuild::FromFileAndArray);
-    XN(float,  float,         Dynamic);
-    XN(float,  svs::Float16,  Dynamic);
+    X(float,  float,         Dynamic, EnableBuild::FromFileAndArray);
+    X(float,  svs::Float16,  Dynamic, EnableBuild::FromFileAndArray);
     // clang-format on
 #undef XN
 #undef X

bindings/python/setup.py

@@ -22,6 +22,8 @@
 cmake_args = [
     # Export compile commands to allow us to explore compiler flags as needed.
     "-DCMAKE_EXPORT_COMPILE_COMMANDS=YES",
+    "-DSVS_EXPERIMENTAL_ENABLE_IVF=YES ",
+    "-DSVS_EXPERIMENTAL_BUILD_CUSTOM_MKL=YES ",
 ]
 
 # Determine the root of the repository

bindings/python/src/ivf.cpp

@@ -521,7 +521,7 @@ void wrap(py::module& m) {
             py::arg("num_centroids") = 1000,
             py::arg("minibatch_size") = 10'000,
             py::arg("num_iterations") = 10,
-            py::arg("is_hierarchical") = false,
+            py::arg("is_hierarchical") = true,
             py::arg("training_fraction") = 0.1,
             py::arg("hierarchical_level1_clusters") = 0,
             py::arg("seed") = 0xc0ffee,

bindings/python/tests/common.py

@@ -40,8 +40,12 @@
 test_groundtruth_cosine = str(TEST_DATASET_DIR.joinpath("groundtruth_cosine.ivecs"))
 test_vamana_reference = str(TEST_DATASET_DIR.joinpath("reference/vamana_reference.toml"))
 
+test_ivf_clustering = str(TEST_DATASET_DIR.joinpath("ivf_clustering"))
+test_ivf_reference = str(TEST_DATASET_DIR.joinpath("reference/ivf_reference.toml"))
+
 test_number_of_vectors = 10000
 test_dimensions = 128
+test_number_of_clusters = 128
 
 #####
 ##### Helper Functions