[c++] ManagedQuery::set_condition should not reset entire instance …

…state (#3418) (#3420) * set_condition should not reset entire instance state * lint * rework ManagedQuery API a bit * lint Co-authored-by: Bruce Martin <[email protected]>
single-cell-data · Dec 11, 2024 · 0f0639c · 0f0639c
1 parent 0535744
commit 0f0639c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/apis/python/src/tiledbsoma/managed_query.cc b/apis/python/src/tiledbsoma/managed_query.cc
@@ -98,8 +98,7 @@ void load_managed_query(py::module& m) {
                              .ptr()
                              .get();
                 }
-                mq.reset();
-                mq.select_columns(column_names);
+                mq.select_columns(column_names, false, true);
 
                 // Release python GIL after we're done accessing python
                 // objects
@@ -115,7 +114,8 @@ void load_managed_query(py::module& m) {
             "select_columns",
             &ManagedQuery::select_columns,
             "names"_a,
-            "if_not_empty"_a = false)
+            "if_not_empty"_a = false,
+            "replace"_a = false)
 
         .def(
             "submit_read",

diff --git a/libtiledbsoma/src/soma/managed_query.cc b/libtiledbsoma/src/soma/managed_query.cc
@@ -107,13 +107,17 @@ void ManagedQuery::set_layout(ResultOrder layout) {
 }
 
 void ManagedQuery::select_columns(
-    const std::vector<std::string>& names, bool if_not_empty) {
+    const std::vector<std::string>& names, bool if_not_empty, bool replace) {
     // Return if we are selecting all columns (columns_ is empty) and we want to
     // continue selecting all columns (if_not_empty == true).
     if (if_not_empty && columns_.empty()) {
         return;
     }
 
+    if (replace) {
+        reset_columns();
+    }
+
     for (auto& name : names) {
         // Name is not an attribute or dimension.
         if (!schema_->has_attribute(name) &&
@@ -128,6 +132,10 @@ void ManagedQuery::select_columns(
     }
 }
 
+void ManagedQuery::reset_columns() {
+    columns_.clear();
+}
+
 void ManagedQuery::setup_read() {
     // If the query is complete, return so we do not submit it again
     auto status = query_->query_status();

diff --git a/libtiledbsoma/src/soma/managed_query.h b/libtiledbsoma/src/soma/managed_query.h
@@ -135,11 +135,22 @@ class ManagedQuery {
      * list of selected columns is empty. This prevents a `select_columns` call
      * from changing an empty list (all columns) to a subset of columns.
      *
+     * NB: you may only select a given column once. Selecting twice will
+     * generate an error in read_next.
+     *
      * @param names Vector of column names
      * @param if_not_empty Prevent changing an "empty" selection of all columns
+     * @param replace Column names will replace any existing selected columns.
      */
     void select_columns(
-        const std::vector<std::string>& names, bool if_not_empty = false);
+        const std::vector<std::string>& names,
+        bool if_not_empty = false,
+        bool replace = false);
+
+    /**
+     * @brief Reset column selection to none, aka "all".
+     */
+    void reset_columns();
 
     /**
      * @brief Returns the column names set by the query.