Add support for an empty tdb_matrix (#268)

Author: jparismorgan

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

@@ -395,14 +395,14 @@ static void declareColMajorMatrixSubclass(
 
   cls.def(
       py::init<
-          const Ctx&,
-          std::string,
-          size_t,
-          size_t,
-          size_t,
-          size_t,
-          size_t,
-          uint64_t>(),
+          const Ctx&,             // ctx
+          std::string,            // uri
+          size_t,                 // first_row
+          std::optional<size_t>,  // last_row
+          size_t,                 // first_col
+          std::optional<size_t>,  // last_col
+          size_t,                 // upper_bound
+          uint64_t>(),            // timestamp
       py::keep_alive<1, 2>());
 
   if constexpr (std::is_same<P, tdbColMajorMatrix<T>>::value) {

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

@@ -11,11 +11,11 @@ def load_as_matrix(
     path: str,
     ctx: "Ctx" = None,
     config: Optional[Mapping[str, Any]] = None,
-    size: int = 0,
+    size: Optional[int] = None,
     timestamp: int = 0,
 ):
     """
-    Load array as Matrix class
+    Load array as Matrix class. We read in all rows (i.e. from 0 to the row domain length).
 
     Parameters
     ----------
@@ -24,7 +24,7 @@ def load_as_matrix(
     ctx: Ctx
         TileDB context
     size: int
-        Size of vectors to load
+        Size of vectors to load. If not set we will read from 0 to the column domain length.
     """
     # If the user passes a tiledb python Config object convert to a dictionary
     if isinstance(config, tiledb.Config):
@@ -35,16 +35,18 @@ def load_as_matrix(
 
     a = tiledb.ArraySchema.load(path, ctx=tiledb.Ctx(config))
     dtype = a.attr(0).dtype
+    # Read all rows from column 0 -> `size`. Set no upper_bound. Note that if `size` is None then 
+    # we'll read to the column domain length.
     if dtype == np.float32:
-        m = tdbColMajorMatrix_f32(ctx, path, 0, 0, 0, size, 0, timestamp)
+        m = tdbColMajorMatrix_f32(ctx, path, 0, None, 0, size, 0, timestamp)
     elif dtype == np.float64:
-        m = tdbColMajorMatrix_f64(ctx, path, 0, 0, 0, size, 0, timestamp)
+        m = tdbColMajorMatrix_f64(ctx, path, 0, None, 0, size, 0, timestamp)
     elif dtype == np.int32:
-        m = tdbColMajorMatrix_i32(ctx, path, 0, 0, 0, size, 0, timestamp)
+        m = tdbColMajorMatrix_i32(ctx, path, 0, None, 0, size, 0, timestamp)
     elif dtype == np.int32:
-        m = tdbColMajorMatrix_i64(ctx, path, 0, 0, 0, size, 0, timestamp)
+        m = tdbColMajorMatrix_i64(ctx, path, 0, None, 0, size, 0, timestamp)
     elif dtype == np.uint8:
-        m = tdbColMajorMatrix_u8(ctx, path, 0, 0, 0, size, 0, timestamp)
+        m = tdbColMajorMatrix_u8(ctx, path, 0, None, 0, size, 0, timestamp)
     # elif dtype == np.uint64:
     #     return tdbColMajorMatrix_u64(ctx, path, size, timestamp)
     else:
@@ -58,6 +60,7 @@ def load_as_array(
     return_matrix: bool = False,
     ctx: "Ctx" = None,
     config: Optional[Mapping[str, Any]] = None,
+    size: Optional[int] = None,
 ):
     """
     Load array as array class
@@ -71,7 +74,7 @@ def load_as_array(
     config: Dict
         TileDB configuration parameters
     """
-    m = load_as_matrix(path, ctx=ctx, config=config)
+    m = load_as_matrix(path, size=size, ctx=ctx, config=config)
     r = np.array(m, copy=False)
 
     # hang on to a copy for testing purposes, for now

apis/python/test/test_api.py

@@ -22,6 +22,24 @@ def test_load_matrix(tmpdir):
     assert np.array_equal(m, data)
     assert np.array_equal(orig_matrix[0, 0], data[0, 0])
 
+def test_load_matrix_specify_size(tmpdir):
+    p = str(tmpdir.mkdir("test").join("test.tdb"))
+    data = np.random.rand(12).astype(np.float32).reshape(3, 4)
+
+    # write some test data with tiledb-py
+    create_array(p, data)
+
+    # test specifying a size
+    m = vs.load_as_array(p, size=data.shape[1])
+    assert np.array_equal(m, data)
+
+    # test specifying a smaller size
+    m = vs.load_as_array(p, size=2)
+    assert np.array_equal(m, data[:, :2])
+
+    # test specifying a size of 0
+    m = vs.load_as_array(p, size=0)
+    assert m.shape == (3, 0)
 
 def test_vector(tmpdir):
     v = vspy._create_vector_u64()

apis/python/test/test_module.py

@@ -13,12 +13,32 @@ def test_tdbMatrix(tmpdir):
     create_array(p, data)
 
     ctx = vspy.Ctx({})
-    m = vspy.tdbColMajorMatrix_f32(ctx, p, 0, 0, 0, 0, 0, 0)
+    # Read all rows and cols automatically.
+    m = vspy.tdbColMajorMatrix_f32(ctx, p, 0, None, 0, None, 0, 0)
     m.load()
     m_array = np.array(m)
     assert m_array.shape == data.shape
     assert np.array_equal(m_array, data)
 
+    # Read all rows and cols by specifying how many rows and cols there are.
+    m = vspy.tdbColMajorMatrix_f32(ctx, p, 0, data.shape[0], 0, data.shape[1], 0, 0)
+    m.load()
+    m_array = np.array(m)
+    assert m_array.shape == data.shape
+    assert np.array_equal(m_array, data)
+
+    # Read all rows and no cols.
+    m_no_fill = vspy.tdbColMajorMatrix_f32(ctx, p, 0, None, 0, 0, 0, 0)
+    m_no_fill.load()
+    m_array = np.array(m_no_fill)
+    assert m_array.shape == (3, 0)
+
+    # Read no rows and no cols.
+    m_no_fill = vspy.tdbColMajorMatrix_f32(ctx, p, 0, 0, 0, 0, 0, 0)
+    m_no_fill.load()
+    m_array = np.array(m_no_fill)
+    assert m_array.shape == (0, 0)
+
     m_array2 = np.array(m, copy=False)  # mutable view
     v = np.random.rand(1).astype(np.float32)
     m_array2[1, 2] = v

src/include/detail/ivf/index.h

@@ -80,11 +80,12 @@ int ivf_index(
   auto non_empty = array.non_empty_domain<int32_t>();
   auto partitions = non_empty[1].second.second + 1;
 
+  // Read all rows from column 0 -> `partitions`. Set no upper_bound.
   auto centroids = tdbColMajorMatrix<centroids_type>(
       ctx,
       centroids_uri,
       0,
-      0,
+      std::nullopt,
       0,
       partitions,
       0,
@@ -218,8 +219,9 @@ int ivf_index(
     size_t end_pos = 0,
     size_t nthreads = 0,
     uint64_t timestamp = 0) {
-  auto db =
-      tdbColMajorMatrix<T>(ctx, db_uri, 0, 0, start_pos, end_pos, 0, timestamp);
+  // Read all rows from column `start_pos` -> `end_pos`. Set no upper_bound.
+  auto db = tdbColMajorMatrix<T>(
+      ctx, db_uri, 0, std::nullopt, start_pos, end_pos, 0, timestamp);
   db.load();
   std::vector<ids_type> external_ids;
   if (external_ids_uri.empty()) {
@@ -261,8 +263,9 @@ int ivf_index(
     size_t end_pos = 0,
     size_t nthreads = 0,
     uint64_t timestamp = 0) {
-  auto db =
-      tdbColMajorMatrix<T>(ctx, db_uri, 0, 0, start_pos, end_pos, 0, timestamp);
+  // Read all rows from column `start_pos` -> `end_pos`. Set no upper_bound.
+  auto db = tdbColMajorMatrix<T>(
+      ctx, db_uri, 0, std::nullopt, start_pos, end_pos, 0, timestamp);
   db.load();
   return ivf_index<T, ids_type, centroids_type>(
       ctx,

src/include/detail/linalg/tdb_matrix.h

@@ -122,26 +122,29 @@ class tdbBlockedMatrix : public MatrixBase {
 
   /**
    * @brief Construct a new tdbBlockedMatrix object, limited to `upper_bound`
-   * vectors. In this case, the `Matrix` is row-major, so the number of vectors
-   * is the number of rows.
+   * vectors. We read rows from 0 -> row domain length and cols from 0 -> col
+   * domain length. In this case, the `Matrix` is column-major, so the number of
+   * vectors is the number of columns.
    *
    * @param ctx The TileDB context to use.
    * @param uri URI of the TileDB array to read.
    */
   tdbBlockedMatrix(const tiledb::Context& ctx, const std::string& uri) noexcept
     requires(std::is_same_v<LayoutPolicy, stdx::layout_left>)
-      : tdbBlockedMatrix(ctx, uri, 0, 0, 0, 0, 0, 0) {
+      : tdbBlockedMatrix(ctx, uri, 0, std::nullopt, 0, std::nullopt, 0, 0) {
   }
 
   /**
    * @brief Construct a new tdbBlockedMatrix object, limited to `upper_bound`
-   * vectors. In this case, the `Matrix` is column-major, so the number of
+   * vectors. We read rows from 0 -> row domain length and cols from 0 -> col
+   * domain length. In this case, the `Matrix` is column-major, so the number of
    * vectors is the number of columns.
    *
    * @param ctx The TileDB context to use.
    * @param uri URI of the TileDB array to read.
-   * @param upper_bound The maximum number of vectors to read.
-   * @param temporal_policy The TemporalPolicy to use for reading the array
+   * @param upper_bound The maximum number of vectors to read. Set to 0 for no
+   * upper bound.
+   * @param timestamp The TemporalPolicy to use for reading the array
    * data.
    */
   tdbBlockedMatrix(
@@ -150,17 +153,33 @@ class tdbBlockedMatrix : public MatrixBase {
       size_t upper_bound,
       size_t timestamp = 0)
     requires(std::is_same_v<LayoutPolicy, stdx::layout_left>)
-      : tdbBlockedMatrix(ctx, uri, 0, 0, 0, 0, upper_bound, timestamp) {
+      : tdbBlockedMatrix(
+            ctx,
+            uri,
+            0,
+            std::nullopt,
+            0,
+            std::nullopt,
+            upper_bound,
+            timestamp) {
   }
 
-  /** General constructor */
+  /** General constructor
+   *
+   * @param first_row The first row to read from.
+   * @param last_row The last row to read to. Read rows from 0 -> row domain
+   * length if nullopt is passed.
+   * @param first_col The first col to read from.
+   * @param last_col The last col to read to. Read rows from 0 -> col domain
+   * length if nullopt is passed.
+   */
   tdbBlockedMatrix(
       const tiledb::Context& ctx,
       const std::string& uri,
       size_t first_row,
-      size_t last_row,
+      std::optional<size_t> last_row,
       size_t first_col,
-      size_t last_col,
+      std::optional<size_t> last_col,
       size_t upper_bound,
       size_t timestamp)
     requires(std::is_same_v<LayoutPolicy, stdx::layout_left>)
@@ -177,14 +196,22 @@ class tdbBlockedMatrix : public MatrixBase {
                  tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp))) {
   }
 
-  /** General constructor */
+  /** General constructor
+   *
+   * @param first_row The first row to read from.
+   * @param last_row The last row to read to. Read rows from 0 -> row domain
+   * length if nullopt is passed.
+   * @param first_col The first col to read from.
+   * @param last_col The last col to read to. Read rows from 0 -> col domain
+   * length if nullopt is passed.
+   */
   tdbBlockedMatrix(
       const tiledb::Context& ctx,
       const std::string& uri,
       size_t first_row,
-      size_t last_row,
+      std::optional<size_t> last_row,
       size_t first_col,
-      size_t last_col,
+      std::optional<size_t> last_col,
       size_t upper_bound,
       tiledb::TemporalPolicy temporal_policy)  // noexcept
     requires(std::is_same_v<LayoutPolicy, stdx::layout_left>)
@@ -194,16 +221,14 @@ class tdbBlockedMatrix : public MatrixBase {
             ctx, uri, TILEDB_READ, temporal_policy))
       , schema_{array_->schema()}
       , first_row_{first_row}
-      , last_row_{last_row}
-      , first_col_{first_col}
-      , last_col_{last_col} {
+      , first_col_{first_col} {
     constructor_timer.stop();
     scoped_timer _{tdb_func__ + " " + uri};
 
-    if (last_row_ < first_row_) {
+    if (last_row && *last_row < first_row_) {
       throw std::runtime_error("last_row < first_row");
     }
-    if (last_col_ < first_col_) {
+    if (last_col && *last_col < first_col_) {
       throw std::runtime_error("last_col < first_col");
     }
 
@@ -228,15 +253,17 @@ class tdbBlockedMatrix : public MatrixBase {
 
     /* The size of the array may not be the size of domain.  Use non-zero value
      * if set in constructor */
-    if (last_row_ == 0) {
-      last_row_ =
-          (row_domain.template domain<row_domain_type>().second -
-           row_domain.template domain<row_domain_type>().first + 1);
+    if (!last_row.has_value()) {
+      last_row_ = row_domain.template domain<row_domain_type>().second -
+                  row_domain.template domain<row_domain_type>().first + 1;
+    } else {
+      last_row_ = *last_row;
     }
-    if (last_col_ == 0) {
-      last_col_ =
-          (col_domain.template domain<col_domain_type>().second -
-           col_domain.template domain<col_domain_type>().first + 1);
+    if (!last_col.has_value()) {
+      last_col_ = col_domain.template domain<col_domain_type>().second -
+                  col_domain.template domain<col_domain_type>().first + 1;
+    } else {
+      last_col_ = *last_col;
     }
 
     size_t dimension = last_row_ - first_row_;
@@ -296,7 +323,7 @@ class tdbBlockedMatrix : public MatrixBase {
         std::min(load_blocksize_, last_col_ - last_resident_col_);
 
     // Return if we're at the end
-    if (elements_to_load == 0) {
+    if (elements_to_load == 0 || dimension == 0) {
       return false;
     }
 
@@ -363,14 +390,15 @@ class tdbPreLoadMatrix : public tdbBlockedMatrix<T, LayoutPolicy, I> {
       const std::string& uri,
       size_t upper_bound = 0,
       uint64_t timestamp = 0)
-      : tdbPreLoadMatrix(ctx, uri, 0, 0, upper_bound, timestamp) {
+      : tdbPreLoadMatrix(
+            ctx, uri, std::nullopt, std::nullopt, upper_bound, timestamp) {
   }
 
   tdbPreLoadMatrix(
       const tiledb::Context& ctx,
       const std::string& uri,
-      size_t num_array_rows,
-      size_t num_array_cols,
+      std::optional<size_t> num_array_rows,
+      std::optional<size_t> num_array_cols,
       size_t upper_bound = 0,
       uint64_t timestamp = 0)
       : Base(
@@ -384,31 +412,6 @@ class tdbPreLoadMatrix : public tdbBlockedMatrix<T, LayoutPolicy, I> {
             timestamp) {
     Base::load();
   }
-
-  tdbPreLoadMatrix(
-      const tiledb::Context& ctx,
-      const std::string& uri,
-      size_t upper_bound,
-      const tiledb::TemporalPolicy& temporal_policy)
-      : tdbPreLoadMatrix(ctx, uri, 0, 0, upper_bound, temporal_policy) {
-  }
-
-  tdbPreLoadMatrix(
-      const tiledb::Context& ctx,
-      const std::string& uri,
-      size_t num_array_rows,
-      size_t num_array_cols,
-      size_t upper_bound,
-      const tiledb::TemporalPolicy& temporal_policy)
-      : Base(
-            ctx,
-            uri,
-            num_array_rows,
-            num_array_cols,
-            upper_bound,
-            temporal_policy) {
-    Base::load();
-  }
 };
 
 /**