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

Co-authored-by: nguyenv <[email protected]>

Author: github-actions[bot]

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> {

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

libtiledbsoma/src/soma/managed_query.cc

@@ -1252,4 +1252,4 @@ bool ManagedQuery::_extend_and_evolve_schema<std::string>(
     }
     return false;
 }
-};  // namespace tiledbsoma
+};  // namespace tiledbsoma

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

libtiledbsoma/src/soma/soma_array.cc

@@ -269,30 +269,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(

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();