Refactor spilling for RowNumber

velox/exec/RowNumber.cpp

@@ -96,42 +96,21 @@ void RowNumber::addInput(RowVectorPtr input) {
   input_ = std::move(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;
-  }
-
-  const auto numInput = input_->size();
-  SelectivityVector rows(numInput);
-
-  VELOX_CHECK(spillConfig_.has_value());
-  table_->prepareForGroupProbe(
-      *lookup_, input_, rows, spillConfig_->startPartitionBit);
-  table_->groupProbe(*lookup_, spillConfig_->startPartitionBit);
-
-  // Initialize new partitions with zeros.
-  for (auto i : lookup_->newGroups) {
-    setNumRows(lookup_->hits[i], 0);
-  }
-}
-
 void RowNumber::noMoreInput() {
   Operator::noMoreInput();
 
   if (inputSpiller_ != nullptr) {
-    inputSpiller_->finishSpill(spillInputPartitionSet_);
-    inputSpiller_.reset();
-    removeEmptyPartitions(spillInputPartitionSet_);
-    restoreNextSpillPartition();
+    finishSpillInputAndRestoreNext();
   }
 }
 
+void RowNumber::finishSpillInputAndRestoreNext() {
+  inputSpiller_->finishSpill(spillInputPartitionSet_);
+  inputSpiller_.reset();
+  removeEmptyPartitions(spillInputPartitionSet_);
+  restoreNextSpillPartition();
+}
+
 void RowNumber::restoreNextSpillPartition() {
   if (spillInputPartitionSet_.empty()) {
     return;
@@ -181,10 +160,11 @@ void RowNumber::restoreNextSpillPartition() {
 
   spillInputPartitionSet_.erase(it);
 
-  spillInputReader_->nextBatch(input_);
-  VELOX_CHECK_NOT_NULL(input_);
+  RowVectorPtr unspilledInput;
+  spillInputReader_->nextBatch(unspilledInput);
+  VELOX_CHECK_NOT_NULL(unspilledInput);
   // NOTE: spillInputReader_ will at least produce one batch output.
-  addSpillInput();
+  addInput(std::move(unspilledInput));
 }
 
 void RowNumber::ensureInputFits(const RowVectorPtr& input) {
@@ -339,19 +319,17 @@ RowVectorPtr RowNumber::getOutput() {
     output = fillOutput(numInput, nullptr);
   }
 
+  input_ = nullptr;
   if (spillInputReader_ != nullptr) {
-    if (spillInputReader_->nextBatch(input_)) {
-      addSpillInput();
+    RowVectorPtr unspilledInput;
+    if (spillInputReader_->nextBatch(unspilledInput)) {
+      addInput(std::move(unspilledInput));
     } else {
-      input_ = nullptr;
       spillInputReader_ = nullptr;
       table_->clear();
       restoreNextSpillPartition();
     }
-  } else {
-    input_ = nullptr;
   }
-
   return output;
 }
 
@@ -522,22 +500,17 @@ void RowNumber::spillInput(
 }
 
 void RowNumber::recursiveSpillInput() {
-  RowVectorPtr input;
-  while (spillInputReader_->nextBatch(input)) {
-    spillInput(input, pool());
+  RowVectorPtr unspilledInput;
+  while (spillInputReader_->nextBatch(unspilledInput)) {
+    spillInput(unspilledInput, pool());
 
     if (operatorCtx_->driver()->shouldYield()) {
       yield_ = true;
       return;
     }
   }
 
-  inputSpiller_->finishSpill(spillInputPartitionSet_);
-  inputSpiller_.reset();
-  spillInputReader_ = nullptr;
-
-  removeEmptyPartitions(spillInputPartitionSet_);
-  restoreNextSpillPartition();
+  finishSpillInputAndRestoreNext();
 }
 
 void RowNumber::setSpillPartitionBits(

velox/exec/RowNumber.h

@@ -64,8 +64,6 @@ class RowNumber : public Operator {
 
   void spill();
 
-  void addSpillInput();
-
   void restoreNextSpillPartition();
 
   SpillPartitionNumSet spillHashTable();
@@ -78,6 +76,10 @@ class RowNumber : public Operator {
 
   FlatVector<int64_t>& getOrCreateRowNumberVector(vector_size_t size);
 
+  // Finishes the current input spilling and restore the next processing
+  // partition.
+  void finishSpillInputAndRestoreNext();
+
   // Used by recursive spill processing to read the spilled input data from the
   // previous spill run through 'spillInputReader_' and then spill them back
   // into a number of sub-partitions. After that, the function restores one of