Add a batch interface for getDistanceByLabel

include/vsag/index.h

@@ -249,6 +249,23 @@ class Index {
         throw std::runtime_error("Index doesn't support get distance by id");
     };
 
+    /**
+     * @brief Calculate the distance between the query and the vector of the given ID for batch.
+     *
+     * @param count is the count of vids
+     * @param vids is the unique identifier of the vector to be calculated in the index.
+     * @param vector is the embedding of query
+     * @param distances is the distances between the query and the vector of the given ID
+     * @return result is valid distance of input vids.
+     */
+    virtual tl::expected<int64_t, Error>
+    CalcBatchDistanceById(int64_t count, 
+                          const int64_t *vids, 
+                          const float* vector,
+                          float *&distances) const {
+        throw std::runtime_error("Index doesn't support get distance by id");
+    };
+
     /**
      * @brief Checks if the specified feature is supported by the index.
      *

src/algorithm/hnswlib/algorithm_interface.h

@@ -69,6 +69,12 @@ class AlgorithmInterface {
     virtual float
     getDistanceByLabel(LabelType label, const void* data_point) = 0;
 
+    virtual int64_t 
+    getBatchDistanceByLabel(int64_t count, 
+                            const int64_t *vids, 
+                            const void *data_point, 
+                            float *&distances) = 0;
+
     virtual const float*
     getDataByLabel(LabelType label) const = 0;
 

src/algorithm/hnswlib/hnswalg.cpp

@@ -171,6 +171,28 @@ HierarchicalNSW::getDistanceByLabel(LabelType label, const void* data_point) {
     return dist;
 }
 
+int64_t
+HierarchicalNSW::getBatchDistanceByLabel(int64_t count, 
+                                         const int64_t *vids, 
+                                         const void* data_point, 
+                                         float *&distances) {
+    std::shared_lock lock_table(label_lookup_lock_);
+    int64_t ret_cnt = 0;
+    distances = (float *)allocator_->Allocate(sizeof(float) * count);
+    for (int i = 0; i < count; i++) {
+        auto search = label_lookup_.find(vids[i]);
+        if (search == label_lookup_.end()) {
+            distances[i] = -1;
+        } else {
+            InnerIdType internal_id = search->second;
+            float dist = fstdistfunc_(data_point, getDataByInternalId(internal_id), dist_func_param_);
+            distances[i] = dist;
+            ret_cnt++;
+        }
+    }
+    return ret_cnt;
+}
+
 bool
 HierarchicalNSW::isValidLabel(LabelType label) {
     std::shared_lock lock_table(label_lookup_lock_);

src/algorithm/hnswlib/hnswalg.h

@@ -146,6 +146,12 @@ class HierarchicalNSW : public AlgorithmInterface<float> {
     float
     getDistanceByLabel(LabelType label, const void* data_point) override;
 
+    int64_t
+    getBatchDistanceByLabel(int64_t count, 
+                            const int64_t *vids, 
+                            const void* data_point,
+                            float *&distances) override;
+
     bool
     isValidLabel(LabelType label) override;
 

src/algorithm/hnswlib/hnswalg_static.h

@@ -81,6 +81,7 @@ class StaticHierarchicalNSW : public AlgorithmInterface<float> {
     void* dist_func_param_{nullptr};
 
     mutable std::mutex label_lookup_lock;  // lock for label_lookup_
+    mutable std::shared_mutex shared_label_lookup_lock;
     std::unordered_map<LabelType, tableint> label_lookup_;
 
     std::default_random_engine level_generator_;
@@ -262,6 +263,28 @@ class StaticHierarchicalNSW : public AlgorithmInterface<float> {
         return dist;
     }
 
+    int64_t
+    getBatchDistanceByLabel(int64_t count, 
+                            const int64_t *vids, 
+                            const void* data_point, 
+                            float *&distances) override {
+        std::shared_lock<std::shared_mutex> lock_table(shared_label_lookup_lock);
+        int64_t ret_cnt = 0;
+        distances = (float *)allocator_->Allocate(sizeof(float) * count);
+        for (int i = 0; i < count; i++) {
+            auto search = label_lookup_.find(vids[i]);
+            if (search == label_lookup_.end()) {
+                distances[i] = -1;
+            } else {
+                InnerIdType internal_id = search->second;
+                float dist = fstdistfunc_(data_point, getDataByInternalId(internal_id), dist_func_param_);
+                distances[i] = dist;
+                ret_cnt++;
+            }
+        }
+        return ret_cnt;
+    }
+
     bool
     isValidLabel(LabelType label) override {
         std::unique_lock<std::mutex> lock_table(label_lookup_lock);

src/index/hnsw.h

@@ -145,6 +145,14 @@ class HNSW : public Index {
         SAFE_CALL(return alg_hnsw_->getDistanceByLabel(id, vector));
     };
 
+    virtual tl::expected<int64_t, Error>
+    CalcBatchDistanceById(int64_t count, 
+                          const int64_t *vids, 
+                          const float* vector,
+                          float *&distances) const override {
+        SAFE_CALL(return alg_hnsw_->getBatchDistanceByLabel(count, vids, vector, distances));
+    };
+
     [[nodiscard]] bool
     CheckFeature(IndexFeature feature) const override;