Skip to content

[MOD-18075] Allow custom HNSW neighbor debug implementations #1024

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
dor-forer wants to merge 1 commit into
base: main
Choose a base branch
from
+14 −58
Draft

[MOD-18075] Allow custom HNSW neighbor debug implementations #1024

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/VecSim/algorithms/hnsw/hnsw.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,7 +285,7 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
return idToMetaData.data() + internal_id;
}
vecsim_stl::vector<graphNodeType> safeCollectAllNodeIncomingNeighbors(idType node_id) const;
VecSimDebugCommandCode getHNSWElementNeighbors(size_t label, int ***neighborsData);
VecSimDebugCommandCode getHNSWElementNeighbors(size_t label, int ***neighborsData) override;
void insertElementToGraph(idType element_id, size_t element_max_level, idType entry_point,
size_t global_max_level, const void *vector_data);
void removeVectorInPlace(idType id);
Expand Down
2 changes: 1 addition & 1 deletion src/VecSim/algorithms/hnsw/hnsw_tiered.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
this->getHNSWIndex()->unlockSharedIndexDataGuard();
}

VecSimDebugCommandCode getHNSWElementNeighbors(size_t label, int ***neighborsData) {
VecSimDebugCommandCode getHNSWElementNeighbors(size_t label, int ***neighborsData) override {
this->mainIndexGuard.lock_shared();
auto res = this->getHNSWIndex()->getHNSWElementNeighbors(label, neighborsData);
this->mainIndexGuard.unlock_shared();
Expand Down
57 changes: 1 addition & 56 deletions src/VecSim/vec_sim_debug.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,65 +8,10 @@
*/
#include "vec_sim_debug.h"
#include "VecSim/vec_sim_index.h"
#include "VecSim/algorithms/hnsw/hnsw.h"
#include "VecSim/algorithms/hnsw/hnsw_tiered.h"
#include "VecSim/types/bfloat16.h"

extern "C" int VecSimDebug_GetElementNeighborsInHNSWGraph(VecSimIndex *index, size_t label,
int ***neighborsData) {

// Set as if we return an error, and upon success we will set the pointers appropriately.
*neighborsData = nullptr;
VecSimIndexBasicInfo info = index->basicInfo();
if (info.algo != VecSimAlgo_HNSWLIB) {
return VecSimDebugCommandCode_BadIndex;
}
if (!info.isTiered) {
if (info.type == VecSimType_FLOAT32) {
return dynamic_cast<HNSWIndex<float, float> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else if (info.type == VecSimType_FLOAT64) {
return dynamic_cast<HNSWIndex<double, double> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else if (info.type == VecSimType_BFLOAT16) {
return dynamic_cast<HNSWIndex<vecsim_types::bfloat16, float> *>(index)
->getHNSWElementNeighbors(label, neighborsData);
} else if (info.type == VecSimType_FLOAT16) {
return dynamic_cast<HNSWIndex<vecsim_types::float16, float> *>(index)
->getHNSWElementNeighbors(label, neighborsData);
} else if (info.type == VecSimType_INT8) {
return dynamic_cast<HNSWIndex<int8_t, float> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else if (info.type == VecSimType_UINT8) {
return dynamic_cast<HNSWIndex<uint8_t, float> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else {
assert(false && "Invalid data type");
}
} else {
if (info.type == VecSimType_FLOAT32) {
return dynamic_cast<TieredHNSWIndex<float, float> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else if (info.type == VecSimType_FLOAT64) {
return dynamic_cast<TieredHNSWIndex<double, double> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else if (info.type == VecSimType_BFLOAT16) {
return dynamic_cast<TieredHNSWIndex<vecsim_types::bfloat16, float> *>(index)
->getHNSWElementNeighbors(label, neighborsData);
} else if (info.type == VecSimType_FLOAT16) {
return dynamic_cast<TieredHNSWIndex<vecsim_types::float16, float> *>(index)
->getHNSWElementNeighbors(label, neighborsData);
} else if (info.type == VecSimType_INT8) {
return dynamic_cast<TieredHNSWIndex<int8_t, float> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else if (info.type == VecSimType_UINT8) {
return dynamic_cast<TieredHNSWIndex<uint8_t, float> *>(index)->getHNSWElementNeighbors(
label, neighborsData);
} else {
assert(false && "Invalid data type");
}
}
return VecSimDebugCommandCode_BadIndex;
return index->getHNSWElementNeighbors(label, neighborsData);
}

extern "C" void VecSimDebug_ReleaseElementNeighborsInHNSWGraph(int **neighborsData) {
Expand Down
11 changes: 11 additions & 0 deletions src/VecSim/vec_sim_interface.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,17 @@ struct VecSimIndexInterface : public VecsimBaseObject {
*/
virtual VecSimDebugInfoIterator *debugInfoIterator() const = 0;

/**
* @brief Return the HNSW graph neighbors for a label, grouped by level.
*
* Non-HNSW indexes return VecSimDebugCommandCode_BadIndex. The caller owns successful
* output and must release it with VecSimDebug_ReleaseElementNeighborsInHNSWGraph.
*/
virtual VecSimDebugCommandCode getHNSWElementNeighbors(size_t label, int ***neighborsData) {
*neighborsData = nullptr;
return VecSimDebugCommandCode_BadIndex;
}

/**
* @brief A function to be implemented by the inheriting index and called by rangeQuery.
* @param queryBlob binary representation of the query vector. Blob size should match the index
Expand Down
Loading