diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index d2c26fddf..6f91e610c 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -94,6 +94,20 @@ template 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 read_nodes(const std::vector &node_ids, + std::vector &coord_buffers, + std::vector> &nbr_buffers); + + DISKANN_DLLEXPORT std::vector 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 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 read_nodes(const std::vector &node_ids, - std::vector &coord_buffers, - std::vector> &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 520cfc59a..1ca8ae64e 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1639,6 +1639,18 @@ template char *PQFlashIndex::getHeaderB } #endif +template +std::vector PQFlashIndex::get_pq_vector(std::uint64_t vid) +{ + std::uint8_t *pqVec = &this->data[vid * this->_n_chunks]; + return std::vector(pqVec, pqVec + this->_n_chunks); +} + +template std::uint64_t PQFlashIndex::get_num_points() +{ + return _num_points; +} + // instantiations template class PQFlashIndex; template class PQFlashIndex;