Add memory tracking for RTree. (#4697)

Add MemoryTracker to RTree using `tdb::pmr` containers.

---
TYPE: NO_HISTORY
DESC: Add memory tracking for RTree.

Author: shaunrd0

tiledb/sm/fragment/fragment_metadata.cc

@@ -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))
@@ -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))
@@ -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();
 }
 

tiledb/sm/fragment/fragment_metadata.h

@@ -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

tiledb/sm/rtree/rtree.cc

@@ -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
@@ -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"
@@ -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              */
 /* ****************************** */
@@ -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();
 }
@@ -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();
 }
 
@@ -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;
@@ -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(
@@ -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

tiledb/sm/rtree/rtree.h

@@ -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
@@ -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"
@@ -48,6 +49,7 @@ namespace sm {
 
 class Buffer;
 class ConstBuffer;
+class MemoryTracker;
 
 enum class Datatype : uint8_t;
 enum class Layout : uint8_t;
@@ -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               */
@@ -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
@@ -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
@@ -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     */
@@ -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
@@ -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.
    *
@@ -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
@@ -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.
@@ -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

tiledb/sm/rtree/test/CMakeLists.txt

@@ -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)

tiledb/sm/rtree/test/compile_rtree_main.cc

@@ -29,6 +29,6 @@
 #include "../rtree.h"
 
 int main() {
-  tiledb::sm::RTree x{};
+  tiledb::sm::RTree x{nullptr, 0, nullptr};
   return 0;
 }