Fix/run Python tests and pip build; add array_to_matrix function (#30)

Author: ihnorton

.github/workflows/ci.yml

@@ -19,6 +19,6 @@ jobs:
       - name: Configure CMake
         run: cmake -S ./src -B ./src/build -DCMAKE_BUILD_TYPE=Debug
       - name: Build
-        run: cmake --build ./src/build
+        run: cmake --build ./src/build -j3
       - name: Run Tests
         run: cmake --build ./src/build --target check-ci

.github/workflows/ci_python.yml

@@ -7,15 +7,14 @@ on:
     branches: [main]
 
 jobs:
-  Run-Tests:
+  run-tests:
     strategy:
       matrix:
         os: [ubuntu-latest]
         python-version: ["3.9"]
     runs-on: ${{ matrix.os }}
     steps:
-      - name: Install OpenBLAS
-        run: sudo apt install libopenblas-dev
+      - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:
@@ -25,10 +24,12 @@ jobs:
           which python
           which pip
           python --version
-      - name: Install pybind11
-        run: pip install "pybind11[global]"
-      - uses: actions/checkout@v3
-      - name: Configure CMake
-        run: cmake -S ./src -B ./src/build -DCMAKE_BUILD_TYPE=Debug -DTILEDB_VS_PYTHON=ON
-      - name: Build
-        run: cmake --build ./src/build
+      - name: Build and test python
+        run: |
+          cd apis/python
+          pip install .[test]
+          pytest -k "not ingest" # TODO: requires token
+          #pip uninstall -y tiledb.vector_search
+          #pip install -e .
+          #pytest -k "not ingest" # TODO: requires token
+        shell: bash -el {0}

apis/python/pyproject.toml

@@ -33,7 +33,8 @@ repository = "https://github.com/TileDB-Inc/tiledb-vector-search"
 [build-system]
 requires = [
     "setuptools>=42", "wheel", "setuptools_scm>=6",
-    "scikit-build>=0.13", "pybind11"
+    "scikit-build>=0.13", "pybind11[global]",
+    "cmake", "ninja"
 ]
 build-backend = "setuptools.build_meta"
 

apis/python/setup.py

@@ -7,6 +7,10 @@ def get_cmake_overrides():
 
     conf = list()
 
+    tiledb_dir = os.environ.get("TILEDB_DIR", None)
+    if tiledb_dir:
+        cmake_args.append(f"-DTileDB_DIR={tiledb_dir}")
+
     key = "CMAKE_OSX_DEPLOYMENT_TARGET"
     val = os.environ.get(key, default=None)
     if val:

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

@@ -2,7 +2,7 @@
 from .ingestion import ingest
 from .module import load_as_array
 from .module import load_as_matrix
-from .module import query_vq
+from .module import query_vq, query_kmeans, validate_top_k, array_to_matrix
 
 __all__ = [
     "FlatIndex",
@@ -11,4 +11,6 @@
     "ingest",
     "query_vq",
     "query_kmeans",
+    "validate_top_k",
+    "array_to_matrix",
 ]

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

@@ -102,6 +102,22 @@ static void declare_kmeans_query(py::module& m, const std::string& suffix) {
         }, py::keep_alive<1,2>());
 }
 
+template <typename T>
+static void declare_pyarray_to_matrix(py::module& m, const std::string& suffix) {
+  m.def(("pyarray_copyto_matrix" + suffix).c_str(),
+      [](py::array_t<T> arr) -> ColMajorMatrix<T> {
+        py::buffer_info info = arr.request();
+        if (info.ndim != 2)
+          throw std::runtime_error("Number of dimensions must be two");
+        if (info.format != py::format_descriptor<T>::format())
+          throw std::runtime_error("Mismatched buffer format!");
+
+        auto data = std::unique_ptr<T[]>{new T[info.shape[0] * info.shape[1]]};
+        auto r = ColMajorMatrix<T>(std::move(data), info.shape[0], info.shape[1]);
+        return r;
+        });
+}
+
 
 // Declarations for typed subclasses of ColMajorMatrix
 template <typename P>
@@ -133,41 +149,48 @@ PYBIND11_MODULE(_tiledbvspy, m) {
     }
   ));
 
-  /* Vector */
+  /* === Vector === */
+
   declareVector<uint32_t>(m, "_u32");
   declareVector<uint64_t>(m, "_u64");
   declareVector<float>(m, "_f32");
   declareVector<double>(m, "_f64");
 
   m.def("read_vector_u32", &read_vector<uint32_t>, "Read a vector from TileDB");
   m.def("read_vector_u64", &read_vector<uint64_t>, "Read a vector from TileDB");
+
+
   /* === Matrix === */
 
   // template specializations
-  //declareTdbMatrix<float>(m, "_f32");
-
   declareColMajorMatrix<uint8_t>(m, "_u8");
   declareColMajorMatrix<float>(m, "_f32");
   declareColMajorMatrix<double>(m, "_f64");
   declareColMajorMatrix<int32_t>(m, "_i32");
   declareColMajorMatrix<int64_t>(m, "_i64");
-//  declareColMajorMatrix<uint64_t>(m, "_u64");
-  declareColMajorMatrix<size_t>(m, "_szt");
+  declareColMajorMatrix<uint32_t>(m, "_u32");
+  declareColMajorMatrix<uint64_t>(m, "_u64");
+  if constexpr (!std::is_same<uint64_t, unsigned long>::value) {
+    // Required for a return type, but these types are equivalent on linux :/
+    declareColMajorMatrix<unsigned long>(m, "_ul");
+  }
 
   declareColMajorMatrixSubclass<tdbColMajorMatrix<uint8_t>>(
       m, "tdbColMajorMatrix", "_u8");
-  declareColMajorMatrixSubclass<tdbColMajorMatrix<size_t>>(
-      m, "tdbColMajorMatrix", "_szt");
+  declareColMajorMatrixSubclass<tdbColMajorMatrix<uint64_t>>(
+      m, "tdbColMajorMatrix", "_u64");
   declareColMajorMatrixSubclass<tdbColMajorMatrix<float>>(
       m, "tdbColMajorMatrix", "_f32");
   declareColMajorMatrixSubclass<tdbColMajorMatrix<int32_t>>(
       m, "tdbColMajorMatrix", "_i32");
   declareColMajorMatrixSubclass<tdbColMajorMatrix<int64_t>>(
       m, "tdbColMajorMatrix", "_i64");
-//  declareColMajorMatrixSubclass<tdbColMajorMatrix<uint64_t>>(
-//      m, "tdbColMajorMatrix", "_u64");
-
 
+  // Converters from pyarray to matrix
+  declare_pyarray_to_matrix<uint8_t>(m, "_u8");
+  declare_pyarray_to_matrix<uint64_t>(m, "_u64");
+  declare_pyarray_to_matrix<float>(m, "_f32");
+  declare_pyarray_to_matrix<double>(m, "_f64");
 
   /* Query API */
 
@@ -176,7 +199,7 @@ PYBIND11_MODULE(_tiledbvspy, m) {
            const ColMajorMatrix<float>& query_vectors,
            int k,
            bool nth,
-           size_t nthreads) {
+           size_t nthreads) -> ColMajorMatrix<uint64_t> {
           auto r = detail::flat::vq_query_heap(data, query_vectors, k, nthreads);
           return r;
         });
@@ -186,13 +209,13 @@ PYBIND11_MODULE(_tiledbvspy, m) {
            const ColMajorMatrix<float>& query_vectors,
            int k,
            bool nth,
-           size_t nthreads) {
+           size_t nthreads) -> ColMajorMatrix<uint64_t> {
           auto r = detail::flat::vq_query_heap(data, query_vectors, k, nthreads);
           return r;
         });
 
-  m.def("validate_top_k",
-      [](const ColMajorMatrix<size_t>& top_k,
+  m.def("validate_top_k_u64",
+      [](const ColMajorMatrix<uint64_t>& top_k,
          const ColMajorMatrix<int32_t>& ground_truth) -> bool {
         return validate_top_k(top_k, ground_truth);
       });

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

@@ -4,6 +4,7 @@
 
 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 = {}):
@@ -146,3 +147,25 @@ def query_kmeans(
         return kmeans_query_u8(*args)
     else:
         raise TypeError("Unknown type!")
+
+
+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)
+    else:
+        raise TypeError("Unknown type for validate_top_k!")
+
+
+def array_to_matrix(array: np.ndarray):
+    if array.dtype == np.float32:
+        return pyarray_copyto_matrix_f32(array)
+    elif array.dtype == np.float64:
+        return pyarray_copyto_matrix_f64(array)
+    elif array.dtype == np.uint8:
+        return pyarray_copyto_matrix_u8(array)
+    elif array.dtype == np.int32:
+        return pyarray_copyto_matrix_i32(array)
+    elif array.dtype == np.uint64:
+        return pyarray_copyto_matrix_u64(array)
+    else:
+        raise TypeError("Unsupported type!")

apis/python/test/test_api.py

@@ -5,6 +5,7 @@
 import tiledb.vector_search as vs
 from tiledb.vector_search import _tiledbvspy as vspy
 
+import pytest
 
 def test_load_matrix(tmpdir):
     p = str(tmpdir.mkdir("test").join("test.tdb"))
@@ -24,7 +25,10 @@ def test_load_matrix(tmpdir):
     assert np.array_equal(orig_matrix[0, 0], data[0, 0])
 
 
-def test_flat_query(tmpdir):
+@pytest.mark.skipif(
+    not os.path.exists(os.path.expanduser("~/work/proj/vector-search/datasets/sift-andrew/")),
+    reason="requires sift dataset")
+def test_flat_query():
     # db_uri = "s3://tiledb-andrew/sift/sift_base"
     # probe_uri = "s3://tiledb-andrew/sift/sift_query"
     # g_uri = "s3://tiledb-andrew/sift/sift_groundtruth"
@@ -54,4 +58,4 @@ def test_flat_query(tmpdir):
     assert np.array_equal(np.sort(ra[:k], axis=0), np.sort(g[:k, :nqueries], axis=0))
 
     g_m = vs.load_as_matrix(g_uri)
-    assert vspy.validate_top_k(r, g_m)
+    assert vspy.validate_top_k_u64(r, g_m)

apis/python/test/test_index.py

@@ -3,15 +3,27 @@
 import tiledb.vector_search as vs
 from tiledb.vector_search.index import FlatIndex
 
+import numpy as np
+import os
+import pytest
 
-def test_flat_index(tmpdir):
+## only run this test if the dataset is available
+@pytest.mark.skipif(
+    not os.path.exists(os.path.expanduser("~/work/proj/vector-search/datasets/sift-andrew/")),
+    reason="requires sift dataset")
+def test_flat_index_local():
     group_uri = "~/work/proj/vector-search/datasets/sift-andrew/"
     query_uri = "~/work/proj/vector-search/datasets/sift-andrew/sift_query"
+    ground_truth_uri = "~/work/proj/vector-search/datasets/sift-andrew/sift_groundtruth"
 
-    query_vectors = vs.load_as_array(query_uri)[:, :10]
+    nquery = 10
+    query_vectors = vs.load_as_array(query_uri)[:, :nquery]
 
     index = FlatIndex(group_uri, dtype="float32", parts_name="sift_base")
     result = index.query(query_vectors)
     assert isinstance(result, np.ndarray)
 
-    # ground_truth = vs.load_as_array(ground_truth_uri)
+    #ground_truth = vs.load_as_array(ground_truth_uri)
+    ground_truth = vs.load_as_matrix(ground_truth_uri, nquery)
+    result_m = vs.array_to_matrix(result)
+    vs.validate_top_k(result_m, ground_truth)

environment.yml

@@ -0,0 +1,12 @@
+name: tiledb-vs-env
+
+dependencies:
+  - python=3.10
+  - numpy
+  - tiledb
+  - tiledb-py
+  - pybind11
+  - cmake
+  - openblas
+  - ninja
+  - pytest