Release memory when reservation triggers spilling (facebookincubator#…

…11063)

Summary:
Spillable operator might trigger spill when it reserve memory from the arbitrator. If that happens, we shall release the reserved memory if the reserved memory is no longer needed after spilling.

Pull Request resolved: facebookincubator#11063

Reviewed By: xiaoxmeng

Differential Revision: D63207407

Pulled By: tanjialiang

fbshipit-source-id: 8a956a25be69ca6fccc55e929a62034d4b9e8c5b

velox/exec/GroupingSet.cpp

@@ -928,6 +928,14 @@ void GroupingSet::ensureOutputFits() {
   {
     memory::ReclaimableSectionGuard guard(nonReclaimableSection_);
     if (pool_.maybeReserve(outputBufferSizeToReserve)) {
+      if (hasSpilled()) {
+        // If reservation triggers spilling on the 'GroupingSet' itself, we will
+        // no longer need the reserved memory for output processing as the
+        // output processing will be conducted from unspilled data through
+        // 'getOutputWithSpill()', and it does not require this amount of memory
+        // to process.
+        pool_.release();
+      }
       return;
     }
   }

velox/exec/HashBuild.cpp

@@ -495,15 +495,21 @@ void HashBuild::ensureInputFits(RowVectorPtr& input) {
 
   {
     Operator::ReclaimableSectionGuard guard(this);
-    if (!pool()->maybeReserve(targetIncrementBytes)) {
-      LOG(WARNING) << "Failed to reserve "
-                   << succinctBytes(targetIncrementBytes) << " for memory pool "
-                   << pool()->name()
-                   << ", usage: " << succinctBytes(pool()->usedBytes())
-                   << ", reservation: "
-                   << succinctBytes(pool()->reservedBytes());
+    if (pool()->maybeReserve(targetIncrementBytes)) {
+      // If above reservation triggers the spilling of 'HashBuild' operator
+      // itself, we will no longer need the reserved memory for building hash
+      // table as the table is spilled, and the input will be directly spilled,
+      // too.
+      if (spiller_->isAllSpilled()) {
+        pool()->release();
+      }
+      return;
     }
   }
+  LOG(WARNING) << "Failed to reserve " << succinctBytes(targetIncrementBytes)
+               << " for memory pool " << pool()->name()
+               << ", usage: " << succinctBytes(pool()->usedBytes())
+               << ", reservation: " << succinctBytes(pool()->reservedBytes());
 }
 
 void HashBuild::spillInput(const RowVectorPtr& input) {
@@ -806,6 +812,12 @@ void HashBuild::ensureTableFits(uint64_t numRows) {
   {
     Operator::ReclaimableSectionGuard guard(this);
     if (pool()->maybeReserve(memoryBytesToReserve)) {
+      // If reservation triggers the spilling of 'HashBuild' operator itself, we
+      // will no longer need the reserved memory for building hash table as the
+      // table is spilled.
+      if (spiller_->isAllSpilled()) {
+        pool()->release();
+      }
       return;
     }
   }

velox/exec/RowNumber.cpp

@@ -98,8 +98,10 @@ void RowNumber::addInput(RowVectorPtr input) {
 
 void RowNumber::addSpillInput() {
   VELOX_CHECK_NOT_NULL(input_);
+  VELOX_CHECK_NULL(inputSpiller_);
   ensureInputFits(input_);
   if (input_ == nullptr) {
+    VELOX_CHECK_NOT_NULL(inputSpiller_);
     // Memory arbitration might be triggered by ensureInputFits() which will
     // spill 'input_'.
     return;
@@ -124,6 +126,7 @@ void RowNumber::noMoreInput() {
 
   if (inputSpiller_ != nullptr) {
     inputSpiller_->finishSpill(spillInputPartitionSet_);
+    inputSpiller_.reset();
     removeEmptyPartitions(spillInputPartitionSet_);
     restoreNextSpillPartition();
   }
@@ -185,7 +188,7 @@ void RowNumber::restoreNextSpillPartition() {
 }
 
 void RowNumber::ensureInputFits(const RowVectorPtr& input) {
-  if (!spillEnabled()) {
+  if (!spillEnabled() || inputSpiller_ != nullptr) {
     // Spilling is disabled.
     return;
   }
@@ -242,6 +245,12 @@ void RowNumber::ensureInputFits(const RowVectorPtr& input) {
   {
     Operator::ReclaimableSectionGuard guard(this);
     if (pool()->maybeReserve(targetIncrementBytes)) {
+      // If reservation triggers the spilling of 'RowNumber' operator itself, we
+      // will no longer need the reserved memory for building hash table as the
+      // table is spilled.
+      if (inputSpiller_ != nullptr) {
+        pool()->release();
+      }
       return;
     }
   }
@@ -524,6 +533,7 @@ void RowNumber::recursiveSpillInput() {
   }
 
   inputSpiller_->finishSpill(spillInputPartitionSet_);
+  inputSpiller_.reset();
   spillInputReader_ = nullptr;
 
   removeEmptyPartitions(spillInputPartitionSet_);