Close tdbMatrix and tdbMatrixWithIds Array's when we have nothing lef…

…t to read (#466)
TileDB-Inc · Jul 26, 2024 · 7adba24 · 7adba24
1 parent 5db39ff
commit 7adba24
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 4 deletions.
diff --git a/src/include/detail/linalg/tdb_matrix.h b/src/include/detail/linalg/tdb_matrix.h
@@ -114,6 +114,10 @@ class tdbBlockedMatrix : public MatrixBase {
   // size_t pending_row_offset{0};
   // size_t pending_col_offset{0};
 
+  size_t get_elements_to_load() const {
+    return std::min(load_blocksize_, last_col_ - last_resident_col_);
+  }
+
  public:
   tdbBlockedMatrix(tdbBlockedMatrix&& rhs) = default;
 
@@ -324,11 +328,11 @@ class tdbBlockedMatrix : public MatrixBase {
     }
 
     size_t dimension = last_row_ - first_row_;
-    auto elements_to_load =
-        std::min(load_blocksize_, last_col_ - last_resident_col_);
+    auto elements_to_load = get_elements_to_load();
 
     // Return if we're at the end
     if (elements_to_load == 0 || dimension == 0) {
+      array_->close();
       return false;
     }
 
@@ -363,6 +367,10 @@ class tdbBlockedMatrix : public MatrixBase {
       throw std::runtime_error("Query status is not complete");
     }
 
+    if (get_elements_to_load() == 0) {
+      array_->close();
+    }
+
     num_loads_++;
     return true;
   }

diff --git a/src/include/detail/linalg/tdb_matrix_with_ids.h b/src/include/detail/linalg/tdb_matrix_with_ids.h
@@ -167,6 +167,7 @@ class tdbBlockedMatrixWithIds
   bool load() {
     scoped_timer _{tdb_func__ + " " + this->ids_uri_};
     if (!Base::load()) {
+      ids_array_->close();
       return false;
     }
 
@@ -214,6 +215,10 @@ class tdbBlockedMatrixWithIds
       throw std::runtime_error("Query status for IDs is not complete");
     }
 
+    if (this->get_elements_to_load() == 0) {
+      ids_array_->close();
+    }
+
     return true;
   }
 };  // tdbBlockedMatrixWithIds

diff --git a/src/include/test/unit_tdb_matrix.cc b/src/include/test/unit_tdb_matrix.cc
@@ -59,7 +59,10 @@ TEMPLATE_TEST_CASE("constructors", "[tdb_matrix]", float, uint8_t) {
   write_matrix(ctx, X, tmp_matrix_uri);
 
   auto Y = tdbColMajorMatrix<TestType>(ctx, tmp_matrix_uri);
-  Y.load();
+  CHECK(Y.load() == true);
+  for (int i = 0; i < 5; ++i) {
+    CHECK(Y.load() == false);
+  }
 
   auto Z = tdbColMajorMatrix<TestType>(std::move(Y));
 

diff --git a/src/include/test/unit_tdb_matrix_with_ids.cc b/src/include/test/unit_tdb_matrix_with_ids.cc
@@ -63,7 +63,10 @@ TEMPLATE_TEST_CASE(
 
   auto Y = tdbColMajorMatrixWithIds<TestType, TestType>(
       ctx, tmp_matrix_uri, tmp_ids_uri);
-  Y.load();
+  CHECK(Y.load() == true);
+  for (int i = 0; i < 5; ++i) {
+    CHECK(Y.load() == false);
+  }
   CHECK(num_vectors(Y) == num_vectors(X));
   CHECK(dimensions(Y) == dimensions(X));
   CHECK(std::equal(