Merge pull request #67 from TileDB-Inc/npapa/cpp_ingestion

Use native vector ingestion in Python taskgraphs

Author: NikolaosPapailiou

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

@@ -1,16 +1,19 @@
-from .index import FlatIndex
+from .index import FlatIndex, IVFFlatIndex
 from .ingestion import ingest
 from .module import load_as_array
 from .module import load_as_matrix
-from .module import query_vq, query_kmeans, validate_top_k, array_to_matrix
+from .module import query_vq, query_kmeans, validate_top_k, array_to_matrix, ivf_index, ivf_index_tdb
 
 __all__ = [
     "FlatIndex",
+    "IVFFlatIndex",
     "load_as_array",
     "load_as_matrix",
     "ingest",
     "query_vq",
     "query_kmeans",
     "validate_top_k",
+    "ivf_index",
+    "ivf_index_tdb",
     "array_to_matrix",
 ]

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

@@ -1,7 +1,6 @@
 import os
 
 import numpy as np
-
 from tiledb.vector_search.module import *
 
 

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

@@ -124,6 +124,56 @@ static void declare_kmeans_query(py::module& m, const std::string& suffix) {
         }, py::keep_alive<1,2>());
 }
 
+template <typename T>
+static void declare_ivf_index(py::module& m, const std::string& suffix) {
+  m.def(("ivf_index_" + suffix).c_str(),
+      [](tiledb::Context& ctx,
+        const ColMajorMatrix<T>& db,
+        const std::string& centroids_uri,
+        const std::string& parts_uri,
+        const std::string& index_uri,
+        const std::string& id_uri,
+        size_t start_pos,
+        size_t end_pos,
+        size_t nthreads) -> int {
+            return detail::ivf::ivf_index<T, uint64_t, float>(
+                ctx,
+                db,
+                centroids_uri,
+                parts_uri,
+                index_uri,
+                id_uri,
+                start_pos,
+                end_pos,
+                nthreads);
+        }, py::keep_alive<1,2>());
+}
+
+template <typename T>
+static void declare_ivf_index_tdb(py::module& m, const std::string& suffix) {
+  m.def(("ivf_index_tdb_" + suffix).c_str(),
+      [](tiledb::Context& ctx,
+        const std::string& db_uri,
+        const std::string& centroids_uri,
+        const std::string& parts_uri,
+        const std::string& index_uri,
+        const std::string& id_uri,
+        size_t start_pos,
+        size_t end_pos,
+        size_t nthreads) -> int {
+            return detail::ivf::ivf_index<T, uint64_t, float>(
+                ctx,
+                db_uri,
+                centroids_uri,
+                parts_uri,
+                index_uri,
+                id_uri,
+                start_pos,
+                end_pos,
+                nthreads);
+        }, py::keep_alive<1,2>());
+}
+
 
 // Declarations for typed subclasses of ColMajorMatrix
 template <typename P>
@@ -273,4 +323,9 @@ PYBIND11_MODULE(_tiledbvspy, m) {
   declare_kmeans_query<uint8_t>(m, "u8");
   declare_kmeans_query<float>(m, "f32");
 
+  declare_ivf_index<uint8_t>(m, "u8");
+  declare_ivf_index<float>(m, "f32");
+  declare_ivf_index_tdb<uint8_t>(m, "u8");
+  declare_ivf_index_tdb<float>(m, "f32");
+
 }

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

@@ -4,7 +4,6 @@
 
 import tiledb
 from tiledb.vector_search._tiledbvspy import *
-from tiledb.vector_search import _tiledbvspy as cc
 
 
 def load_as_matrix(path: str, nqueries: int = 0, config: Dict = {}):
@@ -82,6 +81,84 @@ def query_vq(db: "colMajorMatrix", *args):
         raise TypeError("Unknown type!")
 
 
+def ivf_index_tdb(
+    dtype: np.dtype,
+    db_uri: str,
+    centroids_uri: str,
+    parts_uri: str,
+    index_uri: str,
+    id_uri: str,
+    start: int = 0,
+    end: int = 0,
+    nthreads: int = 0,
+    config: Dict = None,
+):
+    if config is None:
+        ctx = Ctx({})
+    else:
+        ctx = Ctx(config)
+
+    args = tuple(
+        [
+            ctx,
+            db_uri,
+            centroids_uri,
+            parts_uri,
+            index_uri,
+            id_uri,
+            start,
+            end,
+            nthreads
+        ]
+    )
+
+    if dtype == np.float32:
+        return ivf_index_tdb_f32(*args)
+    elif dtype == np.uint8:
+        return ivf_index_tdb_u8(*args)
+    else:
+        raise TypeError("Unknown type!")
+
+
+def ivf_index(
+    dtype: np.dtype,
+    db: "colMajorMatrix",
+    centroids_uri: str,
+    parts_uri: str,
+    index_uri: str,
+    id_uri: str,
+    start: int = 0,
+    end: int = 0,
+    nthreads: int = 0,
+    config: Dict = None,
+):
+    if config is None:
+        ctx = Ctx({})
+    else:
+        ctx = Ctx(config)
+
+    args = tuple(
+        [
+            ctx,
+            db,
+            centroids_uri,
+            parts_uri,
+            index_uri,
+            id_uri,
+            start,
+            end,
+            nthreads
+        ]
+    )
+
+    if dtype == np.float32:
+        return ivf_index_f32(*args)
+    elif dtype == np.uint8:
+        return ivf_index_u8(*args)
+    else:
+        raise TypeError("Unknown type!")
+
+
 def query_kmeans(
     dtype: np.dtype,
     parts_db: "colMajorMatrix",
@@ -150,7 +227,7 @@ def query_kmeans(
 
 def validate_top_k(results: np.ndarray, ground_truth: np.ndarray):
     if results.dtype == np.uint64:
-        return cc.validate_top_k_u64(results, ground_truth)
+        return validate_top_k_u64(results, ground_truth)
     else:
         raise TypeError("Unknown type for validate_top_k!")
 

apis/python/src/tiledb/vector_search/module.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)