Add constructor argument to Query to override the memory budget. (#4968)

[SC-45187](https://app.shortcut.com/tiledb-inc/story/45187/enable-to-set-reader-budget-on-consolidation)

This PR adds an optional argument to the internal `Query` class, that
overrides the total memory budget. This will be useful for better memory
management during consolidation.

---
TYPE: NO_HISTORY

test/src/unit-Reader.cc

@@ -175,6 +175,7 @@ TEST_CASE_METHOD(
       context.storage_manager(),
       array.opened_array(),
       config,
+      nullopt,
       buffers,
       aggregate_buffers,
       subarray,

tiledb/sm/query/query.cc

@@ -74,7 +74,8 @@ namespace sm {
 Query::Query(
     StorageManager* storage_manager,
     shared_ptr<Array> array,
-    optional<std::string> fragment_name)
+    optional<std::string> fragment_name,
+    optional<uint64_t> memory_budget)
     : query_memory_tracker_(
           storage_manager->resources().create_memory_tracker())
     , array_shared_(array)
@@ -105,6 +106,7 @@ Query::Query(
     , is_dimension_label_ordered_read_(false)
     , dimension_label_increasing_(true)
     , fragment_size_(std::numeric_limits<uint64_t>::max())
+    , memory_budget_(memory_budget)
     , query_remote_buffer_storage_(std::nullopt)
     , default_channel_{make_shared<QueryChannel>(HERE(), *this, 0)} {
   assert(array->is_open());
@@ -1756,6 +1758,7 @@ Status Query::create_strategy(bool skip_checks_serialization) {
       storage_manager_,
       opened_array_,
       config_,
+      memory_budget_,
       buffers_,
       aggregate_buffers_,
       subarray_,

tiledb/sm/query/query.h

@@ -155,11 +155,14 @@ class Query {
    * writes.
    * @param fragment_base_uri Optional base name for new fragment. Only used for
    *     writes and only if fragment_uri is empty.
+   * @param memory_budget Total memory budget. If set to nullopt, the value
+   *     will be obtained from the sm.mem.total_budget config option.
    */
   Query(
       StorageManager* storage_manager,
       shared_ptr<Array> array,
-      optional<std::string> fragment_name = nullopt);
+      optional<std::string> fragment_name = nullopt,
+      optional<uint64_t> memory_budget = nullopt);
 
   /** Destructor. */
   virtual ~Query();
@@ -1089,6 +1092,12 @@ class Query {
    */
   uint64_t fragment_size_;
 
+  /**
+   * Memory budget. If set to nullopt, the value will be obtained from the
+   * sm.mem.total_budget config option.
+   */
+  optional<uint64_t> memory_budget_;
+
   /** Already written buffers. */
   std::unordered_set<std::string> written_buffers_;
 

tiledb/sm/query/readers/dense_reader.cc

@@ -72,7 +72,9 @@ DenseReader::DenseReader(
     shared_ptr<Logger> logger,
     StrategyParams& params,
     bool remote_query)
-    : ReaderBase(stats, logger->clone("DenseReader", ++logger_id_), params) {
+    : ReaderBase(stats, logger->clone("DenseReader", ++logger_id_), params)
+    , memory_budget_(params.memory_budget().value_or(0))
+    , memory_budget_from_query_(params.memory_budget()) {
   elements_mode_ = false;
 
   // Sanity checks.
@@ -121,9 +123,11 @@ QueryStatusDetailsReason DenseReader::status_incomplete_reason() const {
 void DenseReader::refresh_config() {
   // Get config values.
   bool found = false;
-  throw_if_not_ok(
-      config_.get<uint64_t>("sm.mem.total_budget", &memory_budget_, &found));
-  assert(found);
+  if (!memory_budget_from_query_.has_value()) {
+    throw_if_not_ok(
+        config_.get<uint64_t>("sm.mem.total_budget", &memory_budget_, &found));
+    assert(found);
+  }
   throw_if_not_ok(config_.get<uint64_t>(
       "sm.mem.tile_upper_memory_limit", &tile_upper_memory_limit_, &found));
   assert(found);

tiledb/sm/query/readers/dense_reader.h

@@ -158,6 +158,9 @@ class DenseReader : public ReaderBase, public IQueryStrategy {
   /** Total memory budget. */
   uint64_t memory_budget_;
 
+  /** Total memory budget if overridden by the query. */
+  optional<uint64_t> memory_budget_from_query_;
+
   /** Target upper memory limit for tiles. */
   uint64_t tile_upper_memory_limit_;
 

tiledb/sm/query/readers/sparse_index_reader_base.cc

@@ -72,7 +72,7 @@ SparseIndexReaderBase::SparseIndexReaderBase(
     : ReaderBase(stats, logger, params)
     , read_state_(array_->fragment_metadata().size())
     , tmp_read_state_(array_->fragment_metadata().size())
-    , memory_budget_(config_, reader_string)
+    , memory_budget_(config_, reader_string, params.memory_budget())
     , include_coords_(include_coords)
     , memory_used_for_coords_total_(0)
     , deletes_consolidation_no_purge_(

tiledb/sm/query/readers/sparse_index_reader_base.h

@@ -110,7 +110,12 @@ class MemoryBudget {
 
   MemoryBudget() = delete;
 
-  MemoryBudget(Config& config, std::string reader_string) {
+  MemoryBudget(
+      Config& config,
+      std::string reader_string,
+      optional<uint64_t> total_budget)
+      : total_budget_(total_budget.value_or(0))
+      , memory_budget_from_query_(total_budget) {
     refresh_config(config, reader_string);
   }
 
@@ -128,8 +133,10 @@ class MemoryBudget {
    * @param reader_string String to identify the reader settings to load.
    */
   void refresh_config(Config& config, std::string reader_string) {
-    total_budget_ =
-        config.get<uint64_t>("sm.mem.total_budget", Config::must_find);
+    if (!memory_budget_from_query_.has_value()) {
+      total_budget_ =
+          config.get<uint64_t>("sm.mem.total_budget", Config::must_find);
+    }
 
     ratio_coords_ = config.get<double>(
         "sm.mem.reader." + reader_string + ".ratio_coords", Config::must_find);
@@ -192,6 +199,9 @@ class MemoryBudget {
   /** Total memory budget. */
   uint64_t total_budget_;
 
+  /** Total memory budget if overridden by the query. */
+  optional<uint64_t> memory_budget_from_query_;
+
   /** How much of the memory budget is reserved for coords. */
   double ratio_coords_;
 

tiledb/sm/query/strategy_base.h

@@ -70,6 +70,7 @@ class StrategyParams {
       StorageManager* storage_manager,
       shared_ptr<OpenedArray> array,
       Config& config,
+      optional<uint64_t> memory_budget,
       std::unordered_map<std::string, QueryBuffer>& buffers,
       std::unordered_map<std::string, QueryBuffer>& aggregate_buffers,
       Subarray& subarray,
@@ -82,6 +83,7 @@ class StrategyParams {
       , storage_manager_(storage_manager)
       , array_(array)
       , config_(config)
+      , memory_budget_(memory_budget)
       , buffers_(buffers)
       , aggregate_buffers_(aggregate_buffers)
       , subarray_(subarray)
@@ -119,6 +121,11 @@ class StrategyParams {
     return config_;
   };
 
+  /** Return the memory budget, if set. */
+  inline optional<uint64_t> memory_budget() {
+    return memory_budget_;
+  }
+
   /** Return the buffers. */
   inline std::unordered_map<std::string, QueryBuffer>& buffers() {
     return buffers_;
@@ -174,6 +181,12 @@ class StrategyParams {
   /** Config for query-level parameters only. */
   Config& config_;
 
+  /**
+   * Memory budget for the query. If set to nullopt, the value will be obtained
+   * from the sm.mem.total_budget config option.
+   */
+  optional<uint64_t> memory_budget_;
+
   /** Buffers. */
   std::unordered_map<std::string, QueryBuffer>& buffers_;