Add support for int8 type indexes (#342)

This adds support for `int8` type vectors. 

`int8` scalar quantization is the most popular configuration for scalar quantization of signed floating point values. 
https://huggingface.co/blog/embedding-quantization

Author: NikolaosPapailiou

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

@@ -25,16 +25,19 @@ bool enable_stats = false;
 std::vector<json> core_stats;
 
 PYBIND11_MAKE_OPAQUE(std::vector<uint8_t>);
+PYBIND11_MAKE_OPAQUE(std::vector<int8_t>);
 PYBIND11_MAKE_OPAQUE(std::vector<uint32_t>);
 PYBIND11_MAKE_OPAQUE(std::vector<uint64_t>);
 PYBIND11_MAKE_OPAQUE(std::vector<float>);
 PYBIND11_MAKE_OPAQUE(std::vector<double>);
 PYBIND11_MAKE_OPAQUE(std::list<std::vector<uint8_t>>);
+PYBIND11_MAKE_OPAQUE(std::list<std::vector<int8_t>>);
 PYBIND11_MAKE_OPAQUE(std::list<std::vector<uint32_t>>);
 PYBIND11_MAKE_OPAQUE(std::list<std::vector<uint64_t>>);
 PYBIND11_MAKE_OPAQUE(std::list<std::vector<float>>);
 PYBIND11_MAKE_OPAQUE(std::list<std::vector<double>>);
 PYBIND11_MAKE_OPAQUE(std::vector<std::list<uint8_t>>);
+PYBIND11_MAKE_OPAQUE(std::vector<std::list<int8_t>>);
 PYBIND11_MAKE_OPAQUE(std::vector<std::list<uint32_t>>);
 PYBIND11_MAKE_OPAQUE(std::vector<std::list<uint64_t>>);
 PYBIND11_MAKE_OPAQUE(std::vector<std::list<float>>);
@@ -627,6 +630,7 @@ PYBIND11_MODULE(_tiledbvspy, m) {
   declareStdVector<float>(m, "f32");
   declareStdVector<double>(m, "f64");
   declareStdVector<uint8_t>(m, "u8");
+  declareStdVector<int8_t>(m, "i8");
   declareStdVector<uint32_t>(m, "u32");
   declareStdVector<uint64_t>(m, "u64");
   if constexpr (!std::is_same_v<uint64_t, size_t>) {
@@ -664,6 +668,7 @@ PYBIND11_MODULE(_tiledbvspy, m) {
   /* === Matrix === */
 
   declareColMajorMatrix<uint8_t>(m, "_u8");
+  declareColMajorMatrix<int8_t>(m, "_i8");
   declareColMajorMatrix<float>(m, "_f32");
   declareColMajorMatrix<double>(m, "_f64");
   declareColMajorMatrix<int32_t>(m, "_i32");
@@ -677,6 +682,8 @@ PYBIND11_MODULE(_tiledbvspy, m) {
 
   declareColMajorMatrixSubclass<tdbColMajorMatrix<uint8_t>>(
       m, "tdbColMajorMatrix", "_u8");
+  declareColMajorMatrixSubclass<tdbColMajorMatrix<int8_t>>(
+      m, "tdbColMajorMatrix", "_i8");
   declareColMajorMatrixSubclass<tdbColMajorMatrix<uint64_t>>(
       m, "tdbColMajorMatrix", "_u64");
   declareColMajorMatrixSubclass<tdbColMajorMatrix<float>>(
@@ -688,6 +695,7 @@ PYBIND11_MODULE(_tiledbvspy, m) {
 
   // Converters from pyarray to matrix
   declare_pyarray_to_matrix<uint8_t>(m, "_u8");
+  declare_pyarray_to_matrix<int8_t>(m, "_i8");
   declare_pyarray_to_matrix<uint64_t>(m, "_u64");
   declare_pyarray_to_matrix<float>(m, "_f32");
   declare_pyarray_to_matrix<double>(m, "_f64");
@@ -716,6 +724,17 @@ PYBIND11_MODULE(_tiledbvspy, m) {
         return r;
       });
 
+  m.def(
+      "query_vq_i8",
+      [](tdbColMajorMatrix<int8_t>& data,
+         ColMajorMatrix<float>& query_vectors,
+         int k,
+         size_t nthreads)
+          -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
+        auto r = detail::flat::vq_query_heap(data, query_vectors, k, nthreads);
+        return r;
+      });
+
   m.def(
       "validate_top_k_u64",
       [](const ColMajorMatrix<uint64_t>& top_k,
@@ -724,33 +743,45 @@ PYBIND11_MODULE(_tiledbvspy, m) {
       });
 
   declare_vq_query_heap<uint8_t>(m, "u8");
+  declare_vq_query_heap<int8_t>(m, "i8");
   declare_vq_query_heap<float>(m, "f32");
   declare_vq_query_heap_pyarray<uint8_t>(m, "u8");
+  declare_vq_query_heap_pyarray<int8_t>(m, "i8");
   declare_vq_query_heap_pyarray<float>(m, "f32");
 
   declare_qv_query_heap_infinite_ram<uint8_t>(m, "u8");
+  declare_qv_query_heap_infinite_ram<int8_t>(m, "i8");
   declare_qv_query_heap_infinite_ram<float>(m, "f32");
   declare_qv_query_heap_finite_ram<uint8_t>(m, "u8");
+  declare_qv_query_heap_finite_ram<int8_t>(m, "i8");
   declare_qv_query_heap_finite_ram<float>(m, "f32");
   declare_nuv_query_heap_infinite_ram<uint8_t>(m, "u8");
+  declare_nuv_query_heap_infinite_ram<int8_t>(m, "i8");
   declare_nuv_query_heap_infinite_ram<float>(m, "f32");
   declare_nuv_query_heap_finite_ram<uint8_t>(m, "u8");
+  declare_nuv_query_heap_finite_ram<int8_t>(m, "i8");
   declare_nuv_query_heap_finite_ram<float>(m, "f32");
 
   declare_ivf_index<uint8_t>(m, "u8");
+  declare_ivf_index<int8_t>(m, "i8");
   declare_ivf_index<float>(m, "f32");
   declare_ivf_index_tdb<uint8_t>(m, "u8");
+  declare_ivf_index_tdb<int8_t>(m, "i8");
   declare_ivf_index_tdb<float>(m, "f32");
 
   declarePartitionIvfIndex<uint8_t>(m, "u8");
+  declarePartitionIvfIndex<int8_t>(m, "i8");
   declarePartitionIvfIndex<float>(m, "f32");
 
   declarePartitionedMatrix<uint8_t, uint64_t, uint64_t, uint64_t>(
       m, "tdbPartitionedMatrix", "u8");
+  declarePartitionedMatrix<int8_t, uint64_t, uint64_t, uint64_t>(
+      m, "tdbPartitionedMatrix", "i8");
   declarePartitionedMatrix<float, uint64_t, uint64_t, uint64_t>(
       m, "tdbPartitionedMatrix", "f32");
 
   declare_dist_qv<uint8_t>(m, "u8");
+  declare_dist_qv<int8_t>(m, "i8");
   declare_dist_qv<float>(m, "f32");
   declareFixedMinPairHeap(m);
 
@@ -770,6 +801,7 @@ PYBIND11_MODULE(_tiledbvspy, m) {
   m.def("stats_dump", []() { return json{core_stats}.dump(); });
 
   declare_debug_slice<uint8_t>(m, "_u8");
+  declare_debug_slice<int8_t>(m, "_i8");
   declare_debug_slice<float>(m, "_f32");
   declare_debug_slice<uint64_t>(m, "_u64");
 

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

@@ -48,6 +48,8 @@ def load_as_matrix(
         m = tdbColMajorMatrix_i64(ctx, path, 0, None, 0, size, 0, timestamp)
     elif dtype == np.uint8:
         m = tdbColMajorMatrix_u8(ctx, path, 0, None, 0, size, 0, timestamp)
+    elif dtype == np.int8:
+        m = tdbColMajorMatrix_i8(ctx, path, 0, None, 0, size, 0, timestamp)
     # elif dtype == np.uint64:
     #     return tdbColMajorMatrix_u64(ctx, path, size, timestamp)
     else:
@@ -91,6 +93,8 @@ def debug_slice(m: "colMajorMatrix", name: str):
         return debug_slice_f32(m, name)
     elif dtype == np.uint8:
         return debug_slice_u8(m, name)
+    elif dtype == np.int8:
+        return debug_slice_i8(m, name)
     elif dtype == np.uint64:
         return debug_slice_u64(m, name)
     else:
@@ -112,6 +116,8 @@ def query_vq_nth(db: "colMajorMatrix", *args):
         return query_vq_f32(db, *args)
     elif db.dtype == np.uint8:
         return query_vq_u8(db, *args)
+    elif db.dtype == np.int8:
+        return query_vq_i8(db, *args)
     else:
         raise TypeError("Unknown type!")
 
@@ -131,6 +137,8 @@ def query_vq_heap(db: "colMajorMatrix", *args):
         return vq_query_heap_f32(db, *args)
     elif db.dtype == np.uint8:
         return vq_query_heap_u8(db, *args)
+    elif db.dtype == np.int8:
+        return vq_query_heap_i8(db, *args)
     else:
         raise TypeError("Unknown type!")
 
@@ -150,6 +158,8 @@ def query_vq_heap_pyarray(db: "colMajorMatrix", *args):
         return vq_query_heap_pyarray_f32(db, *args)
     elif db.dtype == np.uint8:
         return vq_query_heap_pyarray_u8(db, *args)
+    elif db.dtype == np.int8:
+        return vq_query_heap_pyarray_i8(db, *args)
     else:
         raise TypeError("Unknown type!")
 
@@ -195,6 +205,8 @@ def ivf_index_tdb(
         return ivf_index_tdb_f32(*args)
     elif dtype == np.uint8:
         return ivf_index_tdb_u8(*args)
+    elif dtype == np.int8:
+        return ivf_index_tdb_i8(*args)
     else:
         raise TypeError("Unknown type!")
 
@@ -240,6 +252,8 @@ def ivf_index(
         return ivf_index_f32(*args)
     elif dtype == np.uint8:
         return ivf_index_u8(*args)
+    elif dtype == np.int8:
+        return ivf_index_i8(*args)
     else:
         raise TypeError("Unknown type!")
 
@@ -312,6 +326,11 @@ def ivf_query_ram(
             return nuv_query_heap_infinite_ram_reg_blocked_u8(*args)
         else:
             return qv_query_heap_infinite_ram_u8(*args)
+    elif dtype == np.int8:
+        if use_nuv_implementation:
+            return nuv_query_heap_infinite_ram_reg_blocked_i8(*args)
+        else:
+            return qv_query_heap_infinite_ram_i8(*args)
     else:
         raise TypeError("Unknown type!")
 
@@ -394,6 +413,11 @@ def ivf_query(
             return nuv_query_heap_finite_ram_reg_blocked_u8(*args)
         else:
             return qv_query_heap_finite_ram_u8(*args)
+    elif dtype == np.int8:
+        if use_nuv_implementation:
+            return nuv_query_heap_finite_ram_reg_blocked_i8(*args)
+        else:
+            return qv_query_heap_finite_ram_i8(*args)
     else:
         raise TypeError("Unknown type!")
 
@@ -403,6 +427,8 @@ def partition_ivf_index(centroids, query, nprobe=1, nthreads=0):
         return partition_ivf_index_f32(centroids, query, nprobe, nthreads)
     elif query.dtype == np.uint8:
         return partition_ivf_index_u8(centroids, query, nprobe, nthreads)
+    elif query.dtype == np.int8:
+        return partition_ivf_index_i8(centroids, query, nprobe, nthreads)
     else:
         raise TypeError("Unsupported type!")
 
@@ -442,6 +468,8 @@ def dist_qv(
         return dist_qv_f32(*args)
     elif dtype == np.uint8:
         return dist_qv_u8(*args)
+    elif dtype == np.int8:
+        return dist_qv_i8(*args)
     else:
         raise TypeError("Unsupported type!")
 
@@ -464,6 +492,8 @@ def array_to_matrix(array: np.ndarray):
         return pyarray_copyto_matrix_i32(array)
     elif array.dtype == np.uint64:
         return pyarray_copyto_matrix_u64(array)
+    elif array.dtype == np.int8:
+        return pyarray_copyto_matrix_i8(array)
     else:
         raise TypeError("Unsupported type!")
 

apis/python/test/common.py

@@ -351,3 +351,15 @@ def move_local_index_to_new_location(index_uri):
     shutil.copytree(index_uri, copied_index_uri)
     shutil.rmtree(index_uri)
     return copied_index_uri
+
+
+def quantize_embeddings_int8(
+    embeddings: np.ndarray,
+) -> np.ndarray:
+    """
+    Quantizes embeddings to a lower precision.
+    """
+    ranges = np.vstack((np.min(embeddings, axis=0), np.max(embeddings, axis=0)))
+    starts = ranges[0, :]
+    steps = (ranges[1, :] - ranges[0, :]) / 255
+    return ((embeddings - starts) / steps - 128).astype(np.int8)

apis/python/test/test_ingestion.py

@@ -337,6 +337,54 @@ def test_ingestion_numpy(tmp_path):
         assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
 
+def test_ingestion_numpy_i8(tmp_path):
+    source_uri = siftsmall_inputs_file
+    queries_uri = siftsmall_query_file
+    gt_uri = siftsmall_groundtruth_file
+    index_uri = os.path.join(tmp_path, "array")
+    k = 100
+    partitions = 100
+    nqueries = 100
+    nprobe = 20
+
+    input_vectors = quantize_embeddings_int8(load_fvecs(source_uri))
+
+    queries = quantize_embeddings_int8(load_fvecs(queries_uri)).astype(np.float32)
+    gt_i, gt_d = get_groundtruth_ivec(gt_uri, k=k, nqueries=nqueries)
+
+    for index_type, index_class in zip(INDEXES, INDEX_CLASSES):
+        # TODO(paris): Fix Vamana bug and re-enable:
+        # RuntimeError: IndexError: index 100 is out of bounds for axis 0 with size 100
+        if index_type == "VAMANA":
+            continue
+
+        index_uri = os.path.join(tmp_path, f"array_{index_type}")
+        index = ingest(
+            index_type=index_type,
+            index_uri=index_uri,
+            input_vectors=input_vectors,
+            partitions=partitions,
+        )
+        _, result = index.query(queries, k=k, nprobe=nprobe)
+        assert accuracy(result, gt_i) > MINIMUM_ACCURACY
+
+        index_uri = move_local_index_to_new_location(index_uri)
+        index_ram = index_class(uri=index_uri)
+        _, result = index_ram.query(queries, k=k, nprobe=nprobe)
+        assert accuracy(result, gt_i) > MINIMUM_ACCURACY
+
+        _, result = index_ram.query(
+            queries,
+            k=k,
+            nprobe=nprobe,
+            use_nuv_implementation=True,
+        )
+        assert accuracy(result, gt_i) > MINIMUM_ACCURACY
+
+        _, result = index_ram.query(queries, k=k, nprobe=nprobe, mode=Mode.LOCAL)
+        assert accuracy(result, gt_i) > MINIMUM_ACCURACY
+
+
 def test_ingestion_multiple_workers(tmp_path):
     source_uri = siftsmall_inputs_file
     queries_uri = siftsmall_query_file

src/include/detail/scoring/l2_distance.h

@@ -135,7 +135,8 @@ inline float naive_sum_of_squares(const V& a, const W& b) {
  */
 template <feature_vector V>
   requires std::same_as<typename V::value_type, float> ||
-           std::same_as<typename V::value_type, uint8_t>
+           std::same_as<typename V::value_type, uint8_t> ||
+           std::same_as<typename V::value_type, int8_t>
 inline float unroll4_sum_of_squares(const V& a) {
   size_t size_a = size(a);
   size_t stop = 4 * (size_a / 4);
@@ -187,7 +188,8 @@ inline float unroll4_sum_of_squares(const V& a, const W& b) {
  */
 template <feature_vector V, feature_vector W>
   requires std::same_as<typename V::value_type, float> &&
-           std::same_as<typename W::value_type, uint8_t>
+           (std::same_as<typename W::value_type, uint8_t> ||
+            std::same_as<typename W::value_type, int8_t>)
 inline float unroll4_sum_of_squares(const V& a, const W& b) {
   size_t size_a = size(a);
   size_t stop = 4 * (size_a / 4);
@@ -212,8 +214,9 @@ inline float unroll4_sum_of_squares(const V& a, const W& b) {
  * Unrolled l2 distance between vector of uint8_t and vector of float
  */
 template <feature_vector V, feature_vector W>
-  requires std::same_as<typename V::value_type, uint8_t> &&
-           std::same_as<typename W::value_type, float>
+  requires(std::same_as<typename V::value_type, uint8_t> ||
+           std::same_as<typename V::value_type, int8_t>) &&
+          std::same_as<typename W::value_type, float>
 inline float unroll4_sum_of_squares(const V& a, const W& b) {
   size_t size_a = size(a);
   size_t stop = 4 * (size_a / 4);
@@ -238,8 +241,10 @@ inline float unroll4_sum_of_squares(const V& a, const W& b) {
  * Unrolled l2 distance between vector of uint8_t and vector of uint8_t
  */
 template <feature_vector V, feature_vector W>
-  requires std::same_as<typename V::value_type, uint8_t> &&
-           std::same_as<typename W::value_type, uint8_t>
+  requires(std::same_as<typename V::value_type, uint8_t> ||
+           std::same_as<typename V::value_type, int8_t>) &&
+          (std::same_as<typename W::value_type, uint8_t> ||
+           std::same_as<typename W::value_type, int8_t>)
 inline float unroll4_sum_of_squares(const V& a, const W& b) {
   size_t size_a = size(a);
   size_t stop = 4 * (size_a / 4);