Use more pmr containers for tracking

TileDB-Inc · Mar 12, 2024 · 5240b1b · 5240b1b
1 parent 72f7d63
commit 5240b1b
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 24 deletions.
diff --git a/tiledb/sm/query/writers/global_order_writer.cc b/tiledb/sm/query/writers/global_order_writer.cc
@@ -112,6 +112,13 @@ GlobalOrderWriter::GlobalOrderWriter(
 GlobalOrderWriter::~GlobalOrderWriter() {
 }
 
+GlobalOrderWriter::GlobalWriteState::GlobalWriteState(
+    shared_ptr<MemoryTracker> memory_tracker)
+    : last_tiles_(memory_tracker->get_resource(MemoryType::TILE_WRITER_DATA))
+    , last_var_offsets_(memory_tracker->get_resource(MemoryType::TILE_OFFSETS))
+    , cells_written_(memory_tracker->get_resource(MemoryType::TILE_SUMS)) {
+}
+
 /* ****************************** */
 /*               API              */
 /* ****************************** */
@@ -174,7 +181,7 @@ Status GlobalOrderWriter::alloc_global_write_state() {
     return logger_->status(
         Status_WriterError("Cannot initialize global write state; State not "
                            "properly finalized"));
-  global_write_state_.reset(new GlobalWriteState);
+  global_write_state_.reset(new GlobalWriteState(query_memory_tracker_));
 
   // Alloc FragmentMetadata object
   global_write_state_->frag_meta_ = this->create_fragment_metadata();
@@ -749,7 +756,8 @@ Status GlobalOrderWriter::global_write() {
     RETURN_CANCEL_OR_ERROR(compute_coord_dups(&coord_dups));
   }
 
-  std::unordered_map<std::string, WriterTileTupleVector> tiles;
+  tdb::pmr::unordered_map<std::string, WriterTileTupleVector> tiles(
+      query_memory_tracker_->get_resource(MemoryType::TILE_WRITER_DATA));
   RETURN_CANCEL_OR_ERROR(prepare_full_tiles(coord_dups, &tiles));
 
   // Find number of tiles and gather stats
@@ -859,7 +867,7 @@ void GlobalOrderWriter::nuke_global_write_state() {
 
 Status GlobalOrderWriter::prepare_full_tiles(
     const std::set<uint64_t>& coord_dups,
-    std::unordered_map<std::string, WriterTileTupleVector>* tiles) const {
+    tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* tiles) const {
   auto timer_se = stats_->start_timer("prepare_tiles");
 
   // Initialize attribute and coordinate tiles
@@ -1378,7 +1386,7 @@ Status GlobalOrderWriter::prepare_full_tiles_var(
 uint64_t GlobalOrderWriter::num_tiles_to_write(
     uint64_t start,
     uint64_t tile_num,
-    std::unordered_map<std::string, WriterTileTupleVector>& tiles) {
+    tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles) {
   // Cache variables to prevent map lookups.
   const auto buf_names = buffer_names();
   std::vector<bool> var_size;

diff --git a/tiledb/sm/query/writers/global_order_writer.h b/tiledb/sm/query/writers/global_order_writer.h
@@ -58,6 +58,16 @@ class GlobalOrderWriter : public WriterBase {
    * by successive query submissions until the query is finalized.
    */
   struct GlobalWriteState {
+    /** Deleted Default Constructor. */
+    GlobalWriteState() = delete;
+
+    /**
+     * Constructor.
+     *
+     * @param memory_tracker The memory tracker for the underlying containers.
+     */
+    explicit GlobalWriteState(shared_ptr<MemoryTracker> memory_tracker);
+
     /**
      * Stores the last tile of each attribute/dimension for each write
      * operation. The key is the attribute/dimension name. For fixed-sized
@@ -66,7 +76,7 @@ class GlobalOrderWriter : public WriterBase {
      * second tile is the values tile. In both cases, the third tile stores a
      * validity tile for nullable attributes.
      */
-    std::unordered_map<std::string, WriterTileTupleVector> last_tiles_;
+    tdb::pmr::unordered_map<std::string, WriterTileTupleVector> last_tiles_;
 
     /**
      * Stores the last offset into the var size tile buffer for var size
@@ -75,13 +85,13 @@ class GlobalOrderWriter : public WriterBase {
      * Note: Once tiles are created with the correct size from the beginning,
      * this variable can go awaty.
      */
-    std::unordered_map<std::string, uint64_t> last_var_offsets_;
+    tdb::pmr::unordered_map<std::string, uint64_t> last_var_offsets_;
 
     /**
      * Stores the number of cells written for each attribute/dimension across
      * the write operations.
      */
-    std::unordered_map<std::string, uint64_t> cells_written_;
+    tdb::pmr::unordered_map<std::string, uint64_t> cells_written_;
 
     /** The fragment metadata that the writer will focus on. */
     shared_ptr<FragmentMetadata> frag_meta_;
@@ -301,7 +311,7 @@ class GlobalOrderWriter : public WriterBase {
    */
   Status prepare_full_tiles(
       const std::set<uint64_t>& coord_dups,
-      std::unordered_map<std::string, WriterTileTupleVector>* tiles) const;
+      tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* tiles) const;
 
   /**
    * Applicable only to write in global order. It prepares only full
@@ -373,7 +383,7 @@ class GlobalOrderWriter : public WriterBase {
   uint64_t num_tiles_to_write(
       uint64_t start,
       uint64_t tile_num,
-      std::unordered_map<std::string, WriterTileTupleVector>& tiles);
+      tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles);
 
   /**
    * Close the current fragment and start a new one. The closed fragment will

diff --git a/tiledb/sm/query/writers/unordered_writer.cc b/tiledb/sm/query/writers/unordered_writer.cc
@@ -89,6 +89,8 @@ UnorderedWriter::UnorderedWriter(
           remote_query,
           fragment_name)
     , frag_uri_(std::nullopt)
+    , cell_pos_(
+          query_memory_tracker_->get_resource(MemoryType::TILE_WRITER_DATA))
     , written_buffers_(written_buffers)
     , is_coords_pass_(true) {
   // Check the layout is unordered.
@@ -372,7 +374,7 @@ Status UnorderedWriter::compute_coord_dups() {
 }
 
 Status UnorderedWriter::prepare_tiles(
-    std::unordered_map<std::string, WriterTileTupleVector>* tiles) const {
+    tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* tiles) const {
   auto timer_se = stats_->start_timer("prepare_tiles");
 
   // Initialize attribute tiles
@@ -675,7 +677,8 @@ Status UnorderedWriter::unordered_write() {
   frag_uri_ = frag_meta_->fragment_uri();
 
   // Prepare tiles
-  std::unordered_map<std::string, WriterTileTupleVector> tiles;
+  tdb::pmr::unordered_map<std::string, WriterTileTupleVector> tiles(
+      query_memory_tracker_->get_resource(MemoryType::TILE_WRITER_DATA));
   RETURN_CANCEL_OR_ERROR(prepare_tiles(&tiles));
 
   // No tiles

diff --git a/tiledb/sm/query/writers/unordered_writer.h b/tiledb/sm/query/writers/unordered_writer.h
@@ -88,7 +88,7 @@ class UnorderedWriter : public WriterBase {
   Status alloc_frag_meta();
 
   /** Returns the cell position vector. */
-  std::vector<uint64_t>& cell_pos() {
+  tdb::pmr::vector<uint64_t>& cell_pos() {
     return cell_pos_;
   }
 
@@ -119,7 +119,7 @@ class UnorderedWriter : public WriterBase {
    * The positions that resulted from sorting and according to which the cells
    * must be re-arranged.
    */
-  std::vector<uint64_t> cell_pos_;
+  tdb::pmr::vector<uint64_t> cell_pos_;
 
   /** The set with the positions of duplicate coordinates/cells. */
   std::set<uint64_t> coord_dups_;
@@ -173,7 +173,7 @@ class UnorderedWriter : public WriterBase {
    * @return Status
    */
   Status prepare_tiles(
-      std::unordered_map<std::string, WriterTileTupleVector>* tiles) const;
+      tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* tiles) const;
 
   /**
    * It prepares the tiles for the input attribute or dimension, re-organizing

diff --git a/tiledb/sm/query/writers/writer_base.cc b/tiledb/sm/query/writers/writer_base.cc
@@ -623,7 +623,8 @@ Status WriterBase::close_files(shared_ptr<FragmentMetadata> meta) const {
 }
 
 std::vector<NDRange> WriterBase::compute_mbrs(
-    const std::unordered_map<std::string, WriterTileTupleVector>& tiles) const {
+    const tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles)
+    const {
   auto timer_se = stats_->start_timer("compute_coord_meta");
 
   // Applicable only if there are coordinates
@@ -669,7 +670,7 @@ std::vector<NDRange> WriterBase::compute_mbrs(
 void WriterBase::set_coords_metadata(
     const uint64_t start_tile_idx,
     const uint64_t end_tile_idx,
-    const std::unordered_map<std::string, WriterTileTupleVector>& tiles,
+    const tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles,
     const std::vector<NDRange>& mbrs,
     shared_ptr<FragmentMetadata> meta) const {
   // Applicable only if there are coordinates
@@ -701,7 +702,7 @@ void WriterBase::set_coords_metadata(
 
 Status WriterBase::compute_tiles_metadata(
     uint64_t tile_num,
-    std::unordered_map<std::string, WriterTileTupleVector>& tiles) const {
+    tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles) const {
   auto compute_tp = storage_manager_->compute_tp();
 
   // Parallelize over attributes?
@@ -805,7 +806,7 @@ Status WriterBase::create_fragment(
 }
 
 Status WriterBase::filter_tiles(
-    std::unordered_map<std::string, WriterTileTupleVector>* tiles) {
+    tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* tiles) {
   auto timer_se = stats_->start_timer("filter_tiles");
   auto status = parallel_for(
       storage_manager_->compute_tp(), 0, tiles->size(), [&](uint64_t i) {
@@ -1048,7 +1049,7 @@ Status WriterBase::write_tiles(
     const uint64_t start_tile_idx,
     const uint64_t end_tile_idx,
     shared_ptr<FragmentMetadata> frag_meta,
-    std::unordered_map<std::string, WriterTileTupleVector>* const tiles) {
+    tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* const tiles) {
   auto timer_se = stats_->start_timer("write_num_tiles");
 
   assert(!tiles->empty());

diff --git a/tiledb/sm/query/writers/writer_base.h b/tiledb/sm/query/writers/writer_base.h
@@ -249,7 +249,7 @@ class WriterBase : public StrategyBase, public IQueryStrategy {
    * @return MBRs.
    */
   std::vector<NDRange> compute_mbrs(
-      const std::unordered_map<std::string, WriterTileTupleVector>& tiles)
+      const tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles)
       const;
 
   /**
@@ -265,7 +265,7 @@ class WriterBase : public StrategyBase, public IQueryStrategy {
   void set_coords_metadata(
       const uint64_t start_tile_idx,
       const uint64_t end_tile_idx,
-      const std::unordered_map<std::string, WriterTileTupleVector>& tiles,
+      const tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles,
       const std::vector<NDRange>& mbrs,
       shared_ptr<FragmentMetadata> meta) const;
 
@@ -279,7 +279,7 @@ class WriterBase : public StrategyBase, public IQueryStrategy {
    */
   Status compute_tiles_metadata(
       uint64_t tile_num,
-      std::unordered_map<std::string, WriterTileTupleVector>& tiles) const;
+      tdb::pmr::unordered_map<std::string, WriterTileTupleVector>& tiles) const;
 
   /**
    * Returns the i-th coordinates in the coordinate buffers in string
@@ -305,7 +305,7 @@ class WriterBase : public StrategyBase, public IQueryStrategy {
    * of the pipeline.
    */
   Status filter_tiles(
-      std::unordered_map<std::string, WriterTileTupleVector>* tiles);
+      tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* tiles);
 
   /**
    * Runs the input tiles for the input attribute through the filter pipeline.
@@ -441,7 +441,7 @@ class WriterBase : public StrategyBase, public IQueryStrategy {
       const uint64_t start_tile_idx,
       const uint64_t end_tile_idx,
       shared_ptr<FragmentMetadata> frag_meta,
-      std::unordered_map<std::string, WriterTileTupleVector>* tiles);
+      tdb::pmr::unordered_map<std::string, WriterTileTupleVector>* tiles);
 
   /**
    * Writes the input tiles for the input attribute/dimension to storage.