diff --git a/tiledb/sm/query/writers/global_order_writer.cc b/tiledb/sm/query/writers/global_order_writer.cc index eda3eed9e95..59de7d01255 100644 --- 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 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 tiles; + tdb::pmr::unordered_map 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& coord_dups, - std::unordered_map* tiles) const { + tdb::pmr::unordered_map* 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& tiles) { + tdb::pmr::unordered_map& tiles) { // Cache variables to prevent map lookups. const auto buf_names = buffer_names(); std::vector var_size; diff --git a/tiledb/sm/query/writers/global_order_writer.h b/tiledb/sm/query/writers/global_order_writer.h index bf48eafb232..c15b81f67c8 100644 --- 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 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 last_tiles_; + tdb::pmr::unordered_map 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 last_var_offsets_; + tdb::pmr::unordered_map last_var_offsets_; /** * Stores the number of cells written for each attribute/dimension across * the write operations. */ - std::unordered_map cells_written_; + tdb::pmr::unordered_map cells_written_; /** The fragment metadata that the writer will focus on. */ shared_ptr frag_meta_; @@ -301,7 +311,7 @@ class GlobalOrderWriter : public WriterBase { */ Status prepare_full_tiles( const std::set& coord_dups, - std::unordered_map* tiles) const; + tdb::pmr::unordered_map* 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& tiles); + tdb::pmr::unordered_map& 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 index 94d3bd71899..f6af441af35 100644 --- 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* tiles) const { + tdb::pmr::unordered_map* 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 tiles; + tdb::pmr::unordered_map 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 index 2136ba509e7..5a9871c44b0 100644 --- 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& cell_pos() { + tdb::pmr::vector& 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 cell_pos_; + tdb::pmr::vector cell_pos_; /** The set with the positions of duplicate coordinates/cells. */ std::set coord_dups_; @@ -173,7 +173,7 @@ class UnorderedWriter : public WriterBase { * @return Status */ Status prepare_tiles( - std::unordered_map* tiles) const; + tdb::pmr::unordered_map* 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 index 262a7133193..af536be9336 100644 --- 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 meta) const { } std::vector WriterBase::compute_mbrs( - const std::unordered_map& tiles) const { + const tdb::pmr::unordered_map& tiles) + const { auto timer_se = stats_->start_timer("compute_coord_meta"); // Applicable only if there are coordinates @@ -669,7 +670,7 @@ std::vector WriterBase::compute_mbrs( void WriterBase::set_coords_metadata( const uint64_t start_tile_idx, const uint64_t end_tile_idx, - const std::unordered_map& tiles, + const tdb::pmr::unordered_map& tiles, const std::vector& mbrs, shared_ptr 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& tiles) const { + tdb::pmr::unordered_map& 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* tiles) { + tdb::pmr::unordered_map* 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 frag_meta, - std::unordered_map* const tiles) { + tdb::pmr::unordered_map* 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 index 99be55ce2ca..10a4464f5c9 100644 --- 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 compute_mbrs( - const std::unordered_map& tiles) + const tdb::pmr::unordered_map& 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& tiles, + const tdb::pmr::unordered_map& tiles, const std::vector& mbrs, shared_ptr meta) const; @@ -279,7 +279,7 @@ class WriterBase : public StrategyBase, public IQueryStrategy { */ Status compute_tiles_metadata( uint64_t tile_num, - std::unordered_map& tiles) const; + tdb::pmr::unordered_map& 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* tiles); + tdb::pmr::unordered_map* 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 frag_meta, - std::unordered_map* tiles); + tdb::pmr::unordered_map* tiles); /** * Writes the input tiles for the input attribute/dimension to storage.