[C API] LeanVec training data handling refactoring

bindings/c/include/svs/c/svs_c.h

@@ -645,23 +645,33 @@ SVS_API bool svs_storage_get_kind(
 SVS_API void svs_storage_free(svs_storage_h storage);
 
 /// @brief Train LeanVec dimensionality-reduction matrices from a data sample
-/// @param dim The dimensionality of the data (and training queries)
+/// @param builder The index builder handle
+/// @param leanvec_dims The reduced number of LeanVec dimensions
 /// @param num_vectors The number of data vectors in x
 /// @param x Pointer to the data vectors [num_vectors x dim] (float array)
 /// @param num_queries The number of training queries in x_q (0 for in-distribution)
 /// @param x_q Pointer to the training queries [num_queries x dim], or NULL. When
 /// provided, matrices are trained out-of-distribution (OOD) using these queries;
 /// when num_queries is 0 or x_q is NULL, in-distribution (PCA) matrices are computed.
-/// @param leanvec_dims The reduced number of LeanVec dimensions
 /// @param out_err An optional error handle to capture errors
 /// @return A handle to the trained LeanVec matrices
+/// @remarks
+/// * The training data build process depends on the index builder's dimension and thread
+/// pool configuration, so the builder must be configured with the correct dimension and
+/// thread pool before calling this function.
+/// * The training data is copied into the returned handle, so the caller may
+/// free or modify @p x and @p x_q once this call returns.
+/// * The training data built by this function is used to create a LeanVec storage
+/// configuration via svs_storage_create_leanvec_trained(). The training data handle may be
+/// freed once svs_storage_create_leanvec_trained() returns, as the storage retains a
+/// reference to the trained matrices.
 SVS_API svs_leanvec_training_data_h svs_leanvec_training_data_build(
-    size_t dim,
+    svs_index_builder_h builder,
+    size_t leanvec_dims,
     size_t num_vectors,
     const float* x,
     size_t num_queries,
-    const float* x_q /*=NULL*/,
-    size_t leanvec_dims,
+    const float* x_q,   /*=NULL*/
     svs_error_h out_err /*=NULL*/
 );
 

bindings/c/src/data_builder/leanvec.hpp

@@ -19,6 +19,7 @@
 
 #include "svs/c/svs_c.h"
 
+#include "leanvec_training_data.hpp"
 #include "storage.hpp"
 #include "types_support.hpp"
 
@@ -53,11 +54,12 @@ class LeanVecDataBuilder {
     std::optional<svs::leanvec::LeanVecMatrices<svs::Dynamic>> matrices_;
 
   public:
-    LeanVecDataBuilder(
-        size_t leanvec_dims,
-        std::optional<svs::leanvec::LeanVecMatrices<svs::Dynamic>> matrices = std::nullopt
-    )
+    LeanVecDataBuilder(size_t leanvec_dims)
         : leanvec_dims_(leanvec_dims)
+        , matrices_(std::nullopt) {}
+
+    LeanVecDataBuilder(svs::leanvec::LeanVecMatrices<svs::Dynamic> matrices)
+        : leanvec_dims_(matrices.num_cols())
         , matrices_(std::move(matrices)) {}
 
     using data_type = svs::leanvec::LeanDataset<
@@ -106,8 +108,9 @@ struct lib::
         // `leanvec_dims` is taken from the training data at storage construction,
         // so it is authoritative in both cases.
         if (leanvec->training_data) {
-            return To{leanvec->leanvec_dims, leanvec->training_data->matrices()};
+            return To{leanvec->training_data->matrices()};
         }
+        assert(leanvec->leanvec_dims > 0);
         return To{leanvec->leanvec_dims};
     }
 };

bindings/c/src/leanvec_training_data.hpp

@@ -49,16 +49,14 @@ class LeanVecTrainingData {
         size_t leanvec_dims,
         svs::threads::ThreadPoolHandle& pool
     )
-        : leanvec_dims_{leanvec_dims}
-        , matrices_{
+        : matrices_{
               queries.size() == 0 ? compute_pca(data, leanvec_dims, pool)
                                   : compute_ood(data, queries, leanvec_dims, pool)} {}
 
-    size_t leanvec_dims() const { return leanvec_dims_; }
+    size_t leanvec_dims() const { return matrices_.num_cols(); }
     const matrices_type& matrices() const { return matrices_; }
 
   private:
-    size_t leanvec_dims_;
     matrices_type matrices_;
 
     static matrices_type compute_pca(
@@ -91,4 +89,16 @@ class LeanVecTrainingData {
 
 } // namespace svs::c_runtime
 
+#else // SVS_RUNTIME_ENABLE_LVQ_LEANVEC
+namespace svs::c_runtime {
+class LeanVecTrainingData {
+  public:
+    template <typename... Args> LeanVecTrainingData(Args&&...) {
+        throw svs::c_runtime::not_implemented(
+            "LeanVec training is not implemented in this build"
+        );
+    }
+};
+} // namespace svs::c_runtime
+
 #endif // SVS_RUNTIME_ENABLE_LVQ_LEANVEC

bindings/c/src/storage.hpp

@@ -25,8 +25,6 @@
 #include <svs/lib/type_traits.h>
 
 #ifdef SVS_RUNTIME_ENABLE_LVQ_LEANVEC
-#include "leanvec_training_data.hpp"
-
 #include <svs/cpuid.h>
 #endif
 
@@ -58,17 +56,16 @@ struct StorageSimple : public Storage {
     }
 };
 
+class LeanVecTrainingData;
 struct StorageLeanVec : public Storage {
     size_t leanvec_dims;
     size_t primary_bits;
     size_t secondary_bits;
-#ifdef SVS_RUNTIME_ENABLE_LVQ_LEANVEC
     // Pre-trained reduction matrices; when set, they are used instead of PCA
     // matrices computed at build time (enables out-of-distribution LeanVec).
     // Fixed at construction: `leanvec_dims` is taken from the training data, so
     // the two can never disagree.
-    std::shared_ptr<const LeanVecTrainingData> training_data;
-#endif
+    std::shared_ptr<const LeanVecTrainingData> training_data = nullptr;
 
     StorageLeanVec(size_t leanvec_dims, svs_data_type_t primary, svs_data_type_t secondary)
         : Storage{SVS_STORAGE_KIND_LEANVEC}
@@ -88,18 +85,16 @@ struct StorageLeanVec : public Storage {
 #endif
     }
 
-#ifdef SVS_RUNTIME_ENABLE_LVQ_LEANVEC
     // Construct from pre-trained matrices. `leanvec_dims` is the single value
     // carried by the training data, so no reconciliation is needed.
     StorageLeanVec(
         std::shared_ptr<const LeanVecTrainingData> training_data,
         svs_data_type_t primary,
         svs_data_type_t secondary
     )
-        : StorageLeanVec(training_data->leanvec_dims(), primary, secondary) {
+        : StorageLeanVec(0, primary, secondary) {
         this->training_data = std::move(training_data);
     }
-#endif
 
     static size_t to_bits_number(svs_data_type_t data_type) {
         switch (data_type) {

bindings/c/src/svs_c.cpp

@@ -20,14 +20,11 @@
 #include "error.hpp"
 #include "index.hpp"
 #include "index_builder.hpp"
+#include "leanvec_training_data.hpp"
 #include "storage.hpp"
 #include "threadpool.hpp"
 #include "types_support.hpp"
 
-#ifdef SVS_RUNTIME_ENABLE_LVQ_LEANVEC
-#include "leanvec_training_data.hpp"
-#endif
-
 #include <cstddef>
 #include <filesystem>
 #include <memory>
@@ -61,9 +58,7 @@ struct svs_storage {
 };
 
 struct svs_leanvec_training_data {
-#ifdef SVS_RUNTIME_ENABLE_LVQ_LEANVEC
     std::shared_ptr<const svs::c_runtime::LeanVecTrainingData> impl;
-#endif
 };
 
 extern "C" uint32_t svs_get_version() { return SVS_C_API_VERSION; }
@@ -342,7 +337,6 @@ extern "C" svs_storage_h svs_storage_create_leanvec_trained(
     using namespace svs::c_runtime;
     return wrap_exceptions(
         [&]() -> svs_storage_h {
-#ifdef SVS_RUNTIME_ENABLE_LVQ_LEANVEC
             EXPECT_ARG_NOT_NULL(training_data);
             INVALID_ARGUMENT_IF(
                 (training_data->impl == nullptr), "training_data holds no trained matrices"
@@ -356,14 +350,6 @@ extern "C" svs_storage_h svs_storage_create_leanvec_trained(
             auto result = new svs_storage;
             result->impl = storage;
             return result;
-#else
-            (void)training_data;
-            (void)primary;
-            (void)secondary;
-            throw svs::c_runtime::not_implemented(
-                "LeanVec storage is not implemented in this build"
-            );
-#endif
         },
         out_err
     );
@@ -431,22 +417,23 @@ extern "C" SVS_API bool svs_storage_get_kind(
 extern "C" void svs_storage_free(svs_storage_h storage) { delete storage; }
 
 extern "C" svs_leanvec_training_data_h svs_leanvec_training_data_build(
-    size_t dim,
+    svs_index_builder_h builder,
+    size_t leanvec_dims,
     size_t num_vectors,
     const float* x,
     size_t num_queries,
     const float* x_q,
-    size_t leanvec_dims,
     svs_error_h out_err
 ) {
     using namespace svs::c_runtime;
     return wrap_exceptions(
         [&]() -> svs_leanvec_training_data_h {
-#ifdef SVS_RUNTIME_ENABLE_LVQ_LEANVEC
-            EXPECT_ARG_GT_THAN(dim, 0);
+            EXPECT_ARG_NOT_NULL(builder);
             EXPECT_ARG_GT_THAN(num_vectors, 0);
             EXPECT_ARG_NOT_NULL(x);
             EXPECT_ARG_GT_THAN(leanvec_dims, 0);
+
+            const auto dim = builder->impl->dimension;
             EXPECT_ARG_GE_THAN(dim, leanvec_dims);
             INVALID_ARGUMENT_IF(
                 (num_queries > 0 && x_q == nullptr),
@@ -459,25 +446,14 @@ extern "C" svs_leanvec_training_data_h svs_leanvec_training_data_build(
                 x_q, (x_q == nullptr) ? 0 : num_queries, dim
             );
 
-            auto pool = ThreadPoolBuilder{}.build();
+            auto pool = builder->impl->pool_builder.build();
             auto training_data = std::make_shared<const LeanVecTrainingData>(
                 data, queries, leanvec_dims, pool
             );
 
             auto result = new svs_leanvec_training_data;
             result->impl = std::move(training_data);
             return result;
-#else
-            (void)dim;
-            (void)num_vectors;
-            (void)x;
-            (void)num_queries;
-            (void)x_q;
-            (void)leanvec_dims;
-            throw svs::c_runtime::not_implemented(
-                "LeanVec training data is not implemented in this build"
-            );
-#endif
         },
         out_err
     );

bindings/c/tests/c_api_index.cpp

@@ -257,24 +257,23 @@ CATCH_TEST_CASE("C API Index Build and Search", "[c_api][index][build][search]")
 
     CATCH_SECTION("Index Build and Search with pre-trained LeanVec") {
         svs_error_h error = svs_error_create();
+        svs_algorithm_h algorithm = svs_algorithm_create_vamana(16, 32, 50, error);
+        CATCH_REQUIRE(algorithm != nullptr);
+
+        svs_index_builder_h builder = svs_index_builder_create(
+            SVS_DISTANCE_METRIC_EUCLIDEAN, DIMENSION, algorithm, error
+        );
+        CATCH_REQUIRE(builder != nullptr);
+
+        bool success = svs_index_builder_set_threadpool(
+            builder, SVS_THREADPOOL_KIND_NATIVE, NUM_THREADS, error
+        );
+        CATCH_REQUIRE(success);
 
         auto run_build_and_search = [&](svs_storage_h storage) {
             CATCH_REQUIRE(storage != nullptr);
             CATCH_REQUIRE(svs_error_ok(error));
 
-            svs_algorithm_h algorithm = svs_algorithm_create_vamana(16, 32, 50, error);
-            CATCH_REQUIRE(algorithm != nullptr);
-
-            svs_index_builder_h builder = svs_index_builder_create(
-                SVS_DISTANCE_METRIC_EUCLIDEAN, DIMENSION, algorithm, error
-            );
-            CATCH_REQUIRE(builder != nullptr);
-
-            bool success = svs_index_builder_set_threadpool(
-                builder, SVS_THREADPOOL_KIND_NATIVE, NUM_THREADS, error
-            );
-            CATCH_REQUIRE(success);
-
             success = svs_index_builder_set_storage(builder, storage, error);
             CATCH_REQUIRE(success);
             CATCH_REQUIRE(svs_error_ok(error));
@@ -296,16 +295,14 @@ CATCH_TEST_CASE("C API Index Build and Search", "[c_api][index][build][search]")
 
             svs_search_results_free(&results);
             svs_index_free(index);
-            svs_index_builder_free(builder);
-            svs_algorithm_free(algorithm);
             svs_storage_free(storage);
         };
 
         // Builds a storage from matrices trained up front, then drops the training
         // data handle to confirm the storage keeps its own reference to the matrices.
         auto build_trained_storage = [&](size_t num_queries, const float* x_q) {
             svs_leanvec_training_data_h training_data = svs_leanvec_training_data_build(
-                DIMENSION, NUM_VECTORS, data.data(), num_queries, x_q, DIMENSION / 2, error
+                builder, DIMENSION / 2, NUM_VECTORS, data.data(), num_queries, x_q, error
             );
             CATCH_REQUIRE(check_training_data_support(training_data, error) == true);
             if (!training_data_usable(training_data)) {
@@ -332,6 +329,8 @@ CATCH_TEST_CASE("C API Index Build and Search", "[c_api][index][build][search]")
             run_build_and_search(storage);
         }
 
+        svs_index_builder_free(builder);
+        svs_algorithm_free(algorithm);
         svs_error_free(error);
     }