close the array through the ManagedQuery

single-cell-data · Oct 2, 2023 · 764a157 · 764a157
1 parent 8585eea
commit 764a157
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/libtiledbsoma/src/soma/managed_query.cc b/libtiledbsoma/src/soma/managed_query.cc
@@ -54,6 +54,13 @@ ManagedQuery::ManagedQuery(
     reset();
 }
 
+void ManagedQuery::close() {
+    if (query_future_.valid()) {
+        query_future_.wait();
+    }
+    array_->close();
+}
+
 void ManagedQuery::reset() {
     query_ = std::make_unique<Query>(schema_->context(), *array_);
     subarray_ = std::make_unique<Subarray>(schema_->context(), *array_);
@@ -160,9 +167,11 @@ void ManagedQuery::submit_read() {
     if (!is_empty_query()) {
         // Submit query in a separate thread, so we can return immediately
         query_future_ = std::async(std::launch::async, [&]() {
-            LOG_DEBUG("[ManagedQuery] submit thread start");
+            LOG_DEBUG(
+                fmt::format("[ManagedQuery] [{}] submit thread start", name_));
             query_->submit();
-            LOG_DEBUG("[ManagedQuery] submit thread done");
+            LOG_DEBUG(
+                fmt::format("[ManagedQuery] [{}] submit thread done", name_));
         });
     }
     query_submitted_ = true;
@@ -180,7 +189,7 @@ std::shared_ptr<ArrayBuffers> ManagedQuery::results() {
 
     if (!query_submitted_) {
         throw TileDBSOMAError(fmt::format(
-            "[ManagedQuery][{}] submit query before reading results", name_));
+            "[ManagedQuery] [{}] submit query before reading results", name_));
     }
     query_submitted_ = false;
 

diff --git a/libtiledbsoma/src/soma/managed_query.h b/libtiledbsoma/src/soma/managed_query.h
@@ -69,6 +69,13 @@ class ManagedQuery {
     ManagedQuery(ManagedQuery&&) = default;
     ~ManagedQuery() = default;
 
+    /**
+     * @brief Close the array after waiting for any asynchronous queries to
+     * complete.
+     *
+     */
+    void close();
+
     /**
      * @brief Reset the state of this ManagedQuery object to prepare for a new
      * query, while holding the array open.

diff --git a/libtiledbsoma/src/soma/soma_array.cc b/libtiledbsoma/src/soma/soma_array.cc
@@ -180,7 +180,9 @@ void SOMAArray::open(
 }
 
 void SOMAArray::close() {
-    arr_->close();
+    // Close the array through the managed query to ensure any pending queries
+    // are completed.
+    mq_->close();
 }
 
 void SOMAArray::reset(