Add FVEC, IVEC and BVEC inputs

Author: Nikos Papailiou

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

@@ -70,6 +70,7 @@ def ingest(
     import math
     from typing import Any, Mapping, Optional
     import multiprocessing
+    import os
 
     import numpy as np
 
@@ -151,6 +152,33 @@ def read_source_metadata(
                 size = int.from_bytes(f.read(4), "little")
                 dimensions = int.from_bytes(f.read(4), "little")
                 return size, dimensions, np.float32
+        elif source_type == "FVEC":
+            vfs = tiledb.VFS()
+            with vfs.open(source_uri, "rb") as f:
+                dimensions = int.from_bytes(f.read(4), "little")
+                vector_size = 4 + dimensions * 4
+                f.seek(0, os.SEEK_END)
+                file_size = f.tell()
+                size = int(file_size / vector_size)
+                return size, dimensions, np.float32
+        elif source_type == "IVEC":
+            vfs = tiledb.VFS()
+            with vfs.open(source_uri, "rb") as f:
+                dimensions = int.from_bytes(f.read(4), "little")
+                vector_size = 4 + dimensions * 4
+                f.seek(0, os.SEEK_END)
+                file_size = f.tell()
+                size = int(file_size / vector_size)
+                return size, dimensions, np.int32
+        elif source_type == "BVEC":
+            vfs = tiledb.VFS()
+            with vfs.open(source_uri, "rb") as f:
+                dimensions = int.from_bytes(f.read(4), "little")
+                vector_size = 4 + dimensions
+                f.seek(0, os.SEEK_END)
+                file_size = f.tell()
+                size = int(file_size / vector_size)
+                return size, dimensions, np.uint8
         else:
             raise ValueError(f"Not supported source_type {source_type}")
 
@@ -419,6 +447,40 @@ def read_input_vectors(
                     ).astype(vector_type),
                     (read_size, dimensions),
                 )
+        elif source_type == "FVEC" or source_type == "IVEC":
+            vfs = tiledb.VFS()
+            vector_values = 1 + dimensions
+            vector_size = vector_values * 4
+            read_size = end_pos - start_pos
+            read_offset = start_pos * vector_size
+            with vfs.open(source_uri, "rb") as f:
+                f.seek(read_offset)
+                return np.delete(
+                    np.reshape(
+                        np.frombuffer(
+                            f.read(read_size * vector_size),
+                            count=read_size * vector_values,
+                            dtype=vector_type,
+                        ).astype(vector_type),
+                        (read_size, dimensions + 1),
+                    ), 0, axis=1)
+        elif source_type == "BVEC":
+            vfs = tiledb.VFS()
+            vector_values = 1 + dimensions
+            vector_size = vector_values * 1
+            read_size = end_pos - start_pos
+            read_offset = start_pos * vector_size
+            with vfs.open(source_uri, "rb") as f:
+                f.seek(read_offset)
+                return np.delete(
+                    np.reshape(
+                        np.frombuffer(
+                            f.read(read_size * vector_size),
+                            count=read_size * vector_values,
+                            dtype=vector_type,
+                        ).astype(vector_type),
+                        (read_size, dimensions + 1),
+                    ), 0, axis=1)
 
     # --------------------------------------------------------------------
     # UDFs
@@ -1182,6 +1244,7 @@ def consolidate_and_vacuum(
             else:
                 raise err
         group = tiledb.Group(array_uri, "w")
+        group.meta["dataset_type"] = "vector_search"
 
         in_size, dimensions, vector_type = read_source_metadata(
             source_uri=source_uri, source_type=source_type, logger=logger

apis/python/test/common.py

@@ -11,6 +11,44 @@ def xbin_mmap(fname, dtype):
     return np.memmap(fname, dtype=dtype, mode="r", offset=8, shape=(n, d))
 
 
+def get_queries_fvec(file, dimensions, nqueries=None):
+    vfs = tiledb.VFS()
+    vector_values = 1 + dimensions
+    vector_size = vector_values * 4
+    read_size = nqueries
+    read_offset = 0
+    with vfs.open(file, "rb") as f:
+        f.seek(read_offset)
+        return np.delete(
+            np.reshape(
+                np.frombuffer(
+                    f.read(read_size * vector_size),
+                    count=read_size * vector_values,
+                    dtype=np.float32,
+                ).astype(np.float32),
+                (read_size, dimensions + 1),
+            ), 0, axis=1)
+
+
+def get_groundtruth_ivec(file, k=None, nqueries=None):
+    vfs = tiledb.VFS()
+    vector_values = 1 + k
+    vector_size = vector_values * 4
+    read_size = nqueries
+    read_offset = 0
+    with vfs.open(file, "rb") as f:
+        f.seek(read_offset)
+        return np.delete(
+            np.reshape(
+                np.frombuffer(
+                    f.read(read_size * vector_size),
+                    count=read_size * vector_values,
+                    dtype=np.int32,
+                ).astype(np.int32),
+                (read_size, k + 1),
+            ), 0, axis=1), None
+
+
 def get_queries(dataset_dir, dtype, nqueries=None):
     fname = os.path.join(dataset_dir, "queries")
     x = xbin_mmap(fname, dtype=dtype)

apis/python/test/data/siftsmall/siftsmall_base.fvecs

@@ -93,3 +93,28 @@ def test_ivf_flat_ingestion_f32(tmp_path):
     )
     result = np.transpose(index.query(np.transpose(query_vectors), k=k, nprobe=partitions))
     assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+
+
+def test_ivf_flat_ingestion_fvec(tmp_path):
+    source_uri = "test/data/siftsmall/siftsmall_base.fvecs"
+    queries_uri = "test/data/siftsmall/siftsmall_query.fvecs"
+    gt_uri = "test/data/siftsmall/siftsmall_groundtruth.ivecs"
+    source_type = "FVEC"
+    array_uri = os.path.join(tmp_path, "array")
+    k = 100
+    dimensions = 128
+    partitions = 1000
+    nqueries = 100
+
+    query_vectors = get_queries_fvec(queries_uri, dimensions=dimensions, nqueries=nqueries)
+    gt_i, gt_d = get_groundtruth_ivec(gt_uri, k=k, nqueries=nqueries)
+
+    index = ingest(
+        index_type="IVF_FLAT",
+        array_uri=array_uri,
+        source_uri=source_uri,
+        source_type=source_type,
+        partitions=partitions,
+    )
+    result = np.transpose(index.query(np.transpose(query_vectors), k=k, nprobe=partitions))
+    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))

apis/python/test/data/siftsmall/siftsmall_groundtruth.ivecs

@@ -111,7 +111,7 @@ void write_matrix(
                    TILEDB_COL_MAJOR;
   query.set_layout(order)
       .set_data_buffer(
-          "values", &A(0, 0), (int)A.num_rows() * (int)A.num_cols())
+          "values", &A(0, 0), (uint64_t)A.num_rows() * (uint64_t)A.num_cols())
       .set_subarray(subarray);
   query.submit();