feature: get_distance function and test in dynamic_index.h (#102)

This PR solves issue
#99
Added new get_distance function in dynamic_index.h to get distance
between a vector in index and a normal vecor

Author: yuejiaointel

bindings/python/include/svs/python/vamana.h

@@ -213,6 +213,29 @@ Calling this method should not affect recall.)"
     )"
     );
 
+    manager.def(
+        "get_distance",
+        [](Manager& self, size_t external_id, py_contiguous_array_t<float> query_array) {
+            // Get raw pointer + size from the Python array
+            const float* data_ptr = query_array.data();
+            size_t n = query_array.size();
+            std::vector<float> vec(data_ptr, data_ptr + n);
+            return self.get_distance(external_id, vec);
+        },
+        pybind11::arg("external_id"),
+        pybind11::arg("query_vector"),
+        R"(
+    Compute the distance between the stored vector at `external_id` and the provided `query_vector`.
+
+    Args:
+        external_id: the external ID of the vector in the index
+        query_vector: a 1-D contiguous array whose length must match the index dimensionality
+
+    Returns:
+        float: the computed distance
+    )"
+    );
+
     ///// Experiemntal Interfaces
     add_experimental_calibration<svs::Float16>(manager);
     add_experimental_calibration<float>(manager);

bindings/python/src/flat.cpp

@@ -274,6 +274,30 @@ queries-per-second).
 
 See also: `svs.FlatSearchParameters`.)"
     );
+
+    for_standard_specializations([&flat]<typename Q, typename T, size_t N>() {
+        flat.def(
+            "get_distance",
+            [](svs::Flat& self, size_t external_id, py_contiguous_array_t<T> query_array) {
+                const T* data_ptr = query_array.data();
+                size_t len = query_array.size();
+                std::vector<T> vec(data_ptr, data_ptr + len);
+                return self.get_distance(external_id, vec);
+            },
+            py::arg("external_id"),
+            py::arg("query_vector"),
+            R"(
+    Compute the distance between the stored vector at `external_id` and the provided `query_vector`.
+
+    Args:
+        external_id: the external ID of the vector in the index
+        query_vector: a 1-D contiguous array whose length must match the index dimensionality
+
+    Returns:
+        float: the computed distance
+    )"
+        );
+    });
 }
 
 } // namespace svs::python::flat

bindings/python/tests/common.py

@@ -123,3 +123,53 @@ def test_threading(f, *args, validate = None, iters = 4, print_times = False):
     # For short lived processes, we generally see closer to a 3x speedup than a 4x
     # speedup when using 4 threads.
     testcase.assertTrue(1.3 * new_time < base_time)
+
+def test_get_distance(index, distance, data = svs.read_vecs(test_data_vecs), test_distance = True):
+    """
+    Test the get_distance method of an index by comparing its results with direct distance computation.
+
+    Arguments:
+        index: The SVS index object with get_distance method
+        distance: The distance type
+        data: The dataset used to build the index
+        test_distance: Whether to perform the distance test
+    """
+    # Skip get_distance_test if flag is set
+    if not test_distance:
+        return
+
+    tolerance=1e-2
+    query_id = 10
+    index_id = 100
+    dt = data.dtype
+    query_vector_raw = np.array(data[query_id], dtype=dt)
+    indexed_vector_raw = np.array(data[index_id], dtype=dt)
+    index_distance  = index.get_distance(index_id, query_vector_raw)
+    # Up cast to avoid overflow
+    query_vector = query_vector_raw.astype(np.float32)
+    indexed_vector = indexed_vector_raw.astype(np.float32)
+
+    # Compute distance based on distance type
+    if distance == svs.DistanceType.L2:
+        expected_distance = np.sum((query_vector - indexed_vector) ** 2)
+    elif distance == svs.DistanceType.MIP:
+        expected_distance = np.dot(query_vector, indexed_vector)
+    elif distance == svs.DistanceType.Cosine:
+        qn = np.linalg.norm(query_vector)
+        vn = np.linalg.norm(indexed_vector)
+        if qn == 0 or vn == 0:
+            expected_distance = 0.0
+        else:
+            expected_distance = (np.dot(query_vector, indexed_vector) / (qn * vn))
+    else:
+        raise ValueError(f"Unsupported DistanceType: {distance}")
+
+    relative_diff = abs((index_distance - expected_distance) / expected_distance)
+    assert relative_diff < tolerance
+
+    # Test out of bounds ID
+    try:
+        index.get_distance(index_id + 99999, query_vector_raw)
+        assert False, "Should have exception"
+    except Exception as e:
+        pass

bindings/python/tests/test_flat.py

@@ -29,7 +29,8 @@
     test_groundtruth_l2, \
     test_groundtruth_mip, \
     test_number_of_vectors, \
-    test_dimensions
+    test_dimensions, \
+    test_get_distance
 
 class FlatTester(unittest.TestCase):
     """
@@ -54,7 +55,7 @@ def _loaders(self, file: svs.VectorDataLoader):
             }),
         ]
 
-    def _do_test(self, flat, queries, groundtruth, expected_recall = 1.0):
+    def _do_test(self, flat, queries, groundtruth, distance, data = svs.read_vecs(test_data_vecs), expected_recall = 1.0, test_distance = True):
         """
         Perform a series of tests on a Flat index to test its conformance to expectations.
         Parameters:
@@ -67,6 +68,9 @@ def _do_test(self, flat, queries, groundtruth, expected_recall = 1.0):
             - Results of `search` are within acceptable margins of the groundtruth.
             - The number of threads can be changed with an observable side-effect.
         """
+        # Test get distance
+        test_get_distance(flat, distance, data, test_distance)
+
         # Data interface
         self.assertEqual(flat.size, test_number_of_vectors)
         self.assertEqual(flat.dimensions, test_dimensions)
@@ -117,7 +121,7 @@ def _do_test_from_file(self, distance: svs.DistanceType, queries, groundtruth):
             svs.VectorDataLoader(
                 test_data_svs, svs.DataType.float32, dims = test_data_dims
             )
-        );
+        )
         for loader, recall in loaders:
             index = svs.Flat(
                 loader,
@@ -126,7 +130,7 @@ def _do_test_from_file(self, distance: svs.DistanceType, queries, groundtruth):
             )
 
             self.assertEqual(index.num_threads, num_threads)
-            self._do_test(index, queries, groundtruth, expected_recall = recall[distance])
+            self._do_test(index, queries, groundtruth, distance, expected_recall = recall[distance])
 
     def test_from_file(self):
         """
@@ -154,21 +158,22 @@ def test_from_array(self):
         # Test `float32`
         print("Flat, From Array, Float32")
         flat = svs.Flat(data_f32, svs.DistanceType.L2)
-        self._do_test(flat, queries_f32, groundtruth)
+        self._do_test(flat, queries_f32, groundtruth, svs.DistanceType.L2, data_f32)
 
         # Test `float16`
         print("Flat, From Array, Float16")
         data_f16 = data_f32.astype('float16')
         queries_f16 = queries_f32.astype('float16')
         flat = svs.Flat(data_f16, svs.DistanceType.L2)
-        self._do_test(flat, queries_f16, groundtruth)
+        # Do not test get distance for fp16 data as py_contiguous_array_t does not support it
+        self._do_test(flat, queries_f16, groundtruth, svs.DistanceType.L2, data_f16, test_distance = False)
 
         # Test `int8`
         print("Flat, From Array, Int8")
         data_i8 = data_f32.astype('int8')
         queries_i8 = queries_f32.astype('int8')
         flat = svs.Flat(data_i8, svs.DistanceType.L2)
-        self._do_test(flat, queries_i8, groundtruth)
+        self._do_test(flat, queries_i8, groundtruth, svs.DistanceType.L2, data=data_i8)
 
         # Test 'uint8'
         # The dataset is stored as values that can be encoded as `int8`.
@@ -178,4 +183,4 @@ def test_from_array(self):
         data_u8 = (data_f32 + 128).astype('uint8')
         queries_u8 = (queries_f32 + 128).astype('uint8')
         flat = svs.Flat(data_u8, svs.DistanceType.L2)
-        self._do_test(flat, queries_u8, groundtruth)
+        self._do_test(flat, queries_u8, groundtruth, svs.DistanceType.L2, data=data_u8)

bindings/python/tests/test_vamana.py

@@ -39,7 +39,8 @@
     test_vamana_reference, \
     test_number_of_vectors, \
     test_dimensions, \
-    get_test_set
+    get_test_set, \
+    test_get_distance
 
 from .dataset import UncompressedMatcher
 
@@ -305,6 +306,9 @@ def _test_build(
         vamana = svs.Vamana.build(params, loader, distance, num_threads = num_threads)
         print(f"Building: {vamana.experimental_backend_string}")
 
+        # Test get distance
+        test_get_distance(vamana, distance)
+
         groundtruth_map = self._groundtruth_map()
         # Load the queries and groundtruth
         queries = svs.read_vecs(test_queries)

include/svs/index/flat/flat.h

@@ -70,6 +70,39 @@ data::GetDatumAccessor svs_invoke(svs::tag_t<accessor>, const Data& SVS_UNUSED(d
     return data::GetDatumAccessor{};
 }
 
+/////
+///// Distance
+/////
+struct ComputeDistanceType {
+    template <typename Data, typename Distance, typename Query>
+    double operator()(
+        const Data& data, const Distance& distance, size_t internal_id, const Query& query
+    ) const {
+        return svs_invoke(*this, data, distance, internal_id, query);
+    }
+};
+// CPO for distance computation
+inline constexpr ComputeDistanceType get_distance_ext{};
+template <typename Data, typename Distance, typename Query>
+double svs_invoke(
+    svs::tag_t<get_distance_ext>,
+    const Data& data,
+    const Distance& distance,
+    size_t internal_id,
+    const Query& query
+) {
+    // Get distance
+    auto dist_f = extensions::distance(data, distance);
+    svs::distance::maybe_fix_argument(dist_f, query);
+
+    // Get the vector from the index
+    auto indexed_span = data.get_datum(internal_id);
+
+    // Compute the distance using the appropriate distance function
+    auto dist = svs::distance::compute(dist_f, query, indexed_span);
+
+    return static_cast<double>(dist);
+}
 } // namespace extensions
 
 // The flat index is "special" because we wish to enable the `FlatIndex` to either:
@@ -455,6 +488,30 @@ class FlatIndex {
     /// @brief Return the current thread pool handle.
     ///
     threads::ThreadPoolHandle& get_threadpool_handle() { return threadpool_; }
+
+    ///// Distance
+
+    /// @brief Compute the distance between an external vector and a vector in the index.
+    template <typename Query> double get_distance(size_t id, const Query& query) const {
+        // Check if id is valid
+        if (id >= size()) {
+            throw ANNEXCEPTION("ID {} is out of bounds for index of size {}!", id, size());
+        }
+
+        // Verify dimensions match
+        const size_t query_size = query.size();
+        const size_t index_vector_size = dimensions();
+        if (query_size != index_vector_size) {
+            throw ANNEXCEPTION(
+                "Incompatible dimensions. Query has {} while the index expects {}.",
+                query_size,
+                index_vector_size
+            );
+        }
+
+        // Call extension for distance computation
+        return svs::index::flat::extensions::get_distance_ext(data_, distance_, id, query);
+    }
 };
 
 ///

include/svs/index/vamana/dynamic_index.h

@@ -23,6 +23,7 @@
 #include "svs/index/flat/flat.h"
 
 // svs
+#include "svs/concepts/distance.h"
 #include "svs/core/data.h"
 #include "svs/core/distance.h"
 #include "svs/core/graph.h"
@@ -1214,6 +1215,35 @@ class MutableVamanaIndex {
             }
         }
     }
+
+    ///// Distance
+
+    /// @brief Compute the distance between an external vector and a vector in the index.
+    template <typename ExternalId, typename Query>
+    double get_distance(const ExternalId& external_id, const Query& query) const {
+        // Check if the external ID exists
+        if (!has_id(external_id)) {
+            throw ANNEXCEPTION(
+                "ID {} is out of bounds for index of size {}!", external_id, size()
+            );
+        }
+        // Verify dimensions match
+        const size_t query_size = query.size();
+        const size_t index_vector_size = dimensions();
+        if (query_size != index_vector_size) {
+            throw ANNEXCEPTION(
+                "Incompatible dimensions. Query has {} while the index expects {}.",
+                query_size,
+                index_vector_size
+            );
+        }
+
+        // Translate external ID to internal ID
+        auto internal_id = translate_external_id(external_id);
+
+        // Call extension for distance computation
+        return extensions::get_distance_ext(data_, distance_, internal_id, query);
+    }
 };
 
 ///// Deduction Guides.

include/svs/index/vamana/extensions.h

@@ -591,4 +591,37 @@ SVS_FORCE_INLINE data::GetDatumAccessor svs_invoke(
     return data::GetDatumAccessor();
 }
 
+/////
+///// Distance
+/////
+struct ComputeDistanceType {
+    template <typename Data, typename Distance, typename Query>
+    double operator()(
+        const Data& data, const Distance& distance, size_t internal_id, const Query& query
+    ) const {
+        return svs_invoke(*this, data, distance, internal_id, query);
+    }
+};
+// CPO for distance computation
+inline constexpr ComputeDistanceType get_distance_ext{};
+template <typename Data, typename Distance, typename Query>
+double svs_invoke(
+    svs::tag_t<get_distance_ext>,
+    const Data& data,
+    const Distance& distance,
+    size_t internal_id,
+    const Query& query
+) {
+    // Get distance
+    auto dist_f = single_search_setup(data, distance);
+    svs::distance::maybe_fix_argument(dist_f, query);
+
+    // Get the vector from the index
+    auto indexed_span = data.get_datum(internal_id);
+
+    // Compute the distance using the appropriate distance function
+    auto dist = svs::distance::compute(dist_f, query, indexed_span);
+    return static_cast<double>(dist);
+}
+
 } // namespace svs::index::vamana::extensions

include/svs/index/vamana/index.h

@@ -860,6 +860,29 @@ class VamanaIndex {
     template <typename F> void experimental_escape_hatch(F&& f) const {
         std::invoke(SVS_FWD(f), graph_, data_, distance_, lib::as_const_span(entry_point_));
     }
+
+    ///// Distance
+
+    /// @brief Compute the distance between a vector in the index and a query vector
+    template <typename Query> double get_distance(size_t id, const Query& query) const {
+        // Check if id is valid
+        if (id >= size()) {
+            throw ANNEXCEPTION("ID {} is out of bounds for index of size {}!", id, size());
+        }
+        // Verify dimensions match
+        const size_t query_size = query.size();
+        const size_t index_vector_size = dimensions();
+        if (query_size != index_vector_size) {
+            throw ANNEXCEPTION(
+                "Incompatible dimensions. Query has {} while the index expects {}.",
+                query_size,
+                index_vector_size
+            );
+        }
+
+        // Call extension for distance computation
+        return extensions::get_distance_ext(data_, distance_, id, query);
+    }
 };
 
 // Shared documentation for assembly methods.