diff --git a/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h b/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h index 2e53b8050..640f7eeba 100644 --- a/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h +++ b/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h @@ -39,6 +39,9 @@ struct SVS_RUNTIME_API DynamicVamanaIndex : public VamanaIndex { // Utility function to check storage kind support static Status check_storage_kind(StorageKind storage_kind) noexcept; + static Status check_params(const VamanaIndex::DynamicIndexParams& dynamic_index_params + ) noexcept; + // Static constructors and destructors static Status build( DynamicVamanaIndex** index, @@ -46,7 +49,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndex : public VamanaIndex { MetricType metric, StorageKind storage_kind, const VamanaIndex::BuildParams& params = {}, - const VamanaIndex::SearchParams& default_search_params = {} + const VamanaIndex::SearchParams& default_search_params = {}, + const VamanaIndex::DynamicIndexParams& dynamic_index_params = {} ) noexcept; static Status destroy(DynamicVamanaIndex* index) noexcept; @@ -58,6 +62,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndex : public VamanaIndex { MetricType metric, StorageKind storage_kind ) noexcept; + + virtual size_t blocksize_bytes() const noexcept = 0; }; struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex { @@ -69,7 +75,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex { StorageKind storage_kind, size_t leanvec_dims, const VamanaIndex::BuildParams& params = {}, - const VamanaIndex::SearchParams& default_search_params = {} + const VamanaIndex::SearchParams& default_search_params = {}, + const VamanaIndex::DynamicIndexParams& dynamic_index_params = {} ) noexcept; // Specialization to build LeanVec-based Vamana index with provided training data @@ -80,7 +87,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex { StorageKind storage_kind, const LeanVecTrainingData* training_data, const VamanaIndex::BuildParams& params = {}, - const VamanaIndex::SearchParams& default_search_params = {} + const VamanaIndex::SearchParams& default_search_params = {}, + const VamanaIndex::DynamicIndexParams& dynamic_index_params = {} ) noexcept; }; diff --git a/bindings/cpp/include/svs/runtime/vamana_index.h b/bindings/cpp/include/svs/runtime/vamana_index.h index b29f44ddc..ba9739fb4 100644 --- a/bindings/cpp/include/svs/runtime/vamana_index.h +++ b/bindings/cpp/include/svs/runtime/vamana_index.h @@ -44,6 +44,10 @@ struct SVS_RUNTIME_API VamanaIndex { size_t prefetch_step = Unspecify(); }; + struct DynamicIndexParams { + size_t blocksize_exp = 30; + }; + virtual Status search( size_t n, const float* x, diff --git a/bindings/cpp/src/dynamic_vamana_index.cpp b/bindings/cpp/src/dynamic_vamana_index.cpp index 1a7723d30..a9cbe027d 100644 --- a/bindings/cpp/src/dynamic_vamana_index.cpp +++ b/bindings/cpp/src/dynamic_vamana_index.cpp @@ -63,6 +63,8 @@ struct DynamicVamanaIndexManagerBase : public DynamicVamanaIndex { }); } + size_t blocksize_bytes() const noexcept { return impl_->blocksize_bytes(); } + Status remove_selected(size_t* num_removed, const IDFilter& selector) noexcept override { return runtime_error_wrapper([&] { @@ -136,19 +138,40 @@ Status DynamicVamanaIndex::check_storage_kind(StorageKind storage_kind) noexcept ); } +Status DynamicVamanaIndex::check_params( + const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params +) noexcept { + constexpr static size_t kMaxBlockSizeExp = 30; // 1GB + constexpr static size_t kMinBlockSizeExp = 12; // 4KB + + if (dynamic_index_params.blocksize_exp > kMaxBlockSizeExp) + return Status(ErrorCode::INVALID_ARGUMENT, "Blocksize is too large"); + + if (dynamic_index_params.blocksize_exp < kMinBlockSizeExp) + return Status(ErrorCode::INVALID_ARGUMENT, "Blocksize is too small"); + + return Status_Ok; +} + Status DynamicVamanaIndex::build( DynamicVamanaIndex** index, size_t dim, MetricType metric, StorageKind storage_kind, const DynamicVamanaIndex::BuildParams& params, - const DynamicVamanaIndex::SearchParams& default_search_params + const DynamicVamanaIndex::SearchParams& default_search_params, + const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params ) noexcept { using Impl = DynamicVamanaIndexImpl; *index = nullptr; + + auto status = DynamicVamanaIndex::check_params(dynamic_index_params); + if (!status.ok()) + return status; + return runtime_error_wrapper([&] { auto impl = std::make_unique( - dim, metric, storage_kind, params, default_search_params + dim, metric, storage_kind, params, default_search_params, dynamic_index_params ); *index = new DynamicVamanaIndexManagerBase{std::move(impl)}; }); @@ -181,13 +204,20 @@ Status DynamicVamanaIndexLeanVec::build( StorageKind storage_kind, size_t leanvec_dims, const DynamicVamanaIndex::BuildParams& params, - const DynamicVamanaIndex::SearchParams& default_search_params + const DynamicVamanaIndex::SearchParams& default_search_params, + const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params ) noexcept { using Impl = DynamicVamanaIndexLeanVecImpl; *index = nullptr; return runtime_error_wrapper([&] { auto impl = std::make_unique( - dim, metric, storage_kind, leanvec_dims, params, default_search_params + dim, + metric, + storage_kind, + leanvec_dims, + params, + default_search_params, + dynamic_index_params ); *index = new DynamicVamanaIndexManagerBase{std::move(impl)}; }); @@ -201,7 +231,8 @@ Status DynamicVamanaIndexLeanVec::build( StorageKind storage_kind, const LeanVecTrainingData* training_data, const DynamicVamanaIndex::BuildParams& params, - const DynamicVamanaIndex::SearchParams& default_search_params + const DynamicVamanaIndex::SearchParams& default_search_params, + const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params ) noexcept { using Impl = DynamicVamanaIndexLeanVecImpl; *index = nullptr; @@ -209,7 +240,13 @@ Status DynamicVamanaIndexLeanVec::build( auto training_data_impl = static_cast(training_data)->impl_; auto impl = std::make_unique( - dim, metric, storage_kind, training_data_impl, params, default_search_params + dim, + metric, + storage_kind, + training_data_impl, + params, + default_search_params, + dynamic_index_params ); *index = new DynamicVamanaIndexManagerBase{std::move(impl)}; }); @@ -218,7 +255,7 @@ Status DynamicVamanaIndexLeanVec::build( #else // SVS_LEANVEC_HEADER // LeanVec storage kind is not supported in this build configuration Status DynamicVamanaIndexLeanVec:: - build(DynamicVamanaIndex**, size_t, MetricType, StorageKind, size_t, const DynamicVamanaIndex::BuildParams&, const DynamicVamanaIndex::SearchParams&) noexcept { + build(DynamicVamanaIndex**, size_t, MetricType, StorageKind, size_t, const DynamicVamanaIndex::BuildParams&, const DynamicVamanaIndex::SearchParams&, const DynamicVamanaIndex::DynamicIndexParams&) noexcept { return Status( ErrorCode::NOT_IMPLEMENTED, "DynamicVamanaIndexLeanVec is not supported in this build configuration." @@ -226,7 +263,7 @@ Status DynamicVamanaIndexLeanVec:: } Status DynamicVamanaIndexLeanVec:: - build(DynamicVamanaIndex**, size_t, MetricType, StorageKind, const LeanVecTrainingData*, const DynamicVamanaIndex::BuildParams&, const DynamicVamanaIndex::SearchParams&) noexcept { + build(DynamicVamanaIndex**, size_t, MetricType, StorageKind, const LeanVecTrainingData*, const DynamicVamanaIndex::BuildParams&, const DynamicVamanaIndex::SearchParams&, const DynamicVamanaIndex::DynamicIndexParams&) noexcept { return Status( ErrorCode::NOT_IMPLEMENTED, "DynamicVamanaIndexLeanVec is not supported in this build configuration." diff --git a/bindings/cpp/src/dynamic_vamana_index_impl.h b/bindings/cpp/src/dynamic_vamana_index_impl.h index 78d44f111..9e4656288 100644 --- a/bindings/cpp/src/dynamic_vamana_index_impl.h +++ b/bindings/cpp/src/dynamic_vamana_index_impl.h @@ -46,13 +46,15 @@ class DynamicVamanaIndexImpl { MetricType metric, StorageKind storage_kind, const VamanaIndex::BuildParams& params, - const VamanaIndex::SearchParams& default_search_params + const VamanaIndex::SearchParams& default_search_params, + const VamanaIndex::DynamicIndexParams& dynamic_index_params ) : dim_{dim} , metric_type_{metric} , storage_kind_{storage_kind} , build_params_{params} - , default_search_params_{default_search_params} { + , default_search_params_{default_search_params} + , dynamic_index_params_{dynamic_index_params} { if (!storage::is_supported_storage_kind(storage_kind)) { throw StatusException{ ErrorCode::INVALID_ARGUMENT, @@ -63,6 +65,8 @@ class DynamicVamanaIndexImpl { size_t size() const { return impl_ ? impl_->size() : 0; } + size_t blocksize_bytes() const { return 1u << dynamic_index_params_.blocksize_exp; } + size_t dimensions() const { return dim_; } MetricType metric_type() const { return metric_type_; } @@ -71,7 +75,8 @@ class DynamicVamanaIndexImpl { void add(data::ConstSimpleDataView data, std::span labels) { if (!impl_) { - return init_impl(data, labels); + auto blocksize_bytes = lib::PowerOfTwo(dynamic_index_params_.blocksize_exp); + return init_impl(data, labels, blocksize_bytes); } impl_->add_points(data, labels); @@ -385,6 +390,7 @@ class DynamicVamanaIndexImpl { const index::vamana::VamanaBuildParameters& parameters, const svs::data::ConstSimpleDataView& data, std::span labels, + svs::lib::PowerOfTwo blocksize_bytes, StorageArgs&&... storage_args ) { auto threadpool = default_threadpool(); @@ -393,6 +399,7 @@ class DynamicVamanaIndexImpl { std::forward(tag), data, threadpool, + blocksize_bytes, std::forward(storage_args)... ); @@ -408,14 +415,18 @@ class DynamicVamanaIndexImpl { }); } - virtual void - init_impl(data::ConstSimpleDataView data, std::span labels) { + virtual void init_impl( + data::ConstSimpleDataView data, + std::span labels, + lib::PowerOfTwo blocksize_bytes + ) { impl_.reset(storage::dispatch_storage_kind( get_storage_kind(), [this]( auto&& tag, data::ConstSimpleDataView data, - std::span labels + std::span labels, + lib::PowerOfTwo blocksize_bytes ) { using Tag = std::decay_t; return build_impl( @@ -423,11 +434,13 @@ class DynamicVamanaIndexImpl { this->metric_type_, this->vamana_build_parameters(), data, - labels + labels, + blocksize_bytes ); }, data, - labels + labels, + blocksize_bytes )); } @@ -522,6 +535,7 @@ class DynamicVamanaIndexImpl { StorageKind storage_kind_; VamanaIndex::BuildParams build_params_; VamanaIndex::SearchParams default_search_params_; + VamanaIndex::DynamicIndexParams dynamic_index_params_; std::unique_ptr impl_; size_t ntotal_soft_deleted{0}; }; diff --git a/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h b/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h index 9eaab2ec2..0ba764d5c 100644 --- a/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h +++ b/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h @@ -54,9 +54,10 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { StorageKind storage_kind, const LeanVecTrainingDataImpl& training_data, const VamanaIndex::BuildParams& params, - const VamanaIndex::SearchParams& default_search_params + const VamanaIndex::SearchParams& default_search_params, + const VamanaIndex::DynamicIndexParams& dynamic_index_params ) - : DynamicVamanaIndexImpl{dim, metric, storage_kind, params, default_search_params} + : DynamicVamanaIndexImpl{dim, metric, storage_kind, params, default_search_params, dynamic_index_params} , leanvec_dims_{training_data.get_leanvec_dims()} , leanvec_matrices_{training_data.get_leanvec_matrices()} { check_storage_kind(storage_kind); @@ -68,9 +69,10 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { StorageKind storage_kind, size_t leanvec_dims, const VamanaIndex::BuildParams& params, - const VamanaIndex::SearchParams& default_search_params + const VamanaIndex::SearchParams& default_search_params, + const VamanaIndex::DynamicIndexParams& dynamic_index_params ) - : DynamicVamanaIndexImpl{dim, metric, storage_kind, params, default_search_params} + : DynamicVamanaIndexImpl{dim, metric, storage_kind, params, default_search_params, dynamic_index_params} , leanvec_dims_{leanvec_dims} , leanvec_matrices_{std::nullopt} { check_storage_kind(storage_kind); @@ -91,15 +93,19 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { } } - void init_impl(data::ConstSimpleDataView data, std::span labels) - override { + void init_impl( + data::ConstSimpleDataView data, + std::span labels, + lib::PowerOfTwo blocksize_bytes + ) override { assert(storage::is_leanvec_storage(this->storage_kind_)); impl_.reset(dispatch_leanvec_storage_kind( this->storage_kind_, [this]( auto&& tag, data::ConstSimpleDataView data, - std::span labels + std::span labels, + lib::PowerOfTwo blocksize_bytes ) { using Tag = std::decay_t; return DynamicVamanaIndexImpl::build_impl( @@ -108,12 +114,14 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { this->vamana_build_parameters(), data, labels, + blocksize_bytes, this->leanvec_dims_, this->leanvec_matrices_ ); }, data, - labels + labels, + blocksize_bytes )); } diff --git a/bindings/cpp/src/svs_runtime_utils.h b/bindings/cpp/src/svs_runtime_utils.h index 8ab0c2884..0688ff257 100644 --- a/bindings/cpp/src/svs_runtime_utils.h +++ b/bindings/cpp/src/svs_runtime_utils.h @@ -177,7 +177,8 @@ template <> struct StorageFactory { template static StorageType init( const svs::data::ConstSimpleDataView& SVS_UNUSED(data), - Pool& SVS_UNUSED(pool) + Pool& SVS_UNUSED(pool), + svs::lib::PowerOfTwo SVS_UNUSED(blocksize_bytes) ) { throw StatusException( ErrorCode::NOT_IMPLEMENTED, "Requested storage kind is not supported" @@ -196,6 +197,27 @@ template <> struct StorageFactory { template struct StorageFactory> { using StorageType = SimpleDatasetType; + template + static StorageType init( + const svs::data::ConstSimpleDataView& data, + Pool& pool, + svs::lib::PowerOfTwo blocksize_bytes + ) { + auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes}; + typename StorageType::allocator_type alloc(parameters); + StorageType result(data.size(), data.dimensions(), alloc); + svs::threads::parallel_for( + pool, + svs::threads::StaticPartition(result.size()), + [&](auto is, auto SVS_UNUSED(tid)) { + for (auto i : is) { + result.set_datum(i, data.get_datum(i)); + } + } + ); + return result; + } + template static StorageType init(const svs::data::ConstSimpleDataView& data, Pool& pool) { StorageType result(data.size(), data.dimensions()); @@ -223,8 +245,14 @@ struct StorageFactory { using StorageType = SQStorageType; template - static StorageType init(const svs::data::ConstSimpleDataView& data, Pool& pool) { - return SQStorageType::compress(data, pool); + static StorageType init( + const svs::data::ConstSimpleDataView& data, + Pool& pool, + svs::lib::PowerOfTwo blocksize_bytes + ) { + auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes}; + typename StorageType::allocator_type alloc(parameters); + return SQStorageType::compress(data, pool, alloc); } template @@ -254,8 +282,14 @@ struct StorageFactory { using StorageType = LVQStorageType; template - static StorageType init(const svs::data::ConstSimpleDataView& data, Pool& pool) { - return LVQStorageType::compress(data, pool, 0); + static StorageType init( + const svs::data::ConstSimpleDataView& data, + Pool& pool, + svs::lib::PowerOfTwo blocksize_bytes + ) { + auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes}; + typename LVQStorageType::allocator_type alloc(parameters); + return LVQStorageType::compress(data, pool, 0, alloc); } template @@ -288,14 +322,17 @@ struct StorageFactory { static StorageType init( const svs::data::ConstSimpleDataView& data, Pool& pool, + svs::lib::PowerOfTwo blocksize_bytes, size_t leanvec_d = 0, std::optional> matrices = std::nullopt ) { if (leanvec_d == 0) { leanvec_d = (data.dimensions() + 1) / 2; } + auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes}; + typename LeanVecStorageType::allocator_type alloc(parameters); return LeanVecStorageType::reduce( - data, std::move(matrices), pool, 0, svs::lib::MaybeStatic{leanvec_d} + data, std::move(matrices), pool, 0, svs::lib::MaybeStatic{leanvec_d}, alloc ); } diff --git a/bindings/cpp/tests/runtime_test.cpp b/bindings/cpp/tests/runtime_test.cpp index 4fef55f1c..ec6f309d9 100644 --- a/bindings/cpp/tests/runtime_test.cpp +++ b/bindings/cpp/tests/runtime_test.cpp @@ -158,8 +158,9 @@ void write_and_read_index( // Helper that writes and reads and index of requested size // Reports memory usage -UsageInfo -run_save_and_load_test(const size_t target_mibytes, size_t d, size_t graph_max_degree) { +UsageInfo run_save_and_load_test( + const size_t target_mibytes, size_t d, size_t graph_max_degree, size_t blocksize_exp +) { // Generate requested MiB of test data const size_t target_bytes = target_mibytes * 1024 * 1024; const size_t mem_test_n = target_bytes / (d * sizeof(float)); @@ -179,12 +180,16 @@ run_save_and_load_test(const size_t target_mibytes, size_t d, size_t graph_max_d size_t mem_before = get_current_rss(); svs::runtime::v0::DynamicVamanaIndex* index = nullptr; svs::runtime::v0::VamanaIndex::BuildParams build_params{graph_max_degree}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{ + blocksize_exp}; svs::runtime::v0::Status status = svs::runtime::v0::DynamicVamanaIndex::build( &index, d, svs::runtime::v0::MetricType::L2, svs::runtime::v0::StorageKind::FP32, - build_params + build_params, + {}, + dynamic_index_params ); CATCH_REQUIRE(status.ok()); CATCH_REQUIRE(index != nullptr); @@ -232,12 +237,15 @@ CATCH_TEST_CASE("WriteAndReadIndexSVS", "[runtime]") { const auto& test_data = get_test_data(); auto build_func = [](svs::runtime::v0::DynamicVamanaIndex** index) { svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{15}; return svs::runtime::v0::DynamicVamanaIndex::build( index, test_d, svs::runtime::v0::MetricType::L2, svs::runtime::v0::StorageKind::FP32, - build_params + build_params, + {}, + dynamic_index_params ); }; write_and_read_index( @@ -249,12 +257,15 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSFP16", "[runtime]") { const auto& test_data = get_test_data(); auto build_func = [](svs::runtime::v0::DynamicVamanaIndex** index) { svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{16}; return svs::runtime::v0::DynamicVamanaIndex::build( index, test_d, svs::runtime::v0::MetricType::L2, svs::runtime::v0::StorageKind::FP16, - build_params + build_params, + {}, + dynamic_index_params ); }; write_and_read_index( @@ -266,12 +277,15 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSSQI8", "[runtime]") { const auto& test_data = get_test_data(); auto build_func = [](svs::runtime::v0::DynamicVamanaIndex** index) { svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{17}; return svs::runtime::v0::DynamicVamanaIndex::build( index, test_d, svs::runtime::v0::MetricType::L2, svs::runtime::v0::StorageKind::SQI8, - build_params + build_params, + {}, + dynamic_index_params ); }; write_and_read_index( @@ -283,12 +297,15 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSLVQ4x4", "[runtime]") { const auto& test_data = get_test_data(); auto build_func = [](svs::runtime::v0::DynamicVamanaIndex** index) { svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{18}; return svs::runtime::v0::DynamicVamanaIndex::build( index, test_d, svs::runtime::v0::MetricType::L2, svs::runtime::v0::StorageKind::LVQ4x4, - build_params + build_params, + {}, + dynamic_index_params ); }; write_and_read_index( @@ -300,13 +317,16 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSVamanaLeanVec4x4", "[runtime]") { const auto& test_data = get_test_data(); auto build_func = [](svs::runtime::v0::DynamicVamanaIndex** index) { svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{19}; return svs::runtime::v0::DynamicVamanaIndexLeanVec::build( index, test_d, svs::runtime::v0::MetricType::L2, svs::runtime::v0::StorageKind::LeanVec4x4, 32, - build_params + build_params, + {}, + dynamic_index_params ); }; write_and_read_index( @@ -347,6 +367,83 @@ CATCH_TEST_CASE("LeanVecWithTrainingData", "[runtime]") { svs::runtime::v0::DynamicVamanaIndex::destroy(index); } +CATCH_TEST_CASE("LeanVecWithTrainingDataCustomBlockSize", "[runtime]") { + const auto& test_data = get_test_data(); + size_t block_size_exp = 17; // block_size_bytes = 2^block_size_exp + // Build LeanVec index with explicit training + svs::runtime::v0::DynamicVamanaIndex* index = nullptr; + svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{block_size_exp}; + svs::runtime::v0::Status status = svs::runtime::v0::DynamicVamanaIndexLeanVec::build( + &index, + test_d, + svs::runtime::v0::MetricType::L2, + svs::runtime::v0::StorageKind::LeanVec4x4, + 32, + build_params, + {}, + dynamic_index_params + ); + if (!svs::runtime::v0::DynamicVamanaIndex::check_storage_kind( + svs::runtime::v0::StorageKind::LeanVec4x4 + ) + .ok()) { + CATCH_REQUIRE(!status.ok()); + CATCH_SKIP("Storage kind is not supported, skipping test."); + } + CATCH_REQUIRE(status.ok()); + CATCH_REQUIRE(index != nullptr); + + // Add data - should work with provided leanvec dims + std::vector labels(test_n); + std::iota(labels.begin(), labels.end(), 0); + + status = index->add(test_n, labels.data(), test_data.data()); + CATCH_REQUIRE(status.ok()); + + CATCH_REQUIRE(index->blocksize_bytes() == 1u << block_size_exp); + + svs::runtime::v0::DynamicVamanaIndex::destroy(index); +} + +CATCH_TEST_CASE("TrainingDataCustomBlockSize", "[runtime]") { + const auto& test_data = get_test_data(); + size_t block_size_exp = 17; // block_size_bytes = 2^block_size_exp + // Build LeanVec index with explicit training + svs::runtime::v0::DynamicVamanaIndex* index = nullptr; + svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::VamanaIndex::DynamicIndexParams dynamic_index_params{block_size_exp}; + svs::runtime::v0::Status status = svs::runtime::v0::DynamicVamanaIndex::build( + &index, + test_d, + svs::runtime::v0::MetricType::L2, + svs::runtime::v0::StorageKind::FP32, + build_params, + {}, + dynamic_index_params + ); + if (!svs::runtime::v0::DynamicVamanaIndex::check_storage_kind( + svs::runtime::v0::StorageKind::FP32 + ) + .ok()) { + CATCH_REQUIRE(!status.ok()); + CATCH_SKIP("Storage kind is not supported, skipping test."); + } + CATCH_REQUIRE(status.ok()); + CATCH_REQUIRE(index != nullptr); + + // Add data - should work with provided leanvec dims + std::vector labels(test_n); + std::iota(labels.begin(), labels.end(), 0); + + status = index->add(test_n, labels.data(), test_data.data()); + CATCH_REQUIRE(status.ok()); + + CATCH_REQUIRE(index->blocksize_bytes() == 1u << block_size_exp); + + svs::runtime::v0::DynamicVamanaIndex::destroy(index); +} + CATCH_TEST_CASE("FlatIndexWriteAndRead", "[runtime]") { const auto& test_data = get_test_data(); auto build_func = [](svs::runtime::v0::FlatIndex** index) { @@ -470,19 +567,19 @@ CATCH_TEST_CASE("MemoryUsageOnLoad", "[runtime][memory]") { constexpr size_t MiB = 1024 * 1024; CATCH_SECTION("SmallIndex") { - auto stats = run_save_and_load_test(10, 128, 64); + auto stats = run_save_and_load_test(10, 128, 64, 30); CATCH_REQUIRE(stats.file_size < file_threshold(10 * MiB, 128, 64)); CATCH_REQUIRE(stats.rss_increase < rss_threshold(10 * MiB, 1024 * MiB)); } CATCH_SECTION("MediumIndex") { - auto stats = run_save_and_load_test(50, 128, 64); + auto stats = run_save_and_load_test(50, 128, 64, 30); CATCH_REQUIRE(stats.file_size < file_threshold(50 * MiB, 128, 64)); CATCH_REQUIRE(stats.rss_increase < rss_threshold(50 * MiB, 1024 * MiB)); } CATCH_SECTION("LargeIndex") { - auto stats = run_save_and_load_test(200, 128, 64); + auto stats = run_save_and_load_test(200, 128, 64, 30); CATCH_REQUIRE(stats.file_size < file_threshold(200 * MiB, 128, 64)); CATCH_REQUIRE(stats.rss_increase < rss_threshold(200 * MiB, 1024 * MiB)); }