SERVER-4507 Group stages now take advantage of sorted input sequences.

1 files changed, 26 insertions, 0 deletions

diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index f9927de1e6f..4e63247bef4 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -128,6 +128,32 @@ BSONObj DocumentSourceLookUp::queryForInput(const Document& input) const {
     return query.obj();
 }
 
+BSONObjSet DocumentSourceLookUp::getOutputSorts() {
+    BSONObjSet out;
+
+    BSONObjSet inputSort = pSource->getOutputSorts();
+    std::string asPath = _as.getPath(false);
+
+    for (auto&& sortObj : inputSort) {
+        // Truncate each sortObj at the '_as' path.
+        BSONObjBuilder outputSort;
+
+        for (BSONElement fieldSort : sortObj) {
+            if (fieldSort.fieldNameStringData() == asPath) {
+                break;
+            }
+            outputSort.append(fieldSort);
+        }
+
+        BSONObj outSortObj = outputSort.obj();
+        if (!outSortObj.isEmpty()) {
+            out.insert(outSortObj);
+        }
+    }
+
+    return out;
+}
+
 boost::optional<Document> DocumentSourceLookUp::unwindResult() {
     const boost::optional<FieldPath> indexPath(_unwindSrc->indexPath());