diff --git a/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h b/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h index 1d487baf..dce2ae81 100644 --- a/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h +++ b/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h @@ -105,7 +105,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex { size_t leanvec_dims, const VamanaIndex::BuildParams& params, const VamanaIndex::SearchParams& default_search_params, - const VamanaIndex::DynamicIndexParams& dynamic_index_params + const VamanaIndex::DynamicIndexParams& dynamic_index_params, + bool leanvec_primary_only = false ) noexcept; // Specialization to build LeanVec-based Vamana index with provided training data @@ -128,7 +129,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex { const LeanVecTrainingData* training_data, const VamanaIndex::BuildParams& params, const VamanaIndex::SearchParams& default_search_params, - const VamanaIndex::DynamicIndexParams& dynamic_index_params + const VamanaIndex::DynamicIndexParams& dynamic_index_params, + bool leanvec_primary_only = false ) noexcept; }; diff --git a/bindings/cpp/include/svs/runtime/vamana_index.h b/bindings/cpp/include/svs/runtime/vamana_index.h index 2149c0d3..3e9424c6 100644 --- a/bindings/cpp/include/svs/runtime/vamana_index.h +++ b/bindings/cpp/include/svs/runtime/vamana_index.h @@ -114,7 +114,8 @@ struct SVS_RUNTIME_API VamanaIndexLeanVec : public VamanaIndex { StorageKind storage_kind, size_t leanvec_dims, const VamanaIndex::BuildParams& params = {}, - const VamanaIndex::SearchParams& default_search_params = {} + const VamanaIndex::SearchParams& default_search_params = {}, + bool leanvec_primary_only = false ) noexcept; // Specialization to build LeanVec-based Vamana index with provided training data @@ -125,7 +126,8 @@ struct SVS_RUNTIME_API VamanaIndexLeanVec : public VamanaIndex { StorageKind storage_kind, const LeanVecTrainingData* training_data, const VamanaIndex::BuildParams& params = {}, - const VamanaIndex::SearchParams& default_search_params = {} + const VamanaIndex::SearchParams& default_search_params = {}, + bool leanvec_primary_only = false ) noexcept; }; diff --git a/bindings/cpp/src/dynamic_vamana_index.cpp b/bindings/cpp/src/dynamic_vamana_index.cpp index 0c1a6a89..85395a0f 100644 --- a/bindings/cpp/src/dynamic_vamana_index.cpp +++ b/bindings/cpp/src/dynamic_vamana_index.cpp @@ -271,7 +271,8 @@ Status DynamicVamanaIndexLeanVec::build( size_t leanvec_dims, const DynamicVamanaIndex::BuildParams& params, const DynamicVamanaIndex::SearchParams& default_search_params, - const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params + const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params, + bool leanvec_primary_only ) noexcept { using Impl = DynamicVamanaIndexLeanVecImpl; *index = nullptr; @@ -288,7 +289,8 @@ Status DynamicVamanaIndexLeanVec::build( leanvec_dims, params, default_search_params, - dynamic_index_params + dynamic_index_params, + leanvec_primary_only ); *index = new DynamicVamanaIndexManagerBase{std::move(impl)}; }); @@ -303,7 +305,8 @@ Status DynamicVamanaIndexLeanVec::build( const LeanVecTrainingData* training_data, const DynamicVamanaIndex::BuildParams& params, const DynamicVamanaIndex::SearchParams& default_search_params, - const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params + const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params, + bool leanvec_primary_only ) noexcept { using Impl = DynamicVamanaIndexLeanVecImpl; *index = nullptr; @@ -322,7 +325,8 @@ Status DynamicVamanaIndexLeanVec::build( training_data_impl, params, default_search_params, - dynamic_index_params + dynamic_index_params, + leanvec_primary_only ); *index = new DynamicVamanaIndexManagerBase{std::move(impl)}; }); @@ -330,16 +334,34 @@ Status DynamicVamanaIndexLeanVec::build( #else // SVS_RUNTIME_HAVE_LVQ_LEANVEC // 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&, const DynamicVamanaIndex::DynamicIndexParams&) noexcept { +Status DynamicVamanaIndexLeanVec::build( + DynamicVamanaIndex**, + size_t, + MetricType, + StorageKind, + size_t, + const DynamicVamanaIndex::BuildParams&, + const DynamicVamanaIndex::SearchParams&, + const DynamicVamanaIndex::DynamicIndexParams&, + bool +) noexcept { return Status( ErrorCode::NOT_IMPLEMENTED, "DynamicVamanaIndexLeanVec is not supported in this build configuration." ); } -Status DynamicVamanaIndexLeanVec:: - build(DynamicVamanaIndex**, size_t, MetricType, StorageKind, const LeanVecTrainingData*, const DynamicVamanaIndex::BuildParams&, const DynamicVamanaIndex::SearchParams&, const DynamicVamanaIndex::DynamicIndexParams&) noexcept { +Status DynamicVamanaIndexLeanVec::build( + DynamicVamanaIndex**, + size_t, + MetricType, + StorageKind, + const LeanVecTrainingData*, + const DynamicVamanaIndex::BuildParams&, + const DynamicVamanaIndex::SearchParams&, + const DynamicVamanaIndex::DynamicIndexParams&, + bool +) noexcept { return Status( ErrorCode::NOT_IMPLEMENTED, "DynamicVamanaIndexLeanVec is not supported in this build configuration." diff --git a/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h b/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h index 4d59281d..1dc5a092 100644 --- a/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h +++ b/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h @@ -54,11 +54,13 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { const LeanVecTrainingDataImpl& training_data, const VamanaIndex::BuildParams& params, const VamanaIndex::SearchParams& default_search_params, - const VamanaIndex::DynamicIndexParams& dynamic_index_params + const VamanaIndex::DynamicIndexParams& dynamic_index_params, + bool leanvec_primary_only = false ) : 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()} { + , leanvec_matrices_{training_data.get_leanvec_matrices()} + , leanvec_primary_only_{leanvec_primary_only} { check_storage_kind(storage_kind); } @@ -69,11 +71,13 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { size_t leanvec_dims, const VamanaIndex::BuildParams& params, const VamanaIndex::SearchParams& default_search_params, - const VamanaIndex::DynamicIndexParams& dynamic_index_params + const VamanaIndex::DynamicIndexParams& dynamic_index_params, + bool leanvec_primary_only = false ) : DynamicVamanaIndexImpl{dim, metric, storage_kind, params, default_search_params, dynamic_index_params} , leanvec_dims_{leanvec_dims} - , leanvec_matrices_{std::nullopt} { + , leanvec_matrices_{std::nullopt} + , leanvec_primary_only_{leanvec_primary_only} { check_storage_kind(storage_kind); } @@ -124,7 +128,8 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { labels, blocksize_bytes, this->leanvec_dims_, - this->leanvec_matrices_ + this->leanvec_matrices_, + this->leanvec_primary_only_ ); }, data, @@ -136,6 +141,7 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { protected: size_t leanvec_dims_; std::optional leanvec_matrices_; + bool leanvec_primary_only_ = false; StorageKind check_storage_kind(StorageKind kind) { if (!storage::is_leanvec_storage(kind)) { diff --git a/bindings/cpp/src/svs_runtime_utils.h b/bindings/cpp/src/svs_runtime_utils.h index b081c91b..6ce3dc36 100644 --- a/bindings/cpp/src/svs_runtime_utils.h +++ b/bindings/cpp/src/svs_runtime_utils.h @@ -355,13 +355,20 @@ struct StorageFactory { Pool& pool, const Alloc& alloc = {}, size_t leanvec_d = 0, - std::optional> matrices = std::nullopt + std::optional> matrices = std::nullopt, + bool leanvec_primary_only = false ) { if (leanvec_d == 0) { leanvec_d = (data.dimensions() + 1) / 2; } return LeanVecStorageType::reduce( - data, std::move(matrices), pool, 0, svs::lib::MaybeStatic{leanvec_d}, alloc + data, + std::move(matrices), + pool, + 0, + svs::lib::MaybeStatic{leanvec_d}, + alloc, + leanvec_primary_only ); } }; diff --git a/bindings/cpp/src/vamana_index.cpp b/bindings/cpp/src/vamana_index.cpp index c015dd21..6c874853 100644 --- a/bindings/cpp/src/vamana_index.cpp +++ b/bindings/cpp/src/vamana_index.cpp @@ -153,14 +153,21 @@ Status VamanaIndexLeanVec::build( StorageKind storage_kind, size_t leanvec_dims, const VamanaIndex::BuildParams& params, - const VamanaIndex::SearchParams& default_search_params + const VamanaIndex::SearchParams& default_search_params, + bool leanvec_primary_only ) noexcept { using Impl = VamanaIndexLeanVecImpl; *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, + leanvec_primary_only ); *index = new VamanaIndexManagerBase{std::move(impl)}; }); @@ -174,7 +181,8 @@ Status VamanaIndexLeanVec::build( StorageKind storage_kind, const LeanVecTrainingData* training_data, const VamanaIndex::BuildParams& params, - const VamanaIndex::SearchParams& default_search_params + const VamanaIndex::SearchParams& default_search_params, + bool leanvec_primary_only ) noexcept { using Impl = VamanaIndexLeanVecImpl; *index = nullptr; @@ -187,7 +195,13 @@ Status VamanaIndexLeanVec::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, + leanvec_primary_only ); *index = new VamanaIndexManagerBase{std::move(impl)}; }); @@ -195,16 +209,32 @@ Status VamanaIndexLeanVec::build( #else // SVS_RUNTIME_HAVE_LVQ_LEANVEC // LeanVec storage kind is not supported in this build configuration -Status VamanaIndexLeanVec:: - build(VamanaIndex**, size_t, MetricType, StorageKind, size_t, const VamanaIndex::BuildParams&, const VamanaIndex::SearchParams&) noexcept { +Status VamanaIndexLeanVec::build( + VamanaIndex**, + size_t, + MetricType, + StorageKind, + size_t, + const VamanaIndex::BuildParams&, + const VamanaIndex::SearchParams&, + bool +) noexcept { return Status( ErrorCode::NOT_IMPLEMENTED, "VamanaIndexLeanVec is not supported in this build configuration." ); } -Status VamanaIndexLeanVec:: - build(VamanaIndex**, size_t, MetricType, StorageKind, const LeanVecTrainingData*, const VamanaIndex::BuildParams&, const VamanaIndex::SearchParams&) noexcept { +Status VamanaIndexLeanVec::build( + VamanaIndex**, + size_t, + MetricType, + StorageKind, + const LeanVecTrainingData*, + const VamanaIndex::BuildParams&, + const VamanaIndex::SearchParams&, + bool +) noexcept { return Status( ErrorCode::NOT_IMPLEMENTED, "VamanaIndexLeanVec is not supported in this build configuration." diff --git a/bindings/cpp/src/vamana_index_impl.h b/bindings/cpp/src/vamana_index_impl.h index 550257d4..ebada6df 100644 --- a/bindings/cpp/src/vamana_index_impl.h +++ b/bindings/cpp/src/vamana_index_impl.h @@ -515,11 +515,13 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl { StorageKind storage_kind, const LeanVecTrainingDataImpl& training_data, const VamanaIndex::BuildParams& params, - const VamanaIndex::SearchParams& default_search_params + const VamanaIndex::SearchParams& default_search_params, + bool leanvec_primary_only = false ) : VamanaIndexImpl{dim, metric, storage_kind, params, default_search_params} , leanvec_dims_{training_data.get_leanvec_dims()} - , leanvec_matrices_{training_data.get_leanvec_matrices()} { + , leanvec_matrices_{training_data.get_leanvec_matrices()} + , leanvec_primary_only_{leanvec_primary_only} { check_storage_kind(storage_kind); } @@ -529,11 +531,13 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl { StorageKind storage_kind, size_t leanvec_dims, const VamanaIndex::BuildParams& params, - const VamanaIndex::SearchParams& default_search_params + const VamanaIndex::SearchParams& default_search_params, + bool leanvec_primary_only = false ) : VamanaIndexImpl{dim, metric, storage_kind, params, default_search_params} , leanvec_dims_{leanvec_dims} - , leanvec_matrices_{std::nullopt} { + , leanvec_matrices_{std::nullopt} + , leanvec_primary_only_{leanvec_primary_only} { check_storage_kind(storage_kind); } @@ -573,7 +577,8 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl { this->vamana_build_parameters(), data, leanvec_dims_, - leanvec_matrices_ + leanvec_matrices_, + leanvec_primary_only_ ); }, data @@ -584,6 +589,7 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl { protected: size_t leanvec_dims_; std::optional leanvec_matrices_; + bool leanvec_primary_only_ = false; StorageKind check_storage_kind(StorageKind kind) { if (!storage::is_leanvec_storage(kind)) {