Let HashJoinBridge be able to reclaim memory in the middle state

tanjialiang · Oct 24, 2024 · ed909c1 · ed909c1
1 parent 3fa7acf
commit ed909c1
Show file tree

Hide file tree

Showing 7 changed files with 304 additions and 22 deletions.
diff --git a/velox/exec/HashBuild.cpp b/velox/exec/HashBuild.cpp
@@ -768,8 +768,26 @@ bool HashBuild::finishHashBuild() {
       RuntimeCounter(timing.wallNanos, RuntimeCounter::Unit::kNanos));
 
   addRuntimeStats();
+
+  // Setup spill function for spilling hash table directly from hash join
+  // bridge after transferring of table ownership.
+  HashJoinTableSpillFunc tableSpillFunc;
+  if (canReclaim()) {
+    VELOX_CHECK_NOT_NULL(spiller_);
+    tableSpillFunc = [hashBitRange = spiller_->hashBits(),
+                      joinNode = joinNode_,
+                      spillConfig = spillConfig(),
+                      spillStats =
+                          &spillStats_](std::shared_ptr<BaseHashTable> table) {
+      return spillHashJoinTable(
+          table, hashBitRange, joinNode, spillConfig, spillStats);
+    };
+  }
   joinBridge_->setHashTable(
-      std::move(table_), std::move(spillPartitions), joinHasNullKeys_);
+      std::move(table_),
+      std::move(spillPartitions),
+      joinHasNullKeys_,
+      tableSpillFunc);
   if (canSpill()) {
     stateCleared_ = true;
   }

diff --git a/velox/exec/HashJoinBridge.cpp b/velox/exec/HashJoinBridge.cpp
@@ -64,6 +64,52 @@ void HashJoinBridge::addBuilder() {
   ++numBuilders_;
 }
 
+bool HashJoinBridge::canReclaim() const {
+  return tableSpillFunc_ != nullptr && !probeStarted_.load() &&
+      buildResult_.has_value() && buildResult_->table != nullptr &&
+      buildResult_->table->numDistinct() != 0;
+}
+
+uint64_t HashJoinBridge::reclaim() {
+  VELOX_CHECK(buildResult_.has_value());
+  VELOX_CHECK_NOT_NULL(buildResult_->table);
+
+  auto computeTableReservedBytes = [](std::vector<RowContainer*> allRows) {
+    uint64_t totalReservedBytes{0};
+    for (const auto* rowContainer : allRows) {
+      totalReservedBytes += rowContainer->pool()->reservedBytes();
+    }
+    return totalReservedBytes;
+  };
+  const auto oldMemUsage =
+      computeTableReservedBytes(buildResult_->table->allRows());
+
+  auto spillPartitionSet = tableSpillFunc_(buildResult_->table);
+  buildResult_->table->clear(true);
+
+  const auto reclaimedBytes =
+      oldMemUsage - computeTableReservedBytes(buildResult_->table->allRows());
+
+  auto spillPartitionIdSet = toSpillPartitionIdSet(spillPartitionSet);
+  if (restoringSpillPartitionId_.has_value()) {
+    for (const auto& id : spillPartitionIdSet) {
+      VELOX_DCHECK_LT(
+          restoringSpillPartitionId_->partitionBitOffset(),
+          id.partitionBitOffset());
+    }
+  }
+  for (auto& partitionEntry : spillPartitionSet) {
+    const auto id = partitionEntry.first;
+    VELOX_CHECK_EQ(spillPartitionSets_.count(id), 0);
+    spillPartitionSets_.emplace(id, std::move(partitionEntry.second));
+  }
+  buildResult_->restoredPartitionId = restoringSpillPartitionId_;
+  buildResult_->spillPartitionIds = spillPartitionIdSet;
+  restoringSpillPartitionId_.reset();
+
+  return reclaimedBytes;
+}
+
 namespace {
 // Create spiller for spilling the row container from one of the sub-table from
 // 'table' to parallelize the table spilling. The function spills all the rows
@@ -188,7 +234,8 @@ SpillPartitionSet spillHashJoinTable(
 void HashJoinBridge::setHashTable(
     std::unique_ptr<BaseHashTable> table,
     SpillPartitionSet spillPartitionSet,
-    bool hasNullKeys) {
+    bool hasNullKeys,
+    const HashJoinTableSpillFunc& tableSpillFunc) {
   VELOX_CHECK_NOT_NULL(table, "setHashTable called with null table");
 
   auto spillPartitionIdSet = toSpillPartitionIdSet(spillPartitionSet);
@@ -199,7 +246,7 @@ void HashJoinBridge::setHashTable(
     VELOX_CHECK(started_);
     VELOX_CHECK(!buildResult_.has_value());
     VELOX_CHECK(restoringSpillShards_.empty());
-
+    tableSpillFunc_ = tableSpillFunc;
     if (restoringSpillPartitionId_.has_value()) {
       for (const auto& id : spillPartitionIdSet) {
         VELOX_DCHECK_LT(
@@ -266,7 +313,7 @@ std::optional<HashJoinBridge::HashBuildResult> HashJoinBridge::tableOrFuture(
       !buildResult_.has_value() ||
       (!restoringSpillPartitionId_.has_value() &&
        restoringSpillShards_.empty()));
-
+  probeStarted_ = true;
   if (buildResult_.has_value()) {
     return buildResult_.value();
   }
@@ -286,6 +333,7 @@ bool HashJoinBridge::probeFinished() {
         !restoringSpillPartitionId_.has_value() &&
         restoringSpillShards_.empty());
     VELOX_CHECK_GT(numBuilders_, 0);
+    probeStarted_ = false;
 
     // NOTE: we are clearing the hash table as it has been fully processed and
     // not needed anymore. We'll wait for the HashBuild operator to build a new
@@ -374,6 +422,11 @@ uint64_t HashJoinMemoryReclaimer::reclaim(
     }
     return !hasReclaimedFromBuild;
   });
+  auto joinBridge = joinBridge_.lock();
+  VELOX_CHECK_NOT_NULL(joinBridge);
+  if (reclaimedBytes == 0 && joinBridge->canReclaim()) {
+    reclaimedBytes = joinBridge->reclaim();
+  }
   return reclaimedBytes;
 }
 

diff --git a/velox/exec/HashJoinBridge.h b/velox/exec/HashJoinBridge.h
@@ -26,6 +26,9 @@ namespace test {
 class HashJoinBridgeTestHelper;
 }
 
+using HashJoinTableSpillFunc =
+    std::function<SpillPartitionSet(std::shared_ptr<BaseHashTable>)>;
+
 /// Hands over a hash table from a multi-threaded build pipeline to a
 /// multi-threaded probe pipeline. This is owned by shared_ptr by all the build
 /// and probe Operator instances concerned. Corresponds to the Presto concept of
@@ -39,13 +42,18 @@ class HashJoinBridge : public JoinBridge {
   /// HashBuild operators to parallelize the restoring operation.
   void addBuilder();
 
+  bool canReclaim() const;
+
+  uint64_t reclaim();
+
   /// Invoked by the build operator to set the built hash table.
   /// 'spillPartitionSet' contains the spilled partitions while building
   /// 'table' which only applies if the disk spilling is enabled.
   void setHashTable(
       std::unique_ptr<BaseHashTable> table,
       SpillPartitionSet spillPartitionSet,
-      bool hasNullKeys);
+      bool hasNullKeys,
+      const HashJoinTableSpillFunc& tableSpillFunc);
 
   /// Invoked by the probe operator to set the spilled hash table while the
   /// probing. The function puts the spilled table partitions into
@@ -144,6 +152,11 @@ class HashJoinBridge : public JoinBridge {
   // memory and engages in recursive spilling.
   SpillPartitionSet spillPartitionSets_;
 
+  // A flag indicating if any probe operator has poked 'this' join bridge to
+  // attempt to get table. It is reset for each table partition processing.
+  std::atomic_bool probeStarted_;
+
+  HashJoinTableSpillFunc tableSpillFunc_{nullptr};
   friend test::HashJoinBridgeTestHelper;
 };
 
@@ -154,9 +167,10 @@ bool isLeftNullAwareJoinWithFilter(
 
 class HashJoinMemoryReclaimer final : public MemoryReclaimer {
  public:
-  static std::unique_ptr<memory::MemoryReclaimer> create() {
+  static std::unique_ptr<memory::MemoryReclaimer> create(
+      std::shared_ptr<HashJoinBridge> joinBridge) {
     return std::unique_ptr<memory::MemoryReclaimer>(
-        new HashJoinMemoryReclaimer());
+        new HashJoinMemoryReclaimer(joinBridge));
   }
 
   uint64_t reclaim(
@@ -166,7 +180,9 @@ class HashJoinMemoryReclaimer final : public MemoryReclaimer {
       memory::MemoryReclaimer::Stats& stats) final;
 
  private:
-  HashJoinMemoryReclaimer() : MemoryReclaimer() {}
+  HashJoinMemoryReclaimer(std::shared_ptr<HashJoinBridge> joinBridge)
+      : MemoryReclaimer(), joinBridge_(joinBridge) {}
+  std::weak_ptr<HashJoinBridge> joinBridge_;
 };
 
 /// Returns true if 'pool' is a hash build operator's memory pool. The check is

diff --git a/velox/exec/Task.cpp b/velox/exec/Task.cpp
@@ -483,7 +483,10 @@ velox::memory::MemoryPool* Task::getOrAddNodePool(
     return nodePools_[planNodeId];
   }
   childPools_.push_back(pool_->addAggregateChild(
-      fmt::format("node.{}", planNodeId), createNodeReclaimer(false)));
+      fmt::format("node.{}", planNodeId), createNodeReclaimer([&]() {
+        return exec::ParallelMemoryReclaimer::create(
+            queryCtx_->spillExecutor());
+      })));
   auto* nodePool = childPools_.back().get();
   nodePools_[planNodeId] = nodePool;
   return nodePool;
@@ -499,22 +502,24 @@ memory::MemoryPool* Task::getOrAddJoinNodePool(
     return nodePools_[nodeId];
   }
   childPools_.push_back(pool_->addAggregateChild(
-      fmt::format("node.{}", nodeId), createNodeReclaimer(true)));
+      fmt::format("node.{}", nodeId), createNodeReclaimer([&]() {
+        return HashJoinMemoryReclaimer::create(
+            getHashJoinBridgeLocked(splitGroupId, planNodeId));
+      })));
   auto* nodePool = childPools_.back().get();
   nodePools_[nodeId] = nodePool;
   return nodePool;
 }
 
 std::unique_ptr<memory::MemoryReclaimer> Task::createNodeReclaimer(
-    bool isHashJoinNode) const {
+    std::function<std::unique_ptr<memory::MemoryReclaimer>()> reclaimerFactory)
+    const {
   if (pool()->reclaimer() == nullptr) {
     return nullptr;
   }
   // Sets memory reclaimer for the parent node memory pool on the first child
   // operator construction which has set memory reclaimer.
-  return isHashJoinNode
-      ? HashJoinMemoryReclaimer::create()
-      : exec::ParallelMemoryReclaimer::create(queryCtx_->spillExecutor());
+  return reclaimerFactory();
 }
 
 std::unique_ptr<memory::MemoryReclaimer> Task::createExchangeClientReclaimer()

diff --git a/velox/exec/Task.h b/velox/exec/Task.h
@@ -766,7 +766,8 @@ class Task : public std::enable_shared_from_this<Task> {
   // customized instance for hash join plan node, otherwise creates a default
   // memory reclaimer.
   std::unique_ptr<memory::MemoryReclaimer> createNodeReclaimer(
-      bool isHashJoinNode) const;
+      std::function<std::unique_ptr<memory::MemoryReclaimer>()>
+          reclaimerFactory) const;
 
   // Creates a memory reclaimer instance for an exchange client if the task
   // memory pool has set memory reclaimer. We don't support to reclaim memory

diff --git a/velox/exec/tests/HashJoinBridgeTest.cpp b/velox/exec/tests/HashJoinBridgeTest.cpp
@@ -171,7 +171,8 @@ TEST_P(HashJoinBridgeTest, withoutSpill) {
     // Can't call any other APIs except addBuilder() before start a join bridge
     // first.
     VELOX_ASSERT_THROW(
-        joinBridge->setHashTable(createFakeHashTable(), {}, false), "");
+        joinBridge->setHashTable(createFakeHashTable(), {}, false, nullptr),
+        "");
     VELOX_ASSERT_THROW(joinBridge->setAntiJoinHasNullKeys(), "");
     VELOX_ASSERT_THROW(joinBridge->probeFinished(), "");
     VELOX_ASSERT_THROW(joinBridge->tableOrFuture(&futures[0]), "");
@@ -204,9 +205,10 @@ TEST_P(HashJoinBridgeTest, withoutSpill) {
     } else {
       auto table = createFakeHashTable();
       rawTable = table.get();
-      joinBridge->setHashTable(std::move(table), {}, false);
+      joinBridge->setHashTable(std::move(table), {}, false, nullptr);
       VELOX_ASSERT_THROW(
-          joinBridge->setHashTable(createFakeHashTable(), {}, false), "");
+          joinBridge->setHashTable(createFakeHashTable(), {}, false, nullptr),
+          "");
     }
     ASSERT_TRUE(helper.buildResult().has_value());
 
@@ -317,10 +319,13 @@ TEST_P(HashJoinBridgeTest, withSpill) {
         if (oneIn(2)) {
           spillPartitionIdSet = toSpillPartitionIdSet(spillPartitionSet);
           joinBridge->setHashTable(
-              createFakeHashTable(), std::move(spillPartitionSet), false);
+              createFakeHashTable(),
+              std::move(spillPartitionSet),
+              false,
+              nullptr);
         } else {
           spillByProber = !spillPartitionSet.empty();
-          joinBridge->setHashTable(createFakeHashTable(), {}, false);
+          joinBridge->setHashTable(createFakeHashTable(), {}, false, nullptr);
         }
         hasMoreSpill = numSpilledPartitions > numRestoredPartitions;
       }
@@ -446,9 +451,13 @@ TEST_P(HashJoinBridgeTest, multiThreading) {
                 auto spillPartitionSet =
                     makeFakeSpillPartitionSet(partitionBitOffset);
                 joinBridge->setHashTable(
-                    createFakeHashTable(), std::move(spillPartitionSet), false);
+                    createFakeHashTable(),
+                    std::move(spillPartitionSet),
+                    false,
+                    nullptr);
               } else {
-                joinBridge->setHashTable(createFakeHashTable(), {}, false);
+                joinBridge->setHashTable(
+                    createFakeHashTable(), {}, false, nullptr);
               }
             }
             for (auto& promise : promises) {