From 39b33304964d543dd95e20c77a3e25576647bcfb Mon Sep 17 00:00:00 2001
From: Philip Adams <35666630+PhilipBAdams@users.noreply.github.com>
Date: Tue, 15 Aug 2023 15:36:57 -0700
Subject: [PATCH] Add convenience functions for parsing the PQ index (#349)

* move read_nodes to public, add get_pq_vector and get_num_points

* clang-format

* Match new private var naming convention

* more private (_) fixes

* VID->vid

* VID->vid cpp
---
 include/pq_flash_index.h | 25 ++++++++++++++-----------
 src/pq_flash_index.cpp   | 12 ++++++++++++
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h
index bc7fe312d..7214bd6a1 100644
--- a/include/pq_flash_index.h
+++ b/include/pq_flash_index.h
@@ -94,6 +94,20 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
 
     DISKANN_DLLEXPORT diskann::Metric get_metric();
 
+    //
+    // node_ids: input list of node_ids to be read
+    // coord_buffers: pointers to pre-allocated buffers that coords need to copied to. If null, dont copy.
+    // nbr_buffers: pre-allocated buffers to copy neighbors into
+    //
+    // returns a vector of bool one for each node_id: true if read is success, else false
+    //
+    DISKANN_DLLEXPORT std::vector<bool> read_nodes(const std::vector<uint32_t> &node_ids,
+                                                   std::vector<T *> &coord_buffers,
+                                                   std::vector<std::pair<uint32_t, uint32_t *>> &nbr_buffers);
+
+    DISKANN_DLLEXPORT std::vector<std::uint8_t> get_pq_vector(std::uint64_t vid);
+    DISKANN_DLLEXPORT uint64_t get_num_points();
+
   protected:
     DISKANN_DLLEXPORT void use_medoids_data_as_centroids();
     DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096);
@@ -121,17 +135,6 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
     // returns region of `node_buf` containing [COORD(T)]
     DISKANN_DLLEXPORT T *offset_to_node_coords(char *node_buf);
 
-    //
-    // node_ids: input list of node_ids to be read
-    // coord_buffers: pointers to pre-allocated buffers that coords need to copied to. If null, dont copy.
-    // nbr_buffers: pre-allocated buffers to copy neighbors into
-    //
-    // returns a vector of bool one for each node_id: true if read is success, else false
-    //
-    DISKANN_DLLEXPORT std::vector<bool> read_nodes(const std::vector<uint32_t> &node_ids,
-                                                   std::vector<T *> &coord_buffers,
-                                                   std::vector<std::pair<uint32_t, uint32_t *>> &nbr_buffers);
-
     // index info for multi-node sectors
     // nhood of node `i` is in sector: [i / nnodes_per_sector]
     // offset in sector: [(i % nnodes_per_sector) * max_node_len]
diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
index 4e5dab7b8..721009830 100644
--- a/src/pq_flash_index.cpp
+++ b/src/pq_flash_index.cpp
@@ -1688,6 +1688,18 @@ template <typename T, typename LabelT> char *PQFlashIndex<T, LabelT>::getHeaderB
 }
 #endif
 
+template <typename T, typename LabelT>
+std::vector<std::uint8_t> PQFlashIndex<T, LabelT>::get_pq_vector(std::uint64_t vid)
+{
+    std::uint8_t *pqVec = &this->data[vid * this->_n_chunks];
+    return std::vector<std::uint8_t>(pqVec, pqVec + this->_n_chunks);
+}
+
+template <typename T, typename LabelT> std::uint64_t PQFlashIndex<T, LabelT>::get_num_points()
+{
+    return _num_points;
+}
+
 // instantiations
 template class PQFlashIndex<uint8_t>;
 template class PQFlashIndex<int8_t>;