Refactor Vamana Iterator: simplified API and performance enhancement (#105)

This PR simplifies the Vamana Iterator API and enhances its performance
by improving buffer management and search behavior.

*Key Changes:*

**Eliminated `soft_clear()` with Extra Buffer Capacity:**
The buffer size is now larger than the search_window_size, removing the
need for `soft_clear()` when calling `next()`.
This ensures smoother search continuation and better performance.

**Search Triggered Exclusively by `next(batch_size)`:**
The iterator is now empty upon creation or after being updated.
Searches are only performed when `next(batch_size)` is explicitly
called, providing better control over search execution. On calling
`next(batch_size)`, both the `search_window_size` and `buffer_capacity`
are increased by `batch_size` and then search is conducted.

**Dynamic Buffer Adjustment with batch_size:**
Introduced the `batch_size` argument in `next()`, allowing the search
buffer size to be dynamically adjusted based on the batch size.

**Simplified API by Removing Schedules:**
Removed the usage of schedules in the iterator, simplifying its design
and making it easier to use.

Author: ibhati

.github/workflows/build-linux.yml

@@ -81,3 +81,8 @@ jobs:
       working-directory: ${{ runner.temp }}/build/tests
       run: ctest -C ${{ matrix.build_type }}
 
+    - name: Run Cpp Examples
+      env:
+          CTEST_OUTPUT_ON_FAILURE: 1
+      working-directory: ${{ runner.temp }}/build/examples/cpp
+      run: ctest -C RelWithDebugInfo

benchmark/include/svs-benchmark/vamana/iterator.h

@@ -30,70 +30,11 @@
 
 namespace svsbenchmark::vamana {
 
-/// Pre-configuration for the linear schedule.
-struct LinearSchedulePrototype {
-    size_t scale_search_window_;
-    size_t scale_buffer_capacity_;
-    int64_t enable_filter_after_;
-    size_t batch_size_start_;
-    size_t scale_batch_size_;
-    // Whether search should be restarted on every iteration.
-    bool restart_searches_;
-
-    ///// Saving and Loading.
-    static constexpr std::string_view serialization_schema =
-        "svsbench_vamana_iter_schedule";
-    static constexpr svs::lib::Version save_version{0, 0, 0};
-
-    svs::lib::SaveTable save() const {
-        return svs::lib::SaveTable{
-            serialization_schema,
-            save_version,
-            {SVS_LIST_SAVE_(scale_search_window),
-             SVS_LIST_SAVE_(scale_buffer_capacity),
-             SVS_LIST_SAVE_(enable_filter_after),
-             SVS_LIST_SAVE_(batch_size_start),
-             SVS_LIST_SAVE_(scale_batch_size),
-             SVS_LIST_SAVE_(restart_searches)}};
-    }
-
-    static LinearSchedulePrototype load(const svs::lib::ContextFreeLoadTable& table) {
-        return LinearSchedulePrototype{
-            SVS_LOAD_MEMBER_AT_(table, scale_search_window),
-            SVS_LOAD_MEMBER_AT_(table, scale_buffer_capacity),
-            SVS_LOAD_MEMBER_AT_(table, enable_filter_after),
-            SVS_LOAD_MEMBER_AT_(table, batch_size_start),
-            SVS_LOAD_MEMBER_AT_(table, scale_batch_size),
-            SVS_LOAD_MEMBER_AT_(table, restart_searches)};
-    }
-
-    // Return several representative examples for the schedule.
-    static std::vector<LinearSchedulePrototype> examples() {
-        return {{10, 20, -1, 10, 0, false}, {10, 10, 3, 10, 5, false}};
-    }
-
-    // Should search be restarted from scratch every iteration.
-    bool restart_every_iteration() const { return restart_searches_; }
-
-    // Materialize an actual schedule given a set of base parameters.
-    // NOTE: This does not propagate the `restart_searches_` flag.
-    svs::index::vamana::LinearSchedule
-    materialize(const svs::index::vamana::VamanaSearchParameters& sp) const {
-        return svs::index::vamana::LinearSchedule{
-            sp,
-            svs::lib::narrow<uint16_t>(scale_search_window_),
-            svs::lib::narrow<uint16_t>(scale_buffer_capacity_),
-            svs::lib::narrow<int16_t>(enable_filter_after_),
-            svs::lib::narrow<uint16_t>(batch_size_start_),
-            svs::lib::narrow<uint16_t>(scale_batch_size_)};
-    }
-};
-
 struct IteratorSearchParameters {
   public:
     ///// Members
-    // The schedules to try.
-    std::vector<LinearSchedulePrototype> schedules_;
+    // Batch sizes to use for the iterator.
+    std::vector<size_t> batch_sizes_{{10, 20}};
     // target recalls relative to base number of neighbors.
     std::vector<svs::lib::Percent> target_recalls_;
     // The number of batches to yield.
@@ -108,7 +49,7 @@ struct IteratorSearchParameters {
 
     static IteratorSearchParameters example() {
         return IteratorSearchParameters{
-            .schedules_ = LinearSchedulePrototype::examples(),
+            .batch_sizes_ = {10},
             .target_recalls_ = {svs::lib::Percent(0.9)},
             .num_batches_ = 5,
             .query_subsample_ = 10,
@@ -119,15 +60,15 @@ struct IteratorSearchParameters {
         return svs::lib::SaveTable{
             serialization_schema,
             save_version,
-            {SVS_LIST_SAVE_(schedules),
+            {SVS_LIST_SAVE_(batch_sizes),
              SVS_LIST_SAVE_(target_recalls),
              SVS_LIST_SAVE_(num_batches),
              SVS_LIST_SAVE_(query_subsample)}};
     }
 
     static IteratorSearchParameters load(const svs::lib::ContextFreeLoadTable& table) {
         return IteratorSearchParameters{
-            SVS_LOAD_MEMBER_AT_(table, schedules),
+            SVS_LOAD_MEMBER_AT_(table, batch_sizes),
             SVS_LOAD_MEMBER_AT_(table, target_recalls),
             SVS_LOAD_MEMBER_AT_(table, num_batches),
             SVS_LOAD_MEMBER_AT_(table, query_subsample)};
@@ -279,30 +220,25 @@ struct YieldedResult {
 
 // TODO: Make the dependence on `Report` looser.
 template <typename Index> struct QueryIteratorResult {
-    LinearSchedulePrototype schedule_;
+    size_t batch_size_;
     size_t num_batches_;
     double target_recall_;
     search::RunReport<Index> report_;
-    // The search parameters used for each iteration.
-    // Must be the same for all queries in the batch.
-    std::vector<svs::index::vamana::VamanaSearchParameters> iteration_parameters_;
     // Outer vector: Results for each query.
     // Inner vector: Results within a query.
     std::vector<std::vector<YieldedResult>> results_;
 
     ///// Constructor
     QueryIteratorResult(
-        const LinearSchedulePrototype& schedule,
+        size_t batch_size,
         double target_recall,
         search::RunReport<Index> report,
-        std::vector<svs::index::vamana::VamanaSearchParameters> iteration_parameters,
         std::vector<std::vector<YieldedResult>> results
     )
-        : schedule_{schedule}
-        , num_batches_{iteration_parameters.size()}
+        : batch_size_{batch_size}
+        , num_batches_{results.at(0).size()}
         , target_recall_{target_recall}
         , report_{std::move(report)}
-        , iteration_parameters_{std::move(iteration_parameters)}
         , results_{std::move(results)} {
         // Ensure all the yielded results have the correct size.
         for (size_t i = 0, imax = results_.size(); i < imax; ++i) {
@@ -326,11 +262,10 @@ template <typename Index> struct QueryIteratorResult {
         return svs::lib::SaveTable{
             serialization_schema,
             save_version,
-            {SVS_LIST_SAVE_(schedule),
+            {SVS_LIST_SAVE_(batch_size),
              SVS_LIST_SAVE_(num_batches),
              SVS_LIST_SAVE_(target_recall),
              SVS_LIST_SAVE_(report),
-             SVS_LIST_SAVE_(iteration_parameters),
              SVS_LIST_SAVE_(results)}};
     }
 };
@@ -360,15 +295,14 @@ std::vector<QueryIteratorResult<Index>> tune_and_search_iterator(
 
     // Loop over each batchsize.
     auto query_iterator_results = std::vector<QueryIteratorResult<Index>>{};
-    for (const auto& schedule : parameters.schedules_) {
-        auto initial_batch_size = schedule.batch_size_start_;
+    for (const auto& batch_size : parameters.batch_sizes_) {
         for (auto target_recall : parameters.target_recalls_) {
             // Calibrate the index for the given recall.
             auto config = traits::calibrate(
                 index,
                 query_set.training_set_,
                 query_set.training_set_groundtruth_,
-                initial_batch_size,
+                batch_size,
                 target_recall.value(),
                 context,
                 extra
@@ -379,7 +313,7 @@ std::vector<QueryIteratorResult<Index>> tune_and_search_iterator(
                 index,
                 query_set.test_set_,
                 query_set.test_set_groundtruth_,
-                initial_batch_size,
+                batch_size,
                 target_recall.value(),
                 svsbenchmark::CalibrateContext::TestSetTune,
                 config,
@@ -389,7 +323,7 @@ std::vector<QueryIteratorResult<Index>> tune_and_search_iterator(
             // Now we have a calibrated configuration - obtain a baseline report for
             // searching with this batchsize.
             auto report = svsbenchmark::search::search_with_config(
-                index, config, query_test, groundtruth_test, initial_batch_size
+                index, config, query_test, groundtruth_test, batch_size
             );
 
             // `resuilt_buffer`: All results that have been returned by the iterator.
@@ -452,38 +386,34 @@ std::vector<QueryIteratorResult<Index>> tune_and_search_iterator(
 
                 // The first call to `iterator` kick-starts graph search.
                 auto tic = svs::lib::now();
-                auto iterator = make_iterator(index, query, config, schedule);
+                auto iterator = make_iterator(index, query);
+                iterator.next(config.buffer_config_.get_search_window_size());
                 auto elapsed = svs::lib::time_difference(tic);
                 if (i == 0) {
-                    iteration_parameters.push_back(iterator.parameters_for_current_batch());
+                    iteration_parameters.push_back(
+                        iterator.parameters_for_current_iteration()
+                    );
                 }
 
                 timings_for_this_query.push_back(tally(iterator, i, 0, elapsed));
                 for (size_t j = 0; j < parameters.num_batches_; ++j) {
-                    // If requested by the parent schedule, reset search for this
-                    // iteration.
-                    if (schedule.restart_every_iteration()) {
-                        iterator.restart_next_search();
-                    }
-
                     tic = svs::lib::now();
-                    iterator.next();
+                    iterator.next(batch_size);
                     elapsed = svs::lib::time_difference(tic);
                     timings_for_this_query.push_back(tally(iterator, i, j + 1, elapsed));
                     if (i == 0) {
                         iteration_parameters.push_back(
-                            iterator.parameters_for_current_batch()
+                            iterator.parameters_for_current_iteration()
                         );
                     }
                 }
             }
 
             // Finish up summarizing these results.
             query_iterator_results.emplace_back(
-                schedule,
+                batch_size,
                 target_recall.value(),
                 std::move(report),
-                std::move(iteration_parameters),
                 std::move(yielded_results)
             );
             do_checkpoint(query_iterator_results);
@@ -522,9 +452,7 @@ toml::table tune_and_search_iterator(
         job.parameters_,
         query_set,
         svsbenchmark::CalibrateContext::InitialTrainingSet,
-        [](const auto& index, const auto& query, const auto& config, const auto& schedule) {
-            return index.batch_iterator(query, schedule.materialize(config));
-        },
+        [](const auto& index, const auto& query) { return index.batch_iterator(query); },
         do_checkpoint,
         svsbenchmark::IndexTraits<Index>::regression_optimization()
     );

examples/cpp/CMakeLists.txt

@@ -37,7 +37,6 @@ endfunction()
 create_simple_example(saveload test_saveload saveload.cpp)
 create_simple_example(types test_types types.cpp)
 create_simple_example(vamana_iterator test_vamana_iterator vamana_iterator.cpp)
-create_simple_example(custom_thread_pool test_custom_thread_pool custom_thread_pool.cpp)
 
 ## More complicated examples involving more extensive setup.
 
@@ -48,6 +47,7 @@ create_simple_example(custom_thread_pool test_custom_thread_pool custom_thread_p
 configure_file(../../data/test_dataset/data_f32.fvecs . COPYONLY)
 configure_file(../../data/test_dataset/queries_f32.fvecs . COPYONLY)
 configure_file(../../data/test_dataset/groundtruth_euclidean.ivecs . COPYONLY)
+
 # The vamana test executable.
 add_executable(vamana vamana.cpp)
 target_include_directories(vamana PRIVATE ${CMAKE_CURRENT_LIST_DIR})
@@ -61,6 +61,20 @@ add_test(
         groundtruth_euclidean.ivecs
 )
 
+# The custom thread pool executable.
+add_executable(custom_thread_pool custom_thread_pool.cpp)
+target_include_directories(custom_thread_pool PRIVATE ${CMAKE_CURRENT_LIST_DIR})
+target_link_libraries(custom_thread_pool ${SVS_LIB} svs_compile_options svs_native_options)
+add_test(
+    NAME test_custom_thread_pool
+    COMMAND
+        custom_thread_pool
+        data_f32.fvecs
+        queries_f32.fvecs
+        groundtruth_euclidean.ivecs
+)
+
+
 #####
 ##### Dispatcher
 #####

examples/cpp/vamana_iterator.cpp

@@ -40,7 +40,7 @@ void check(bool value, std::string_view expr, svs::lib::detail::LineInfo linfo)
 #define CHECK(expr) check(expr, #expr, SVS_LINEINFO);
 
 //! [Example Index Construction]
-[[nodiscard]] svs::Vamana make_example_index() {
+[[nodiscard]] auto make_example_index() {
     // Build the index.
     auto build_parameters = svs::index::vamana::VamanaBuildParameters{
         1.2, // alpha
@@ -68,42 +68,34 @@ void check(bool value, std::string_view expr, svs::lib::detail::LineInfo linfo)
     }
 
     // Build the index.
-    return svs::Vamana::build<float>(build_parameters, std::move(data), svs::DistanceL2{});
+    return svs::index::vamana::auto_build(
+        build_parameters, std::move(data), svs::DistanceL2{}, 1
+    );
 }
 //! [Example Index Construction]
 
-void demonstrate_default_schedule() {
+void demonstrate_iterator() {
     //! [Setup]
     auto index = make_example_index();
 
-    // Base search parameters for the iterator schedule.
-    // This uses a search window size/capacity of 4.
-    auto base_parameters = svs::index::vamana::VamanaSearchParameters{}.buffer_config({4});
-
-    // The default schedule take base parameters and a batch size.
     // Each iteration will yield 3 elements that have not been yielded previously.
     size_t batchsize = 3;
-    auto schedule = svs::index::vamana::DefaultSchedule{base_parameters, batchsize};
 
     // Create a batch iterator over the index for the query.
-    // After the constructor returns, the contents of the first batch will be available.
     auto itr = [&]() {
         // Construct a query a query in a scoped block to demonstrate that the iterator
         // maintains an internal copy.
         auto query = std::vector<float>{3.25, 3.25, 3.25, 3.25};
 
-        // Make a batch iterator for the query using the provided schedule.
-        return index.batch_iterator(svs::lib::as_const_span(query), schedule);
+        // Make a batch iterator for the query.
+        return svs::VamanaIterator(index, svs::lib::as_const_span(query));
     }();
     //! [Setup]
 
-    //! [Initial Checks]
-    // The iterator was configured to yield three neighbors on each invocation.
-    // This information is available through the `size()` method.
+    //! [First Iteration]
+    itr.next(batchsize);
     CHECK(itr.size() == 3);
-
-    // The contents of the iterator are for batch 0.
-    CHECK(itr.batch() == 0);
+    CHECK(itr.batch_number() == 1);
 
     // Obtain a view of the current list candidates.
     std::span<const svs::Neighbor<size_t>> results = itr.results();
@@ -117,18 +109,15 @@ void demonstrate_default_schedule() {
     CHECK(results[0].id() == 3);
     CHECK(results[1].id() == 4);
     CHECK(results[2].id() == 2);
-    //! [Initial Checks]
+    //! [First Iteration]
 
     //! [Next Iteration]
-    // Once we've finished with the current batch of neighbors, we can step the iterator
-    // to the next batch.
-    //
-    // Using the `DefaultSchedule`, we will retrieve at most 3 new candidates.
-    itr.next();
+    // This will yield the next batch of neighbors.
+    itr.next(batchsize);
     CHECK(itr.size() == 3);
 
-    // The contents of the iterator are for batch 1.
-    CHECK(itr.batch() == 1);
+    // The contents of the iterator are for batch 2.
+    CHECK(itr.batch_number() == 2);
 
     // Update and inspect the results.
     results = itr.results();
@@ -144,7 +133,7 @@ void demonstrate_default_schedule() {
     //! [Final Iteration]
     // So far, the iterator has yielded 6 of the 7 vectors in the dataset.
     // This call to `next()` should only yield a single neighbor - the last on in the index.
-    itr.next();
+    itr.next(batchsize);
     CHECK(itr.size() == 1);
     CHECK(itr.done());
     results = itr.results();
@@ -154,15 +143,15 @@ void demonstrate_default_schedule() {
 
     //! [Beyond Final Iteration]
     // Calling `next()` again should yield no more candidates.
-    itr.next();
+    itr.next(batchsize);
     CHECK(itr.size() == 0);
     CHECK(itr.done());
     //! [Beyond Final Iteration]
 }
 
 // Alternative main definition
 int svs_main(std::vector<std::string> SVS_UNUSED(args)) {
-    demonstrate_default_schedule();
+    demonstrate_iterator();
     return 0;
 }