add orderby ut

facebookincubator · Nov 30, 2023 · 551b3b5 · 551b3b5
1 parent 2aec00c
commit 551b3b5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/velox/exec/SortBuffer.cpp b/velox/exec/SortBuffer.cpp
@@ -133,11 +133,7 @@ void SortBuffer::noMoreInput() {
     // for now.
     spill();
 
-    // Finish spill, and we shouldn't get any rows from non-spilled partition as
-    // there is only one hash partition for SortBuffer.
-    VELOX_CHECK_NULL(spillMerger_);
-    auto spillPartition = spiller_->finishSpill();
-    spillMerger_ = spillPartition.createOrderedReader(pool());
+    finishSpill();
   }
 
   // Releases the unused memory reservation after procesing input.
@@ -222,14 +218,9 @@ void SortBuffer::spillOutput() {
   auto spillRows = std::vector<char*>(
       sortedRows_.begin() + numOutputRows_, sortedRows_.end());
   spiller_->spill(std::move(spillRows));
-  // Finish spill, and we shouldn't get any rows from non-spilled partition as
-  // there is only one hash partition for SortBuffer.
-  VELOX_CHECK_NULL(spillMerger_);
-  auto spillPartition = spiller_->finishSpill();
-  spillMerger_ = spillPartition.createOrderedReader(pool());
-
   data_->clear();
   sortedRows_.clear();
+  finishSpill();
 }
 
 std::optional<uint64_t> SortBuffer::estimateOutputRowSize() const {
@@ -404,4 +395,10 @@ void SortBuffer::getOutputWithSpill() {
   numOutputRows_ += output_->size();
 }
 
+void SortBuffer::finishSpill() {
+  VELOX_CHECK_NULL(spillMerger_);
+  auto spillPartition = spiller_->finishSpill();
+  spillMerger_ = spillPartition.createOrderedReader(pool());
+}
+
 } // namespace facebook::velox::exec
diff --git a/velox/exec/SortBuffer.h b/velox/exec/SortBuffer.h
@@ -85,6 +85,10 @@ class SortBuffer {
   void getOutputWithoutSpill();
   void getOutputWithSpill();
 
+  // Finish spill, and we shouldn't get any rows from non-spilled partition as
+  // there is only one hash partition for SortBuffer.
+  void finishSpill();
+
   const RowTypePtr input_;
   const std::vector<CompareFlags> sortCompareFlags_;
   velox::memory::MemoryPool* const pool_;