From 0f0639c6e284dfe5f027efbbf1d70f6c36dea723 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Wed, 11 Dec 2024 12:04:01 -0500
Subject: [PATCH] [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 <bruce.martin@tiledb.com>
---
 apis/python/src/tiledbsoma/managed_query.cc |  6 +++---
 libtiledbsoma/src/soma/managed_query.cc     | 10 +++++++++-
 libtiledbsoma/src/soma/managed_query.h      | 13 ++++++++++++-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/apis/python/src/tiledbsoma/managed_query.cc b/apis/python/src/tiledbsoma/managed_query.cc
index e3865f5237..b6322fa381 100644
--- 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
index 7aad1e67fc..1993cef703 100644
--- 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
index ab76a48482..a097c6dcc4 100644
--- 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.