Vlad/l2 sumofsquares

apis/python/src/tiledb/vector_search/flat_index.py

@@ -151,7 +151,7 @@ def create(
     group_exists: bool = False,
     config: Optional[Mapping[str, Any]] = None,
     storage_version: str = STORAGE_VERSION,
-    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.L2,
+    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     **kwargs,
 ) -> FlatIndex:
     """

apis/python/src/tiledb/vector_search/index.py

@@ -61,12 +61,15 @@ def __init__(
         self.storage_version = self.group.meta.get("storage_version", "0.1")
         try:
             self.distance_metric = vspy.DistanceMetric(
-                self.group.meta.get("distance_metric", vspy.DistanceMetric.L2)
+                self.group.meta.get(
+                    "distance_metric", vspy.DistanceMetric.SUM_OF_SQUARES
+                )
             )
         except ValueError:
             raise ValueError(
                 f"Invalid distance metric in metadata: {self.group.meta.get('distance_metric')}."
             )
+
         if (
             not storage_formats[self.storage_version]["SUPPORT_TIMETRAVEL"]
             and timestamp is not None

apis/python/src/tiledb/vector_search/ingestion.py

@@ -84,7 +84,7 @@ def ingest(
     ] = None,
     write_centroids_resources: Optional[Mapping[str, Any]] = None,
     partial_index_resources: Optional[Mapping[str, Any]] = None,
-    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.L2,
+    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     normalized: bool = False,
     **kwargs,
 ):
@@ -197,7 +197,7 @@ def ingest(
     partial_index_resources: Optional[Mapping[str, Any]]
         Resources to request when performing the computation of partial indexing, only applies to BATCH mode
     distance_metric: vspy.DistanceMetric
-        Distance metric to use for the index, defaults to 'vspy.DistanceMetric.L2'. Options are 'vspy.DistanceMetric.L2', 'vspy.DistanceMetric.INNER_PRODUCT', 'vspy.DistanceMetric.COSINE'.
+        Distance metric to use for the index, defaults to 'vspy.DistanceMetric.SUM_OF_SQUARES'. Options are 'vspy.DistanceMetric.SUM_OF_SQUARES', 'vspy.DistanceMetric.INNER_PRODUCT', 'vspy.DistanceMetric.COSINE', 'vspy.DistanceMetric.L2'.
     """
     import enum
     import json
@@ -1252,7 +1252,7 @@ def init_centroids(
         config: Optional[Mapping[str, Any]] = None,
         verbose: bool = False,
         trace_id: Optional[str] = None,
-        distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.L2,
+        distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     ) -> np.ndarray:
         logger = setup(config, verbose)
         logger.debug(
@@ -1682,7 +1682,7 @@ def ingest_vectors_udf(
         config: Optional[Mapping[str, Any]] = None,
         verbose: bool = False,
         trace_id: Optional[str] = None,
-        distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.L2,
+        distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     ):
         import os
         import random

apis/python/src/tiledb/vector_search/ivf_flat_index.py

@@ -350,7 +350,7 @@ def _taskgraph_query(
         num_partitions: int = -1,
         num_workers: int = -1,
         config: Optional[Mapping[str, Any]] = None,
-        distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.L2,
+        distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     ):
         """
         Query an IVF_FLAT index using TileDB cloud taskgraphs
@@ -536,7 +536,7 @@ def create(
     group_exists: bool = False,
     config: Optional[Mapping[str, Any]] = None,
     storage_version: str = STORAGE_VERSION,
-    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.L2,
+    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     **kwargs,
 ) -> IVFFlatIndex:
     """
@@ -562,7 +562,8 @@ def create(
     """
     validate_storage_version(storage_version)
     if (
-        distance_metric != vspy.DistanceMetric.L2
+        distance_metric != vspy.DistanceMetric.SUM_OF_SQUARES
+        and distance_metric != vspy.DistanceMetric.L2
         and distance_metric != vspy.DistanceMetric.COSINE
     ):
         raise ValueError(

apis/python/src/tiledb/vector_search/ivf_pq_index.py

@@ -138,7 +138,7 @@ def create(
     config: Optional[Mapping[str, Any]] = None,
     storage_version: str = STORAGE_VERSION,
     partitions: Optional[int] = None,
-    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.L2,
+    distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     **kwargs,
 ) -> IVFPQIndex:
     """
@@ -181,7 +181,10 @@ def create(
         raise ValueError(
             f"Number of dimensions ({dimensions}) must be divisible by num_subspaces ({num_subspaces})."
         )
-    if distance_metric != vspy.DistanceMetric.L2:
+    if (
+        distance_metric != vspy.DistanceMetric.SUM_OF_SQUARES
+        and distance_metric != vspy.DistanceMetric.L2
+    ):
         raise ValueError(
             f"Distance metric {distance_metric} is not supported in IVF_PQ"
         )

apis/python/src/tiledb/vector_search/module.cc

@@ -167,13 +167,14 @@ static void declare_qv_query_heap_infinite_ram(
          size_t nprobe,
          size_t k_nn,
          size_t nthreads,
-         DistanceMetric distance_metric = DistanceMetric::L2) -> py::tuple {
+         DistanceMetric distance_metric =
+             DistanceMetric::SUM_OF_SQUARES) -> py::tuple {
         auto mat = ColMajorPartitionedMatrixWrapper<T, Id_Type, Id_Type>(
             parts, ids, indices);
 
         auto top_centroids = detail::ivf::ivf_top_centroids(
             centroids, query_vectors, nprobe, nthreads);
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::ivf::qv_query_heap_infinite_ram(
               top_centroids,
               mat,
@@ -203,6 +204,16 @@ static void declare_qv_query_heap_infinite_ram(
               nthreads,
               cosine_distance_normalized{});
           return make_python_pair(std::move(r));
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::ivf::qv_query_heap_infinite_ram(
+              top_centroids,
+              mat,
+              query_vectors,
+              nprobe,
+              k_nn,
+              nthreads,
+              sqrt_sum_of_squares_distance{});
+          return make_python_pair(std::move(r));
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -227,11 +238,11 @@ static void declare_qv_query_heap_finite_ram(
          size_t upper_bound,
          size_t nthreads,
          uint64_t timestamp,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> py::tuple {  // std::tuple<ColMajorMatrix<float>,
                           // ColMajorMatrix<size_t>> { //
                           // TODO change return type
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::ivf::qv_query_heap_finite_ram<T, Id_Type>(
               ctx,
               parts_uri,
@@ -276,6 +287,21 @@ static void declare_qv_query_heap_finite_ram(
               timestamp,
               cosine_distance_normalized{});
           return make_python_pair(std::move(r));
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::ivf::qv_query_heap_finite_ram<T, Id_Type>(
+              ctx,
+              parts_uri,
+              centroids,
+              query_vectors,
+              indices,
+              ids_uri,
+              nprobe,
+              k_nn,
+              upper_bound,
+              nthreads,
+              timestamp,
+              sqrt_sum_of_squares_distance{});
+          return make_python_pair(std::move(r));
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -296,7 +322,7 @@ static void declare_nuv_query_heap_infinite_ram(
          size_t nprobe,
          size_t k_nn,
          size_t nthreads,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> std::tuple<
               ColMajorMatrix<float>,
               ColMajorMatrix<uint64_t>> {  // TODO change return type
@@ -307,7 +333,7 @@ static void declare_nuv_query_heap_infinite_ram(
             detail::ivf::partition_ivf_flat_index<Id_Type>(
                 centroids, query_vectors, nprobe, nthreads);
 
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::ivf::nuv_query_heap_infinite_ram(
               mat,
               active_partitions,
@@ -337,6 +363,16 @@ static void declare_nuv_query_heap_infinite_ram(
               nthreads,
               cosine_distance_normalized{});
           return r;
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::ivf::nuv_query_heap_infinite_ram(
+              mat,
+              active_partitions,
+              query_vectors,
+              active_queries,
+              k_nn,
+              nthreads,
+              sqrt_sum_of_squares_distance{});
+          return r;
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -360,7 +396,7 @@ static void declare_nuv_query_heap_finite_ram(
          size_t upper_bound,
          size_t nthreads,
          uint64_t timestamp,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> std::tuple<
               ColMajorMatrix<float>,
               ColMajorMatrix<uint64_t>> {  // TODO change return type
@@ -381,7 +417,7 @@ static void declare_nuv_query_heap_finite_ram(
             upper_bound,
             temporal_policy);
 
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::ivf::nuv_query_heap_finite_ram_reg_blocked(
               mat,
               query_vectors,
@@ -411,6 +447,16 @@ static void declare_nuv_query_heap_finite_ram(
               nthreads,
               cosine_distance_normalized{});
           return r;
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::ivf::nuv_query_heap_finite_ram_reg_blocked(
+              mat,
+              query_vectors,
+              active_queries,
+              k_nn,
+              upper_bound,
+              nthreads,
+              sqrt_sum_of_squares_distance{});
+          return r;
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -702,9 +748,9 @@ static void declare_vq_query_heap(py::module& m, const std::string& suffix) {
          const std::vector<uint64_t>& ids,
          int k,
          size_t nthreads,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, ids, k, nthreads, sum_of_squares_distance{});
           return r;
@@ -716,6 +762,15 @@ static void declare_vq_query_heap(py::module& m, const std::string& suffix) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, ids, k, nthreads, cosine_distance{});
           return r;
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::flat::vq_query_heap(
+              data,
+              query_vectors,
+              ids,
+              k,
+              nthreads,
+              sqrt_sum_of_squares_distance{});
+          return r;
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -732,9 +787,9 @@ static void declare_vq_query_heap_pyarray(
          const std::vector<uint64_t>& ids,
          int k,
          size_t nthreads,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, ids, k, nthreads, sum_of_squares_distance{});
           return r;
@@ -746,6 +801,15 @@ static void declare_vq_query_heap_pyarray(
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, ids, k, nthreads, cosine_distance{});
           return r;
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::flat::vq_query_heap(
+              data,
+              query_vectors,
+              ids,
+              k,
+              nthreads,
+              sqrt_sum_of_squares_distance{});
+          return r;
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -869,9 +933,9 @@ PYBIND11_MODULE(_tiledbvspy, m) {
          ColMajorMatrix<float>& query_vectors,
          int k,
          size_t nthreads,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, k, nthreads, sum_of_squares_distance{});
           return r;
@@ -883,6 +947,10 @@ PYBIND11_MODULE(_tiledbvspy, m) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, k, nthreads, cosine_distance{});
           return r;
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::flat::vq_query_heap(
+              data, query_vectors, k, nthreads, sqrt_sum_of_squares_distance{});
+          return r;
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -894,9 +962,9 @@ PYBIND11_MODULE(_tiledbvspy, m) {
          ColMajorMatrix<float>& query_vectors,
          int k,
          size_t nthreads,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, k, nthreads, sum_of_squares_distance{});
           return r;
@@ -908,6 +976,10 @@ PYBIND11_MODULE(_tiledbvspy, m) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, k, nthreads, cosine_distance{});
           return r;
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::flat::vq_query_heap(
+              data, query_vectors, k, nthreads, sqrt_sum_of_squares_distance{});
+          return r;
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -919,9 +991,9 @@ PYBIND11_MODULE(_tiledbvspy, m) {
          ColMajorMatrix<float>& query_vectors,
          int k,
          size_t nthreads,
-         DistanceMetric distance_metric = DistanceMetric::L2)
+         DistanceMetric distance_metric = DistanceMetric::SUM_OF_SQUARES)
           -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
-        if (distance_metric == DistanceMetric::L2) {
+        if (distance_metric == DistanceMetric::SUM_OF_SQUARES) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, k, nthreads, sum_of_squares_distance{});
           return r;
@@ -933,6 +1005,10 @@ PYBIND11_MODULE(_tiledbvspy, m) {
           auto r = detail::flat::vq_query_heap(
               data, query_vectors, k, nthreads, cosine_distance{});
           return r;
+        } else if (distance_metric == DistanceMetric::L2) {
+          auto r = detail::flat::vq_query_heap(
+              data, query_vectors, k, nthreads, sqrt_sum_of_squares_distance{});
+          return r;
         } else {
           throw std::runtime_error("Invalid distance metric");
         }
@@ -1014,9 +1090,10 @@ PYBIND11_MODULE(_tiledbvspy, m) {
   declare_debug_matrix<uint64_t>(m, "_u64");
 
   py::enum_<DistanceMetric>(m, "DistanceMetric")
-      .value("L2", DistanceMetric::L2)
+      .value("SUM_OF_SQUARES", DistanceMetric::SUM_OF_SQUARES)
       .value("INNER_PRODUCT", DistanceMetric::INNER_PRODUCT)
       .value("COSINE", DistanceMetric::COSINE)
+      .value("L2", DistanceMetric::L2)
       .export_values();
 
   /* === Module inits === */