SERVER-49829 - Implement spilling for Top K sort in SBE.

2 files changed, 18 insertions, 2 deletions

diff --git a/src/mongo/db/sorter/sorter.cpp b/src/mongo/db/sorter/sorter.cpp
index 0b3030e2209..b83fe1d6619 100644
--- a/src/mongo/db/sorter/sorter.cpp
+++ b/src/mongo/db/sorter/sorter.cpp
@@ -142,6 +142,8 @@ public:
     template <typename Container>
     InMemIterator(const Container& input) : _data(input.begin(), input.end()) {}
 
+    InMemIterator(std::deque<Data> input) : _data(std::move(input)) {}
+
     void openSource() {}
     void closeSource() {}
 
@@ -590,12 +592,24 @@ public:
             spill();
     }
 
+    void emplace(Key&& key, Value&& val) override {
+        invariant(!_done);
+
+        _memUsed += key.memUsageForSorter();
+        _memUsed += val.memUsageForSorter();
+
+        _data.emplace_back(std::move(key), std::move(val));
+
+        if (_memUsed > _opts.maxMemoryUsageBytes)
+            spill();
+    }
+
     Iterator* done() {
         invariant(!_done);
 
         if (this->_iters.empty()) {
             sort();
-            return new InMemIterator<Key, Value>(_data);
+            return new InMemIterator<Key, Value>(std::move(_data));
         }
 
         spill();
diff --git a/src/mongo/db/sorter/sorter.h b/src/mongo/db/sorter/sorter.h
index a3537c46a7d..1debcf5e318 100644
--- a/src/mongo/db/sorter/sorter.h
+++ b/src/mongo/db/sorter/sorter.h
@@ -240,7 +240,9 @@ public:
                                           const Settings& settings = Settings());
 
     virtual void add(const Key&, const Value&) = 0;
-
+    virtual void emplace(Key&& k, Value&& v) {
+        add(k, v);
+    }
     /**
      * Cannot add more data after calling done().
      *