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

@@ -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);
@@ -120,17 +134,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]

src/pq_flash_index.cpp

@@ -1639,6 +1639,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>;