SERVER-55601 Improve performance of Queries.CoveredBlockingSort benchmark in SBE

2 files changed, 25 insertions, 1 deletions

diff --git a/src/mongo/db/sorter/sorter.cpp b/src/mongo/db/sorter/sorter.cpp
index 4a0aa54156f..330b3b6cc1f 100644
--- a/src/mongo/db/sorter/sorter.cpp
+++ b/src/mongo/db/sorter/sorter.cpp
@@ -159,6 +159,8 @@ public:
     template <typename Container>
     InMemIterator(const Container& input) : _data(input.begin(), input.end()) {}
 
+    InMemIterator(std::deque<Data> data) : _data(std::move(data)) {}
+
     void openSource() {}
     void closeSource() {}
 
@@ -622,6 +624,9 @@ public:
 
         if (this->_iters.empty()) {
             sort();
+            if (this->_opts.moveSortedDataIntoIterator) {
+                return new InMemIterator<Key, Value>(std::move(_data));
+            }
             return new InMemIterator<Key, Value>(_data);
         }
 
@@ -717,6 +722,9 @@ public:
 
     Iterator* done() {
         if (_haveData) {
+            if (this->_opts.moveSortedDataIntoIterator) {
+                return new InMemIterator<Key, Value>(std::move(_best));
+            }
             return new InMemIterator<Key, Value>(_best);
         } else {
             return new InMemIterator<Key, Value>();
@@ -823,6 +831,9 @@ public:
     Iterator* done() {
         if (this->_iters.empty()) {
             sort();
+            if (this->_opts.moveSortedDataIntoIterator) {
+                return new InMemIterator<Key, Value>(std::move(_data));
+            }
             return new InMemIterator<Key, Value>(_data);
         }
 
diff --git a/src/mongo/db/sorter/sorter.h b/src/mongo/db/sorter/sorter.h
index 782d0389f57..423638d7a8a 100644
--- a/src/mongo/db/sorter/sorter.h
+++ b/src/mongo/db/sorter/sorter.h
@@ -114,7 +114,15 @@ struct SortOptions {
     // extSortAllowed is true.
     std::string tempDir;
 
-    SortOptions() : limit(0), maxMemoryUsageBytes(64 * 1024 * 1024), extSortAllowed(false) {}
+    // If set to true and sorted data fits into memory, sorted data will be moved into iterator
+    // instead of copying.
+    bool moveSortedDataIntoIterator;
+
+    SortOptions()
+        : limit(0),
+          maxMemoryUsageBytes(64 * 1024 * 1024),
+          extSortAllowed(false),
+          moveSortedDataIntoIterator(false) {}
 
     // Fluent API to support expressions like SortOptions().Limit(1000).ExtSortAllowed(true)
 
@@ -142,6 +150,11 @@ struct SortOptions {
         dbName = std::move(newDbName);
         return *this;
     }
+
+    SortOptions& MoveSortedDataIntoIterator(bool newMoveSortedDataIntoIterator = true) {
+        moveSortedDataIntoIterator = newMoveSortedDataIntoIterator;
+        return *this;
+    }
 };
 
 /**