SERVER-23469 Match on unwind's indexPath is no longer optimized.

2 files changed, 20 insertions, 2 deletions

diff --git a/src/mongo/db/pipeline/document_source_unwind.cpp b/src/mongo/db/pipeline/document_source_unwind.cpp
index 506627cc4dd..bba51298389 100644
--- a/src/mongo/db/pipeline/document_source_unwind.cpp
+++ b/src/mongo/db/pipeline/document_source_unwind.cpp
@@ -231,8 +231,13 @@ Pipeline::SourceContainer::iterator DocumentSourceUnwind::optimizeAt(
 
     if (auto nextMatch = dynamic_cast<DocumentSourceMatch*>((*std::next(itr)).get())) {
         const bool includeDollarPrefix = false;
-        std::set<std::string> field = {_unwindPath.getPath(includeDollarPrefix)};
-        auto splitMatch = nextMatch->splitSourceBy(field);
+        std::set<std::string> fields = {_unwindPath.getPath(includeDollarPrefix)};
+
+        if (_indexPath) {
+            fields.insert((*_indexPath).getPath(false));
+        }
+
+        auto splitMatch = nextMatch->splitSourceBy(fields);
 
         invariant(splitMatch.first || splitMatch.second);
 
diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp
index e5afb99fb86..1b583ccf70c 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -315,6 +315,18 @@ class MatchShouldNotOptimizeWithElemMatch : public Base {
     }
 };
 
+class MatchShouldNotOptimizeWhenMatchingOnIndexField : public Base {
+    string inputPipeJson() {
+        return "[{$unwind: {path: '$a', includeArrayIndex: 'foo'}}, "
+               " {$match: {foo: 0, b: 1}}]";
+    }
+    string outputPipeJson() {
+        return "[{$match: {b: {$eq: 1}}}, "
+               " {$unwind: {path: '$a', includeArrayIndex: 'foo'}}, "
+               " {$match: {foo: {$eq: 0}}}]";
+    }
+};
+
 class MatchWithNorOnlySplitsIndependentChildren : public Base {
     string inputPipeJson() {
         return "[{$unwind: {path: '$a'}}, "
@@ -758,6 +770,7 @@ public:
         add<Optimizations::Local::LookupShouldNotCoalesceWithUnwindNotOnAs>();
         add<Optimizations::Local::MatchShouldDuplicateItselfBeforeRedact>();
         add<Optimizations::Local::MatchShouldSwapWithUnwind>();
+        add<Optimizations::Local::MatchShouldNotOptimizeWhenMatchingOnIndexField>();
         add<Optimizations::Local::MatchOnPrefixShouldNotSwapOnUnwind>();
         add<Optimizations::Local::MatchShouldNotOptimizeWithElemMatch>();
         add<Optimizations::Local::MatchWithNorOnlySplitsIndependentChildren>();