Add debug api (#1053)

<!--
Thanks for contributing a pull request! Please ensure you have taken a
look at
the contribution guidelines:
https://github.com/microsoft/DiskANN/blob/main/CONTRIBUTING.md
-->
- [ ] Does this PR have a descriptive title that could go in our release
notes?
- [ ] Does this PR add any new dependencies?
- [ ] Does this PR modify any existing APIs?
- [ ] Is the change to the API backwards compatible?
- [ ] Should this result in any changes to our documentation, either
updating existing docs or adding new ones?

#### Reference Issues/PRs
<!--
Example: Fixes #1234. See also #3456.
Please use keywords (e.g., Fixes) to create link to the issues or pull
requests
you resolved, so that they will automatically be closed when your pull
request
is merged. See
https://github.com/blog/1506-closing-issues-via-pull-requests
-->

#### What does this implement/fix? Briefly explain your changes.

#### Any other comments?

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Sanhaoji2

include/abstract_index.h

@@ -6,6 +6,7 @@
 #include "index_config.h"
 #include "index_build_params.h"
 #include "percentile_stats.h"
+#include "debug_utils.h"
 #include <any>
 
 namespace diskann
@@ -90,6 +91,30 @@ class AbstractIndex
                                                       float *distances,
                                                       std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr);
 
+    // Debug interface: retrieve the raw embedding at internal location index.
+    // Caller must pre-allocate vec with at least the index dimension elements.
+    template <typename data_type>
+    void get_embedding(uint32_t location, data_type *vec);
+
+    // Debug search: runs ANN search and records every traversed node in debug_info.
+    template <typename data_type, typename IDType>
+    std::pair<uint32_t, uint32_t> debug_search(
+        const data_type *query, const size_t K, const uint32_t L,
+        IDType *indices, float *distances,
+        DebugTraversalInfo &debug_info,
+        const uint32_t maxLperSeller = 0,
+        std::function<float(const std::uint8_t *, size_t)> rerank_fn = nullptr);
+
+    // Debug filtered search: same as debug_search with label filtering.
+    template <typename data_type, typename IDType>
+    std::pair<uint32_t, uint32_t> debug_search_with_filters(
+        const data_type *query, const std::vector<std::string> &raw_labels,
+        const size_t K, const uint32_t L,
+        IDType *indices, float *distances,
+        DebugTraversalInfo &debug_info,
+        const uint32_t maxLperSeller = 0,
+        std::function<float(const std::uint8_t *, size_t)> rerank_fn = nullptr);
+
     // insert points with labels, labels should be present for filtered index
     template <typename data_type, typename tag_type>
     int insert_point(const data_type *point, const tag_type tag, const std::vector<std::string> &labels);
@@ -148,5 +173,18 @@ class AbstractIndex
                                      const std::vector<std::string>& filter_labels) = 0;
     virtual void _search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) = 0;
     virtual void _set_universal_label(const LabelType universal_label) = 0;
+    virtual void _get_embedding(uint32_t location, DataType &vec) = 0;
+    virtual std::pair<uint32_t, uint32_t> _debug_search(const DataType &query, const size_t K, const uint32_t L,
+                                                        std::any &indices, float *distances,
+                                                        DebugTraversalInfo &debug_info,
+                                                        const uint32_t maxLperSeller,
+                                                        std::function<float(const std::uint8_t *, size_t)> rerank_fn) = 0;
+    virtual std::pair<uint32_t, uint32_t> _debug_search_with_filters(const DataType &query,
+                                                                      const std::vector<std::string> &raw_labels,
+                                                                      const size_t K, const uint32_t L,
+                                                                      std::any &indices, float *distances,
+                                                                      DebugTraversalInfo &debug_info,
+                                                                      const uint32_t maxLperSeller,
+                                                                      std::function<float(const std::uint8_t *, size_t)> rerank_fn) = 0;
 };
 } // namespace diskann

include/debug_utils.h

@@ -0,0 +1,54 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+#include <limits>
+
+namespace diskann
+{
+
+// Reason why a node visited during ANN graph traversal was or was not included
+// in the final result set.
+enum class FilterReason : uint8_t
+{
+    InResult       = 0, // Node was kept in the final top-K result
+    DistanceTooLarge,   // Node was visited but its distance was too large for top-K
+    LabelMismatch,      // Node was skipped because its label did not match the filter
+};
+
+// Collects per-node traversal information during a debug ANN search.
+// Populated by iterate_to_fixed_point / cached_beam_search when a non-null
+// pointer is passed. Each parallel vector entry corresponds to one node
+// encountered during traversal.
+struct DebugTraversalInfo
+{
+    std::vector<uint32_t> ids;            // Internal location index of each encountered node
+    std::vector<float>    distances;      // PQ/exact distance to query; FLT_MAX when label-rejected
+    std::vector<uint8_t>  label_rejected; // 1 if skipped due to label mismatch, 0 if evaluated
+
+    void clear()
+    {
+        ids.clear();
+        distances.clear();
+        label_rejected.clear();
+    }
+
+    void record_label_rejected(uint32_t id)
+    {
+        ids.push_back(id);
+        distances.push_back(std::numeric_limits<float>::max());
+        label_rejected.push_back(1);
+    }
+
+    void record_visited(uint32_t id, float dist)
+    {
+        ids.push_back(id);
+        distances.push_back(dist);
+        label_rejected.push_back(0);
+    }
+};
+
+} // namespace diskann

include/index.h

@@ -28,6 +28,7 @@
 
 #include "quantized_distance.h"
 #include "pq_data_store.h"
+#include "debug_utils.h"
 
 #define OVERHEAD_FACTOR 1.1
 #define EXPAND_IF_FULL 0
@@ -145,7 +146,8 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     template <typename IDType>
     DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> search(const T *query, const size_t K, const uint32_t L,
                                                            IDType *indices, float *distances = nullptr, const uint32_t maxLperSeller = 0,
-                                                           std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr);
+                                                           std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr,
+                                                           DebugTraversalInfo *debug_info = nullptr);
 
     template <typename IDType>
     std::pair<uint32_t, uint32_t> diverse_search(const T* query, const size_t K, const uint32_t L, const uint32_t maxLperSeller, IDType* indices,
@@ -157,6 +159,30 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
                                               float *distances, std::vector<T *> &res_vectors, bool use_filters,
                                               const std::vector<std::string>& filter_labels);
 
+    // Debug interface: retrieve the raw embedding stored at the given internal location index.
+    // Caller must allocate vec with at least get_aligned_dim() elements of type T.
+    DISKANN_DLLEXPORT void get_embedding(uint32_t location, T *vec) const;
+
+    // Debug search: runs ANN search and records every traversed node.
+    // debug_info is populated in traversal order; call FilterReason helpers to classify entries.
+    template <typename IDType>
+    DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> debug_search(
+        const T *query, const size_t K, const uint32_t L,
+        IDType *indices, float *distances,
+        DebugTraversalInfo &debug_info,
+        const uint32_t maxLperSeller = 0,
+        std::function<float(const std::uint8_t *, size_t)> rerank_fn = nullptr);
+
+    // Debug filtered search: same as debug_search but applies label filtering.
+    template <typename IDType>
+    DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> debug_search_with_filters(
+        const T *query, const std::vector<LabelT> &filter_labels,
+        const size_t K, const uint32_t L,
+        IDType *indices, float *distances,
+        DebugTraversalInfo &debug_info,
+        const uint32_t maxLperSeller = 0,
+        std::function<float(const std::uint8_t *, size_t)> rerank_fn = nullptr);
+
     virtual std::pair<uint32_t, uint32_t> _diverse_search(const DataType& query, const size_t K, const uint32_t L, const uint32_t maxLperSeller,
         std::any& indices, float* distances = nullptr,
         std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr) override;
@@ -166,7 +192,8 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> search_with_filters(const T *query, const std::vector<LabelT> &filter_labels,
                                                                         const size_t K, const uint32_t L, const uint32_t maxLperSeller,
                                                                         IndexType *indices, float *distances,
-                                                                        std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr);
+                                                                        std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr,
+                                                                        DebugTraversalInfo *debug_info = nullptr);
 
     // Will fail if tag already in the index or if tag=0.
     DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag);
@@ -235,6 +262,22 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
                                                                float *distances,
                                                                std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr) override;
 
+    virtual void _get_embedding(uint32_t location, DataType &vec) override;
+
+    virtual std::pair<uint32_t, uint32_t> _debug_search(const DataType &query, const size_t K, const uint32_t L,
+                                                        std::any &indices, float *distances,
+                                                        DebugTraversalInfo &debug_info,
+                                                        const uint32_t maxLperSeller,
+                                                        std::function<float(const std::uint8_t *, size_t)> rerank_fn) override;
+
+    virtual std::pair<uint32_t, uint32_t> _debug_search_with_filters(const DataType &query,
+                                                                      const std::vector<std::string> &raw_labels,
+                                                                      const size_t K, const uint32_t L,
+                                                                      std::any &indices, float *distances,
+                                                                      DebugTraversalInfo &debug_info,
+                                                                      const uint32_t maxLperSeller,
+                                                                      std::function<float(const std::uint8_t *, size_t)> rerank_fn) override;
+
     virtual int _insert_point(const DataType &data_point, const TagType tag) override;
     virtual int _insert_point(const DataType &data_point, const TagType tag, const std::vector<std::string> &labels) override;
 
@@ -293,7 +336,9 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     // The query to use is placed in scratch->aligned_query
     std::pair<uint32_t, uint32_t> iterate_to_fixed_point(InMemQueryScratch<T> *scratch, const uint32_t Lindex,
                                                          const std::vector<uint32_t> &init_ids, bool use_filter,
-                                                         const std::vector<LabelT> &filters, bool search_invocation, uint32_t maxLperSeller = 0);
+                                                         const std::vector<LabelT> &filters, bool search_invocation,
+                                                         uint32_t maxLperSeller = 0,
+                                                         DebugTraversalInfo *debug_info = nullptr);
 
     void search_for_point_and_prune(int location, uint32_t Lindex, std::vector<uint32_t> &pruned_list,
                                     InMemQueryScratch<T> *scratch, bool use_filter = false,

include/pq_flash_index.h

@@ -17,6 +17,7 @@
 #include "tsl/robin_set.h"
 #include "label_bitmask.h"
 #include "integer_label_vector.h"
+#include "debug_utils.h"
 
 #define FULL_PRECISION_REORDER_MULTIPLIER 3
 
@@ -81,10 +82,11 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
     DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search,
                                               uint64_t *res_ids, float *res_dists, const uint64_t beam_width,
                                               const bool use_filter, const std::vector<LabelT> &filter_labels,
-                                              const uint32_t io_limit, uint32_t maxLperSeller = 0, 
+                                              const uint32_t io_limit, uint32_t maxLperSeller = 0,
                                               const bool use_reorder_data = false,
                                               std::function<float(const std::uint8_t*, size_t)> rerank_fn = nullptr,
-                                              QueryStats *stats = nullptr);
+                                              QueryStats *stats = nullptr,
+                                              DebugTraversalInfo *debug_info = nullptr);
 
     DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label);
 
@@ -117,6 +119,25 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
     DISKANN_DLLEXPORT std::vector<std::uint8_t> get_pq_vector(std::uint64_t vid);
     DISKANN_DLLEXPORT uint64_t get_num_points();
 
+    // Debug interface: retrieve full-precision embedding for a given internal node ID.
+    // Caller must pre-allocate vec with at least the index dimension elements (get_data_dim()).
+    DISKANN_DLLEXPORT void get_embedding(uint32_t id, T *vec);
+
+    // Debug search: runs ANN search and records every traversed node with a FilterReason.
+    DISKANN_DLLEXPORT void debug_search(
+        const T *query, const uint64_t k_search, const uint64_t l_search,
+        uint64_t *res_ids, float *res_dists, const uint64_t beam_width,
+        DebugTraversalInfo &debug_info,
+        uint32_t maxLperSeller = 0);
+
+    // Debug filtered search: same as debug_search but applies label filtering.
+    DISKANN_DLLEXPORT void debug_search_with_filters(
+        const T *query, const uint64_t k_search, const uint64_t l_search,
+        uint64_t *res_ids, float *res_dists, const uint64_t beam_width,
+        const std::vector<LabelT> &filter_labels,
+        DebugTraversalInfo &debug_info,
+        uint32_t maxLperSeller = 0);
+
   protected:
     DISKANN_DLLEXPORT void use_medoids_data_as_centroids();
     DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096);