Merge branch 'main' into npapa/cpp_ingestion

Author: Nikos Papailiou

.github/workflows/ci_python.yml

@@ -28,8 +28,9 @@ jobs:
         run: |
           cd apis/python
           pip install .[test]
-          pytest -k "not ingest" # TODO: requires token
+          pytest
+          # TODO: fix editable on linux
           #pip uninstall -y tiledb.vector_search
           #pip install -e .
-          #pytest -k "not ingest" # TODO: requires token
+          #pytest
         shell: bash -el {0}

README.md

@@ -14,13 +14,7 @@
 
 # Quick Installation
 
-TileDB-Vector-Search is available from either [PyPI](https://pypi.org/project/tiledb-vector-search/) with ``pip``:
-
-```
-pip install tiledb-vector-search
-```
-
-or from [conda-forge](https://anaconda.org/conda-forge/tiledb-vector-searcg) with
+Pre-built packages are available from the [tiledb conda channel](https://anaconda.org/tiledb/tiledb-vector-search) using
 [conda](https://conda.io/docs/) or [mamba](https://github.com/mamba-org/mamba#installation):
 
 ```
@@ -33,4 +27,4 @@ We welcome contributions. Please see [`Building`](Building.md) for
 development-build instructions. For large new
 features, please open an issue to discuss goals and approach in order
 to ensure a smooth PR integration and review process. All contributions
-must be licensed under the repository's [MIT License](../LICENSE).
+must be licensed under the repository's [MIT License](../LICENSE).

apis/python/CMakeLists.txt

@@ -49,7 +49,7 @@ target_link_libraries(${VSPY_TARGET_NAME}
     kmeans_lib
     TileDB::tiledb_shared)
 
-target_compile_definitions(${VSPY_TARGET_NAME} PRIVATE VERSION_INFO=${PROJECT_VERSION})
+target_compile_definitions(${VSPY_TARGET_NAME} PRIVATE TILEDBVS_ENABLE_STATS VERSION_INFO=${PROJECT_VERSION})
 
 if (APPLE)
   set_target_properties(${VSPY_TARGET_NAME} PROPERTIES INSTALL_RPATH "@loader_path/lib")

apis/python/pyproject.toml

@@ -20,10 +20,11 @@ classifiers = [
 dependencies = [
     "tiledb-cloud>=0.10.5",
     "tiledb>=0.15.2",
+    "typing-extensions" # for tiledb-cloud indirect
 ]
 
 [project.optional-dependencies]
-test = ["pytest"]
+test = ["pytest", "scikit-learn", "tiledb-cloud"]
 
 
 [project.urls]

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

@@ -14,6 +14,9 @@ using Ctx = tiledb::Context;
 bool global_debug = true;
 double global_time_of_interest;
 
+bool enable_stats = false;
+std::vector<json> core_stats;
+
 PYBIND11_MAKE_OPAQUE(std::vector<uint32_t>);
 PYBIND11_MAKE_OPAQUE(std::vector<uint64_t>);
 
@@ -299,6 +302,24 @@ PYBIND11_MODULE(_tiledbvspy, m) {
         return validate_top_k(top_k, ground_truth);
       });
 
+  m.def("stats_enable", []() {
+    enable_stats = true;
+    tiledb::Stats::enable();
+  });
+
+  m.def("stats_disable", []() {
+    enable_stats = false;
+    tiledb::Stats::disable();
+  });
+
+  m.def("stats_reset", []() {
+    core_stats.clear();
+  });
+
+  m.def("stats_dump", []() {
+    return json{core_stats}.dump();
+  });
+
   declare_kmeans_query<uint8_t>(m, "u8");
   declare_kmeans_query<float>(m, "f32");
 

documentation/index.md

@@ -14,17 +14,17 @@
 
 # Quick Installation
 
-TileDB-Vector-Search is available from either [PyPI](https://pypi.org/project/tiledb-vector-search/) with ``pip``:
+TileDB-Vector-Search is available from the [tiledb conda channel](https://anaconda.org/tiledb/tiledb-vector-search) with
+[conda](https://conda.io/docs/) or [mamba](https://github.com/mamba-org/mamba#installation)
 
 ```
-pip install tiledb-vector-search
+conda install -c tiledb tiledb-vector-search
 ```
 
-~~or from the [tiledb conda channel](https://anaconda.org/tiledb/tiledb-vector-search) with
-[conda](https://conda.io/docs/) or [mamba](https://github.com/mamba-org/mamba#installation)~~ (Pending 🚧):
+~~Or [PyPI](https://pypi.org/project/tiledb-vector-search/) with ``pip``~~: (Pending 🚧)
 
 ```
-conda install -c tiledb tiledb-vector-search
+pip install tiledb-vector-search
 ```
 
 # Contributing
@@ -33,4 +33,4 @@ We welcome contributions. Please see [`Building`](Building.md) for
 development-build instructions. For large new
 features, please open an issue to discuss goals and approach in order
 to ensure a smooth PR integration and review process. All contributions
-must be licensed under the repository's [MIT License](../LICENSE).
+must be licensed under the repository's [MIT License](../LICENSE).

src/include/detail/linalg/tdb_helpers.h

@@ -0,0 +1,81 @@
+/**
+ * @file   tiledb_helpers.h
+ *
+ * @section LICENSE
+ *
+ * The MIT License
+ *
+ * @copyright Copyright (c) 2023 TileDB, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ *
+ * @section DESCRIPTION
+ *
+ * Helper functions for certain TileDB operations.
+ *
+ */
+
+#ifndef TILEDB_HELPERS_H
+#define TILEDB_HELPERS_H
+
+#include "stats.h"
+#include <tiledb/tiledb>
+
+namespace tiledb_helpers {
+
+/**
+ * @brief Opens a TileDB array and displays stats to stderr.
+ *
+ * Stats are only collected if the TILEDBVS_ENABLE_STATS symbol is
+ * defined, and a variable named enable_stats is set to true.
+ * The stats are written to a FILE* specified by the variable named stats_file.
+ *
+ * @param function_name The name of the function calling this. You can use the tdb_func__ macro.
+ * @param ctx The TileDB context to use.
+ * @param uri The URI of the array to open.
+ * @param query_type The mode to open the array.
+ */
+inline tiledb::Array open_array(const std::string &function_name,
+                                const tiledb::Context &ctx,
+                                const std::string &uri,
+                                tiledb_query_type_t query_type) {
+  StatsCollectionScope stats_scope(uri, function_name, "open_array");
+  return tiledb::Array(ctx, uri, query_type);
+}
+
+/**
+ * @brief Submits a TileDB query and displays stats to stderr.
+ *
+ * Stats are only collected if the TILEDBVS_ENABLE_STATS symbol is
+ * defined, and a variable named enable_stats is set to true.
+ * The stats are written to a FILE* specified by the variable named stats_file.
+ *
+ * @param function_name The name of the function calling this. You can use the tdb_func__ macro.
+ * @param query The query to submit.
+ */
+inline void submit_query(const std::string &function_name,
+                         const std::string &uri,
+                         tiledb::Query &query) {
+  StatsCollectionScope stats_scope(uri, function_name, "submit_query");
+  query.submit();
+}
+
+} // namespace tiledb_helpers
+
+#endif

src/include/detail/linalg/tdb_io.h

@@ -100,7 +100,8 @@ void write_matrix(
       (int)start_pos + (int)A.num_cols() - 1};
 
   // Open array for writing
-  tiledb::Array array(ctx, uri, TILEDB_WRITE);
+  tiledb::Array array =
+      tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_WRITE);
 
   tiledb::Subarray subarray(ctx, array);
   subarray.set_subarray(subarray_vals);
@@ -113,7 +114,7 @@ void write_matrix(
       .set_data_buffer(
           "values", &A(0, 0), (uint64_t)A.num_rows() * (uint64_t)A.num_cols())
       .set_subarray(subarray);
-  query.submit();
+  tiledb_helpers::submit_query(tdb_func__, uri, query);
 
   assert(tiledb::Query::Status::COMPLETE == query.query_status());
 
@@ -167,7 +168,8 @@ void write_vector(
       (int)start_pos, (int)start_pos + (int)size(v) - 1};
 
   // Open array for writing
-  tiledb::Array array(ctx, uri, TILEDB_WRITE);
+  tiledb::Array array =
+      tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_WRITE);
 
   tiledb::Subarray subarray(ctx, array);
   subarray.set_subarray(subarray_vals);
@@ -176,8 +178,10 @@ void write_vector(
   query.set_layout(TILEDB_ROW_MAJOR)
       .set_data_buffer("values", v)
       .set_subarray(subarray);
+
   query.submit();
   assert(tiledb::Query::Status::COMPLETE == query.query_status());
+  tiledb_helpers::submit_query(tdb_func__, uri, query);
 
   array.close();
 }
@@ -193,7 +197,8 @@ std::vector<T> read_vector(const tiledb::Context& ctx, const std::string& uri) {
     std::cerr << "# Reading std::vector: " << uri << std::endl;
   }
 
-  auto array_ = tiledb::Array{ctx, uri, TILEDB_READ};
+  tiledb::Array array_ =
+      tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_READ);
   auto schema_ = array_.schema();
 
   using domain_type = int32_t;
@@ -225,7 +230,7 @@ std::vector<T> read_vector(const tiledb::Context& ctx, const std::string& uri) {
   tiledb::Query query(ctx, array_);
   query.set_subarray(subarray).set_data_buffer(
       attr_name, data_.data(), vec_rows_);
-  query.submit();
+  tiledb_helpers::submit_query(tdb_func__, uri, query);
   _memory_data.insert_entry(tdb_func__, vec_rows_ * sizeof(T));
 
   array_.close();

src/include/detail/linalg/tdb_matrix.h

@@ -72,6 +72,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
   log_timer constructor_timer{"tdbMatrix constructor"};
 
   std::reference_wrapper<const tiledb::Context> ctx_;
+  std::string uri_;
   tiledb::Array array_;
   tiledb::ArraySchema schema_;
   std::unique_ptr<T[]> backing_data_;
@@ -195,7 +196,8 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
       size_t col_begin,
       size_t col_end)  // noexcept
       : ctx_{ctx}
-      , array_{ctx, uri, TILEDB_READ}
+      , uri_{uri}
+      , array_{tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_READ)}
       , schema_{array_.schema()} {
     constructor_timer.stop();
     scoped_timer _{tdb_func__ + uri};
@@ -278,7 +280,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
     query.set_subarray(subarray)
         .set_layout(layout_order)
         .set_data_buffer(attr_name, data_.get(), num_rows * num_cols);
-    query.submit();
+    tiledb_helpers::submit_query(tdb_func__, uri, query);
     _memory_data.insert_entry(tdb_func__, num_rows * num_cols * sizeof(T));
 
     // assert(tiledb::Query::Status::COMPLETE == query.query_status());
@@ -391,7 +393,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
         query.set_subarray(subarray)
             .set_layout(layout_order)
             .set_data_buffer(attr_name, ptr, num_elements);
-        query.submit();
+        tiledb_helpers::submit_query(tdb_func__, uri, query);
         _memory_data.insert_entry(tdb_func__, num_elements * sizeof(T));
 
         // assert(tiledb::Query::Status::COMPLETE == query.query_status());
@@ -413,7 +415,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
        */
       auto attr_idx = 0;
 
-      auto ids_array_ = tiledb::Array{ctx_, id_uri, TILEDB_READ};
+      tiledb::Array ids_array_ = tiledb_helpers::open_array(tdb_func__, ctx_, id_uri, TILEDB_READ);
       auto ids_schema_ = ids_array_.schema();
 
       auto attr_num{ids_schema_.attribute_num()};
@@ -445,7 +447,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
         auto ptr = part_ids.data() + offset;
         query.set_subarray(subarray).set_data_buffer(
             attr_name, ptr, num_elements);
-        query.submit();
+        tiledb_helpers::submit_query(tdb_func__, uri, query);
         _memory_data.insert_entry(tdb_func__, num_elements * sizeof(T));
 
         if (tiledb::Query::Status::COMPLETE != query.query_status()) {
@@ -552,7 +554,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
     query.set_subarray(subarray)
         .set_layout(layout_order)
         .set_data_buffer(attr_name, this_data, read_size);
-    query.submit();
+    tiledb_helpers::submit_query(tdb_func__, uri_, query);
     _memory_data.insert_entry(tdb_func__, read_size * sizeof(T));
 
     // assert(tiledb::Query::Status::COMPLETE == query.query_status());

src/include/detail/linalg/tdb_partitioned_matrix.h

@@ -106,6 +106,7 @@ class tdbPartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
   log_timer constructor_timer{"tdbPartitionedMatrix constructor"};
 
   std::reference_wrapper<const tiledb::Context> ctx_;
+  std::string uri_;
   tiledb::Array array_;
   tiledb::ArraySchema schema_;
   std::unique_ptr<T[]> backing_data_;
@@ -178,9 +179,10 @@ class tdbPartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
       size_t nthreads)
       : constructor_timer{tdb_func__ + std::string{" constructor"}}
       , ctx_{ctx}
-      , array_{ctx_, uri, TILEDB_READ}
+      , uri_{uri}
+      , array_{tiledb_helpers::open_array(tdb_func__, ctx_, uri, TILEDB_READ)}
       , schema_{array_.schema()}
-      , ids_array_{ctx_, ids_uri, TILEDB_READ}
+      , ids_array_{tiledb_helpers::open_array(tdb_func__, ctx_, ids_uri, TILEDB_READ)}
       , ids_schema_{ids_array_.schema()}
       , indices_{std::move(in_indices)}
       , parts_{in_parts}
@@ -345,7 +347,7 @@ class tdbPartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
       query.set_subarray(subarray)
           .set_layout(layout_order)
           .set_data_buffer(attr_name, ptr, col_count * dimension);
-      query.submit();
+      tiledb_helpers::submit_query(tdb_func__, uri_, query);
       _memory_data.insert_entry(tdb_func__, col_count * dimension * sizeof(T));
 
       // assert(tiledb::Query::Status::COMPLETE == query.query_status());