[c++, r, python] read_next clean-up (#3450)

single-cell-data · Dec 19, 2024 · 345d596 · 345d596
1 parent bda7d97
commit 345d596
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 87 deletions.
diff --git a/apis/python/src/tiledbsoma/managed_query.cc b/apis/python/src/tiledbsoma/managed_query.cc
@@ -61,7 +61,6 @@ void load_managed_query(py::module& m) {
             py::arg("ctx"),
             py::arg("name") = "unnamed")
 
-        .def("setup_read", &ManagedQuery::setup_read)
         .def("is_empty_query", &ManagedQuery::is_empty_query)
         .def("is_complete", &ManagedQuery::is_complete)
 
@@ -117,25 +116,6 @@ void load_managed_query(py::module& m) {
             "if_not_empty"_a = false,
             "replace"_a = false)
 
-        .def(
-            "submit_read",
-            &ManagedQuery::submit_read,
-            py::call_guard<py::gil_scoped_release>())
-        .def(
-            "results",
-            [](ManagedQuery& mq) -> std::optional<py::object> {
-                try {
-                    // Release python GIL before reading data
-                    py::gil_scoped_release release;
-                    auto tbl = mq.results();
-                    // Acquire python GIL before accessing python objects
-                    py::gil_scoped_acquire acquire;
-                    return to_table(std::make_optional(tbl));
-                } catch (const std::exception& e) {
-                    throw TileDBSOMAError(e.what());
-                }
-            })
-
         .def(
             "next",
             [](ManagedQuery& mq) -> std::optional<py::object> {

diff --git a/apis/r/src/riterator.cpp b/apis/r/src/riterator.cpp
@@ -169,16 +169,7 @@ Rcpp::XPtr<tdbs::SOMAArray> sr_setup(
 // [[Rcpp::export]]
 bool sr_complete(Rcpp::XPtr<tdbs::SOMAArray> sr) {
     check_xptr_tag<tdbs::SOMAArray>(sr);
-    bool complt = sr->is_complete(true);
-    bool initial = sr->is_initial_read();
-    bool res = complt && !initial;  // completed transfer if query status
-                                    // complete and query ran once
-    spdl::debug(
-        "[sr_complete] Complete query test {} (compl {} initial {})",
-        res,
-        complt,
-        initial);
-    return res;
+    return sr->results_complete();
 }
 
 //' @noRd

diff --git a/libtiledbsoma/src/soma/managed_query.cc b/libtiledbsoma/src/soma/managed_query.cc
@@ -1253,4 +1253,4 @@ bool ManagedQuery::_extend_and_evolve_schema<std::string>(
     }
     return false;
 }
-};  // namespace tiledbsoma
+};  // namespace tiledbsoma
diff --git a/libtiledbsoma/src/soma/managed_query.h b/libtiledbsoma/src/soma/managed_query.h
@@ -286,12 +286,6 @@ class ManagedQuery {
         std::unique_ptr<ArrowSchema> arrow_schema,
         std::unique_ptr<ArrowArray> arrow_array);
 
-    /**
-     * @brief Configure query and allocate result buffers for reads.
-     *
-     */
-    void setup_read();
-
     std::optional<std::shared_ptr<ArrayBuffers>> read_next();
 
     /**
@@ -380,19 +374,6 @@ class ManagedQuery {
         return buffers_->at(name)->string_view(index);
     }
 
-    /**
-     * @brief Submit the query.
-     *
-     */
-    void submit_read();
-
-    /**
-     * @brief Return results from the query.
-     *
-     * @return std::shared_ptr<ArrayBuffers>
-     */
-    std::shared_ptr<ArrayBuffers> results();
-
     /**
      * @brief Submit the write query.
      *
@@ -445,6 +426,25 @@ class ManagedQuery {
     //= private non-static
     //===================================================================
 
+    /**
+     * @brief Configure query and allocate result buffers for reads.
+     *
+     */
+    void setup_read();
+
+    /**
+     * @brief Submit the query.
+     *
+     */
+    void submit_read();
+
+    /**
+     * @brief Return results from the query.
+     *
+     * @return std::shared_ptr<ArrayBuffers>
+     */
+    std::shared_ptr<ArrayBuffers> results();
+
     /**
      * @brief Check if column name is contained in the query results.
      *
@@ -880,4 +880,4 @@ bool ManagedQuery::_extend_and_evolve_schema<std::string>(
 
 };  // namespace tiledbsoma
 
-#endif
+#endif
diff --git a/libtiledbsoma/src/soma/soma_array.cc b/libtiledbsoma/src/soma/soma_array.cc
@@ -271,30 +271,7 @@ void SOMAArray::reset(
 }
 
 std::optional<std::shared_ptr<ArrayBuffers>> SOMAArray::read_next() {
-    // If the query is complete, return `std::nullopt`
-    if (mq_->is_complete(true)) {
-        return std::nullopt;
-    }
-
-    // Configure query and allocate result buffers
-    mq_->setup_read();
-
-    // Continue to submit the empty query on first read to return empty results
-    if (mq_->is_empty_query()) {
-        if (first_read_next_) {
-            first_read_next_ = false;
-            return mq_->results();
-        } else {
-            return std::nullopt;
-        }
-    }
-
-    first_read_next_ = false;
-
-    mq_->submit_read();
-
-    // Return the results, possibly incomplete
-    return mq_->results();
+    return mq_->read_next();
 }
 
 void SOMAArray::set_column_data(

diff --git a/libtiledbsoma/test/unit_managed_query.cc b/libtiledbsoma/test/unit_managed_query.cc
@@ -128,10 +128,7 @@ TEST_CASE("ManagedQuery: Basic execution test") {
     auto [array, d0, a0, _] = create_array(uri, *ctx);
 
     auto mq = ManagedQuery(array, ctx);
-    mq.setup_read();
-
-    mq.submit_read();
-    auto results = mq.results();
+    auto results = mq.read_next();
     REQUIRE(mq.results_complete());
 
     auto num_cells = mq.total_num_cells();
@@ -152,10 +149,7 @@ TEST_CASE("ManagedQuery: Select test") {
     auto mq = ManagedQuery(array, ctx);
     mq.select_columns({attr_name});
     mq.select_points<std::string>(dim_name, {"a"});
-    mq.setup_read();
-
-    mq.submit_read();
-    auto results = mq.results();
+    auto results = mq.read_next();
     REQUIRE(mq.results_complete());
 
     auto num_cells = mq.total_num_cells();
@@ -178,10 +172,7 @@ TEST_CASE("ManagedQuery: Validity test") {
     auto [array, d0, a0, a0_valids] = create_array(uri, *ctx);
 
     auto mq = ManagedQuery(array, ctx);
-    mq.setup_read();
-
-    mq.submit_read();
-    auto results = mq.results();
+    auto results = mq.read_next();
     REQUIRE(mq.results_complete());
 
     auto num_cells = mq.total_num_cells();