[ENH] Add API for listing all nodes

CMakeLists.txt

@@ -59,5 +59,8 @@ if(CMAKE_PROJECT_NAME STREQUAL PROJECT_NAME)
 
     add_executable(persistent_test tests/cpp/persistent_test.cpp)
     target_link_libraries(persistent_test hnswlib)
+
+    add_executable(api_tests tests/cpp/api_test.cpp)
+    target_link_libraries(api_tests hnswlib)
 endif()
 

hnswlib/hnswalg.h

@@ -1181,6 +1181,45 @@ namespace hnswlib
             return;
         }
 
+        std::pair<size_t, size_t> getLabelCounts() const
+        {
+            size_t non_deleted_labels = 0;
+            size_t deleted_labels = 0;
+            std::unique_lock<std::mutex> label_lock(label_lookup_lock);
+            for (auto it = label_lookup_.begin(); it != label_lookup_.end(); ++it)
+            {
+                if (!isMarkedDeleted(it->second))
+                {
+                    non_deleted_labels += 1;
+                }
+                else
+                {
+                    deleted_labels += 1;
+                }
+            }
+            return std::make_pair(non_deleted_labels, deleted_labels);
+        }
+
+        // Get all labels, segregated by deleted and non deleted.
+        std::pair<std::vector<labeltype>, std::vector<labeltype>> getAllLabels() const
+        {
+            std::vector<labeltype> labels;
+            std::vector<labeltype> deleted_labels;
+            std::unique_lock<std::mutex> label_lock(label_lookup_lock);
+            for (auto it = label_lookup_.begin(); it != label_lookup_.end(); ++it)
+            {
+                if (!isMarkedDeleted(it->second))
+                {
+                    labels.push_back(it->first);
+                }
+                else
+                {
+                    deleted_labels.push_back(it->first);
+                }
+            }
+            return std::make_pair(labels, deleted_labels);
+        }
+
         template <typename data_t>
         std::vector<data_t> getDataByLabel(labeltype label) const
         {

tests/cpp/api_test.cpp

@@ -0,0 +1,78 @@
+#include "../../hnswlib/hnswlib.h"
+
+#include <assert.h>
+
+void testListAllLabels()
+{
+    int d = 1536;
+    hnswlib::labeltype n = 1000;
+
+    std::vector<float> data(n * d);
+
+    std::mt19937 rng;
+    rng.seed(47);
+    std::uniform_real_distribution<> distrib;
+
+    for (auto i = 0; i < n * d; i++)
+    {
+        data[i] = distrib(rng);
+    }
+
+    hnswlib::InnerProductSpace space(d);
+    hnswlib::HierarchicalNSW<float> *alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, n, 16, 200, 100, false, false, true, ".");
+
+    for (size_t i = 0; i < n; i++)
+    {
+        alg_hnsw->addPoint(data.data() + d * i, i);
+    }
+    // Delete odd points.
+    for (size_t i = 1; i < n; i += 2)
+    {
+        alg_hnsw->markDelete(i);
+    }
+    // Get all data.
+    auto res = alg_hnsw->getAllLabels();
+    auto non_deleted = res.first;
+    auto deleted = res.second;
+    assert(non_deleted.size() == n / 2);
+    assert(deleted.size() == n / 2);
+
+    for (auto idx : non_deleted)
+    {
+        assert(idx % 2 == 0);
+    }
+    for (auto idx : deleted)
+    {
+        assert(idx % 2 == 1);
+    }
+
+    // After persisting and reloading the data should be the same.
+    alg_hnsw->persistDirty();
+
+    // Load the index with n elements
+    hnswlib::HierarchicalNSW<float> *alg_hnsw2 = new hnswlib::HierarchicalNSW<float>(&space, ".", false, n, false, false, true);
+
+    // Check that all the data is the same
+    auto res2 = alg_hnsw2->getAllLabels();
+    auto non_deleted2 = res2.first;
+    auto deleted2 = res2.second;
+    assert(non_deleted2.size() == n / 2);
+    assert(deleted2.size() == n / 2);
+
+    for (auto idx : non_deleted2)
+    {
+        assert(idx % 2 == 0);
+    }
+    for (auto idx : deleted2)
+    {
+        assert(idx % 2 == 1);
+    }
+}
+
+int main()
+{
+    std::cout << "Testing ..." << std::endl;
+    testListAllLabels();
+    std::cout << "Test testListAllLabels ok" << std::endl;
+    return 0;
+}