Skip to content

Commit

Permalink
Add memory tracking for RTree. (#4697)
Browse files Browse the repository at this point in the history 
Add MemoryTracker to RTree using `tdb::pmr` containers.

---
TYPE: NO_HISTORY
DESC: Add memory tracking for RTree.
  • Loading branch information
@shaunrd0
shaunrd0 authored Feb 16, 2024
1 parent 14e99dd commit c1e3427
Show file tree
Hide file tree
Showing 8 changed files with 143 additions and 162 deletions.
6 changes: 4 additions & 2 deletions tiledb/sm/fragment/fragment_metadata.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ FragmentMetadata::FragmentMetadata(
ContextResources* resources, shared_ptr<MemoryTracker> memory_tracker)
: resources_(resources)
, memory_tracker_(memory_tracker)
, rtree_(RTree(nullptr, constants::rtree_fanout, memory_tracker_))
, tile_offsets_(memory_tracker_->get_resource(MemoryType::TILE_OFFSETS))
, tile_var_offsets_(memory_tracker_->get_resource(MemoryType::TILE_OFFSETS))
, tile_var_sizes_(memory_tracker_->get_resource(MemoryType::TILE_OFFSETS))
Expand Down Expand Up @@ -119,7 +120,8 @@ FragmentMetadata::FragmentMetadata(
, has_delete_meta_(has_deletes_meta)
, sparse_tile_num_(0)
, meta_file_size_(0)
, rtree_(RTree(&array_schema_->domain(), constants::rtree_fanout))
, rtree_(RTree(
&array_schema_->domain(), constants::rtree_fanout, memory_tracker_))
, tile_index_base_(0)
, tile_offsets_(memory_tracker_->get_resource(MemoryType::TILE_OFFSETS))
, tile_var_offsets_(memory_tracker_->get_resource(MemoryType::TILE_OFFSETS))
Expand Down Expand Up @@ -1576,7 +1578,7 @@ const NDRange& FragmentMetadata::mbr(uint64_t tile_idx) const {
return rtree_.leaf(tile_idx);
}

const std::vector<NDRange>& FragmentMetadata::mbrs() const {
const tdb::pmr::vector<NDRange>& FragmentMetadata::mbrs() const {
return rtree_.leaves();
}

Expand Down
2 changes: 1 addition & 1 deletion tiledb/sm/fragment/fragment_metadata.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -783,7 +783,7 @@ class FragmentMetadata {
const NDRange& mbr(uint64_t tile_idx) const;

/** Returns all the MBRs of all tiles in the fragment. */
const std::vector<NDRange>& mbrs() const;
const tdb::pmr::vector<NDRange>& mbrs() const;

/**
* Retrieves the size of the tile when it is persisted (e.g. the size of the
Expand Down
77 changes: 23 additions & 54 deletions tiledb/sm/rtree/rtree.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
*
* The MIT License
*
* @copyright Copyright (c) 2017-2021 TileDB, Inc.
* @copyright Copyright (c) 2024 TileDB, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
Expand All @@ -32,6 +32,7 @@

#include "tiledb/sm/rtree/rtree.h"
#include "tiledb/common/logger.h"
#include "tiledb/common/memory_tracker.h"
#include "tiledb/sm/array_schema/dimension.h"
#include "tiledb/sm/buffer/buffer.h"
#include "tiledb/sm/enums/datatype.h"
Expand All @@ -53,43 +54,18 @@ namespace sm {
/* CONSTRUCTORS & DESTRUCTORS */
/* ****************************** */

RTree::RTree() {
domain_ = nullptr;
fanout_ = 0;
deserialized_buffer_size_ = 0;
}

RTree::RTree(const Domain* domain, unsigned fanout)
: domain_(domain)
, fanout_(fanout) {
RTree::RTree(
const Domain* domain,
unsigned fanout,
shared_ptr<MemoryTracker> memory_tracker)
: memory_tracker_(memory_tracker)
, domain_(domain)
, fanout_(fanout)
, levels_(memory_tracker_->get_resource(MemoryType::RTREE)) {
}

RTree::~RTree() = default;

RTree::RTree(const RTree& rtree)
: RTree() {
auto clone = rtree.clone();
swap(clone);
}

RTree::RTree(RTree&& rtree) noexcept
: RTree() {
swap(rtree);
}

RTree& RTree::operator=(const RTree& rtree) {
auto clone = rtree.clone();
swap(clone);

return *this;
}

RTree& RTree::operator=(RTree&& rtree) noexcept {
swap(rtree);

return *this;
}

/* ****************************** */
/* API */
/* ****************************** */
Expand Down Expand Up @@ -240,7 +216,7 @@ const NDRange& RTree::leaf(uint64_t leaf_idx) const {
return levels_.back()[leaf_idx];
}

const std::vector<NDRange>& RTree::leaves() const {
const tdb::pmr::vector<NDRange>& RTree::leaves() const {
assert(!levels_.empty());
return levels_.back();
}
Expand Down Expand Up @@ -297,10 +273,10 @@ Status RTree::set_leaf(uint64_t leaf_id, const NDRange& mbr) {
return Status::Ok();
}

Status RTree::set_leaves(const std::vector<NDRange>& mbrs) {
Status RTree::set_leaves(const tdb::pmr::vector<NDRange>& mbrs) {
levels_.clear();
levels_.resize(1);
levels_[0] = mbrs;
levels_[0].assign(mbrs.begin(), mbrs.end());
return Status::Ok();
}

Expand Down Expand Up @@ -328,13 +304,21 @@ void RTree::deserialize(
deserialize_v5(deserializer, domain);
}

void RTree::reset(const Domain* domain, unsigned int fanout) {
domain_ = domain;
fanout_ = fanout;
free_memory();
}

/* ****************************** */
/* PRIVATE METHODS */
/* ****************************** */

RTree::Level RTree::build_level(const Level& level) {
auto cur_mbr_num = (uint64_t)level.size();
Level new_level((uint64_t)std::ceil((double)cur_mbr_num / fanout_));
Level new_level(
(uint64_t)std::ceil((double)cur_mbr_num / fanout_),
memory_tracker_->get_resource(MemoryType::RTREE));
auto new_mbr_num = (uint64_t)new_level.size();

uint64_t mbrs_visited = 0;
Expand All @@ -344,16 +328,7 @@ RTree::Level RTree::build_level(const Level& level) {
domain_->expand_ndrange(level[mbrs_visited], &new_level[i]);
}

return new_level;
}

RTree RTree::clone() const {
RTree clone;
clone.domain_ = domain_;
clone.fanout_ = fanout_;
clone.levels_ = levels_;

return clone;
return {new_level, memory_tracker_->get_resource(MemoryType::RTREE)};
}

void RTree::deserialize_v1_v4(
Expand Down Expand Up @@ -424,11 +399,5 @@ void RTree::deserialize_v5(Deserializer& deserializer, const Domain* domain) {
domain_ = domain;
}

void RTree::swap(RTree& rtree) {
std::swap(domain_, rtree.domain_);
std::swap(fanout_, rtree.fanout_);
std::swap(levels_, rtree.levels_);
}

} // namespace sm
} // namespace tiledb
54 changes: 25 additions & 29 deletions tiledb/sm/rtree/rtree.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
*
* The MIT License
*
* @copyright Copyright (c) 2017-2021 TileDB, Inc.
* @copyright Copyright (c) 2024 TileDB, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
Expand Down Expand Up @@ -36,6 +36,7 @@
#include <vector>

#include "tiledb/common/common.h"
#include "tiledb/common/pmr.h"
#include "tiledb/common/status.h"
#include "tiledb/sm/array_schema/domain.h"
#include "tiledb/sm/misc/tile_overlap.h"
Expand All @@ -48,6 +49,7 @@ namespace sm {

class Buffer;
class ConstBuffer;
class MemoryTracker;

enum class Datatype : uint8_t;
enum class Layout : uint8_t;
Expand All @@ -62,27 +64,19 @@ class RTree {
/* ********************************* */
/* CONSTRUCTORS & DESTRUCTORS */
/* ********************************* */
RTree() = delete;

/** Constructor. */
RTree();

/** Constructor. */
RTree(const Domain* domain, unsigned fanout);
RTree(
const Domain* domain,
unsigned fanout,
shared_ptr<MemoryTracker> memory_tracker);

/** Destructor. */
~RTree();

/** Copy constructor. This performs a deep copy. */
RTree(const RTree& rtree);

/** Move constructor. */
RTree(RTree&& rtree) noexcept;

/** Copy-assign operator. This performs a deep copy. */
RTree& operator=(const RTree& rtree);

/** Move-assign operator. */
RTree& operator=(RTree&& rtree) noexcept;
DISABLE_COPY_AND_COPY_ASSIGN(RTree);
DISABLE_MOVE_AND_MOVE_ASSIGN(RTree);

/* ********************************* */
/* API */
Expand Down Expand Up @@ -125,7 +119,7 @@ class RTree {
const NDRange& leaf(uint64_t leaf_idx) const;

/** Returns the leaves of the tree. */
const std::vector<NDRange>& leaves() const;
const tdb::pmr::vector<NDRange>& leaves() const;

/**
* Returns the number of leaves that are stored in a (full) subtree
Expand Down Expand Up @@ -158,7 +152,7 @@ class RTree {
* Sets the input MBRs as leaves. This will destroy the existing
* RTree.
*/
Status set_leaves(const std::vector<NDRange>& mbrs);
Status set_leaves(const tdb::pmr::vector<NDRange>& mbrs);

/**
* Resizes the leaf level. It destroys the upper levels
Expand All @@ -176,6 +170,14 @@ class RTree {
void deserialize(
Deserializer& deserializer, const Domain* domain, uint32_t version);

/**
* Resets the RTree with the input domain and fanout.
*
* @param domain The domain to use for the RTree.
* @param fanout The fanout of the RTree.
*/
void reset(const Domain* domain, unsigned fanout);

private:
/* ********************************* */
/* PRIVATE TYPE DEFINITIONS */
Expand Down Expand Up @@ -203,7 +205,7 @@ class RTree {
* `levels_`, where the first level is the root. This is how
* we can infer which tree level each `Level` object corresponds to.
*/
typedef std::vector<NDRange> Level;
typedef tdb::pmr::vector<NDRange> Level;

/**
* Defines an R-Tree level entry, which corresponds to a node
Expand All @@ -222,6 +224,9 @@ class RTree {
/* PRIVATE ATTRIBUTES */
/* ********************************* */

/** Memory tracker for the RTree. */
shared_ptr<MemoryTracker> memory_tracker_;

/**
* The domain for which this R-tree provides an index.
*
Expand All @@ -237,7 +242,7 @@ class RTree {
* The tree levels. The first level is the root. Note that the root
* always consists of a single MBR.
*/
std::vector<Level> levels_;
tdb::pmr::vector<Level> levels_;

/**
* Stores the size of the buffer used to deserialize the data, used for
Expand All @@ -252,9 +257,6 @@ class RTree {
/** Builds a single tree level on top of the input level. */
Level build_level(const Level& level);

/** Returns a deep copy of this RTree. */
RTree clone() const;

/**
* Deserializes the contents of the object from the input buffer based
* on the format version.
Expand All @@ -272,12 +274,6 @@ class RTree {
* Applicable to versions >= 5
*/
void deserialize_v5(Deserializer& deserializer, const Domain* domain);

/**
* Swaps the contents (all field values) of this RTree with the
* given ``rtree``.
*/
void swap(RTree& rtree);
};

} // namespace sm
Expand Down
1 change: 1 addition & 0 deletions tiledb/sm/rtree/test/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,5 +28,6 @@ include(unit_test)

commence(unit_test rtree)
this_target_sources(main.cc unit_rtree.cc)
this_target_link_libraries(tiledb_test_support_lib)
this_target_object_libraries(rtree)
conclude(unit_test)
2 changes: 1 addition & 1 deletion tiledb/sm/rtree/test/compile_rtree_main.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,6 @@
#include "../rtree.h"

int main() {
tiledb::sm::RTree x{};
tiledb::sm::RTree x{nullptr, 0, nullptr};
return 0;
}
Loading

0 comments on commit c1e3427

Please sign in to comment.