SERVER-51303 Fix lookup match absorbtion optimization for $typer3.6.21-rc2

3 files changed, 200 insertions, 8 deletions

diff --git a/jstests/aggregation/sources/lookup/lookup_absorb_match.js b/jstests/aggregation/sources/lookup/lookup_absorb_match.js
index 1a6aea31b16..8f65eb22e56 100644
--- a/jstests/aggregation/sources/lookup/lookup_absorb_match.js
+++ b/jstests/aggregation/sources/lookup/lookup_absorb_match.js
@@ -1,6 +1,6 @@
 /**
- * Tests that a $match with a geo expression still returns the correct results if it has been
- * absorbed by a $lookup.
+ * Tests that $match stages with a variety of expressions still return the correct results if they
+ * have been absorbed by $lookup stages.
  *
  * Accessed collections cannot be implicitly sharded because you cannot $lookup into a sharded
  * collection.
@@ -13,8 +13,16 @@
     testDB.dropDatabase();
 
     let locations = testDB.getCollection("locations");
-    assert.writeOK(locations.insert({_id: "doghouse", coordinates: [25.0, 60.0]}));
-    assert.writeOK(locations.insert({_id: "bullpen", coordinates: [-25.0, -60.0]}));
+    assert.writeOK(locations.insert({
+        _id: "doghouse",
+        coordinates: [25.0, 60.0],
+        extra: {breeds: ["terrier", "dachshund", "bulldog"]}
+    }));
+    assert.writeOK(locations.insert({
+        _id: "bullpen",
+        coordinates: [-25.0, -60.0],
+        extra: {breeds: "Scottish Highland", feeling: "bullish"}
+    }));
 
     let animals = testDB.getCollection("animals");
     assert.writeOK(animals.insert({_id: "dog", locationId: "doghouse"}));
@@ -50,7 +58,8 @@
                                    }
                                }
                            }
-                         }
+                         },
+                         {$project: {"location.extra": false}}
                      ])
                      .toArray();
     let expected = [{
@@ -89,7 +98,8 @@
                                }
                            }
                        }
-                     }
+                     },
+                     {$project: {"location.extra": false}}
                  ])
                  .toArray();
     expected = [{
@@ -98,4 +108,169 @@
         location: {_id: "bullpen", coordinates: [-25.0, -60.0]}
     }];
     assert.eq(result, expected);
+
+    // Test that a $match with $type works as expected when absorbed by a $lookup.
+    result = testDB.animals
+                 .aggregate([
+                     {
+                       $lookup: {
+                           from: "locations",
+                           localField: "locationId",
+                           foreignField: "_id",
+                           as: "location"
+                       }
+                     },
+                     {$unwind: "$location"},
+                     {$match: {"location.extra.breeds": {$type: "array"}}},
+                     {$project: {"location.extra": false}}
+                 ])
+                 .toArray();
+    expected = [{
+        _id: "dog",
+        locationId: "doghouse",
+        location: {_id: "doghouse", coordinates: [25.0, 60.0]}
+    }];
+    assert.eq(result, expected);
+
+    // Test that a $match with $jsonSchema works as expected although ineligable for absorbtion by a
+    // $lookup.
+    result = testDB.animals
+                 .aggregate([
+                     {
+                       $lookup: {
+                           from: "locations",
+                           localField: "locationId",
+                           foreignField: "_id",
+                           as: "location"
+                       }
+                     },
+                     {$unwind: "$location"},
+                     {
+                       $match: {
+                           $jsonSchema: {
+                               properties: {
+                                   location: {
+                                       properties:
+                                           {extra: {properties: {breeds: {type: "string"}}}}
+                                   }
+                               }
+                           }
+                       }
+                     },
+                     {$project: {"location.extra": false}}
+                 ])
+                 .toArray();
+    expected = [{
+        _id: "bull",
+        locationId: "bullpen",
+        location: {_id: "bullpen", coordinates: [-25.0, -60.0]}
+    }];
+    assert.eq(result, expected);
+
+    // Test that a more complex $match with $jsonSchema works as expected although ineligable for
+    // absorbtion by a $lookup.
+    result =
+        testDB.animals
+            .aggregate([
+                {
+                  $lookup: {
+                      from: "locations",
+                      localField: "locationId",
+                      foreignField: "_id",
+                      as: "location"
+                  }
+                },
+                {$unwind: "$location"},
+                {
+                  $match: {
+                      $jsonSchema:
+                          {properties: {location: {properties: {extra: {minProperties: 2}}}}}
+                  }
+                },
+                {$project: {"location.extra": false}}
+            ])
+            .toArray();
+    expected = [{
+        _id: "bull",
+        locationId: "bullpen",
+        location: {_id: "bullpen", coordinates: [-25.0, -60.0]}
+    }];
+    assert.eq(result, expected);
+
+    // Test that a $match with $alwaysTrue works as expected although ineligable for absorbtion by a
+    // $lookup.
+    result = testDB.animals
+                 .aggregate([
+                     {
+                       $lookup: {
+                           from: "locations",
+                           localField: "locationId",
+                           foreignField: "_id",
+                           as: "location"
+                       }
+                     },
+                     {$unwind: "$location"},
+                     {$match: {$alwaysTrue: 1}},
+                     {$project: {"location.extra": false}},
+                     {$sort: {_id: 1}}
+                 ])
+                 .toArray();
+    expected = [
+        {
+          _id: "bull",
+          locationId: "bullpen",
+          location: {_id: "bullpen", coordinates: [-25.0, -60.0]}
+        },
+        {
+          _id: "dog",
+          locationId: "doghouse",
+          location: {_id: "doghouse", coordinates: [25.0, 60.0]}
+        }
+    ];
+    assert.eq(result, expected);
+
+    // Test that a $match with $alwaysFalse works as expected although ineligable for absorbtion by
+    // a
+    // $lookup.
+    result = testDB.animals
+                 .aggregate([
+                     {
+                       $lookup: {
+                           from: "locations",
+                           localField: "locationId",
+                           foreignField: "_id",
+                           as: "location"
+                       }
+                     },
+                     {$unwind: "$location"},
+                     {$match: {$alwaysFalse: 1}},
+                     {$project: {"location.extra": false}},
+                 ])
+                 .toArray();
+    expected = [];
+    assert.eq(result, expected);
+
+    // Test that a $match with $expr works as expected although ineligable for absorbtion by a
+    // $lookup.
+    result = testDB.animals
+                 .aggregate([
+                     {
+                       $lookup: {
+                           from: "locations",
+                           localField: "locationId",
+                           foreignField: "_id",
+                           as: "location"
+                       }
+                     },
+                     {$unwind: "$location"},
+                     {$match: {$expr: {$in: [25.0, "$location.coordinates"]}}},
+                     {$project: {"location.extra": false}},
+                 ])
+                 .toArray();
+    expected = [{
+        _id: "dog",
+        locationId: "doghouse",
+        location: {_id: "doghouse", coordinates: [25.0, 60.0]}
+    }];
+    assert.eq(result, expected);
 }());
diff --git a/src/mongo/db/pipeline/document_source_match.cpp b/src/mongo/db/pipeline/document_source_match.cpp
index 29ec41f2bbd..b1a2c942372 100644
--- a/src/mongo/db/pipeline/document_source_match.cpp
+++ b/src/mongo/db/pipeline/document_source_match.cpp
@@ -447,8 +447,7 @@ boost::intrusive_ptr<DocumentSourceMatch> DocumentSourceMatch::descendMatchOnPat
         invariant(expression::isPathPrefixOf(descendOn, leafPath));
 
         auto newPath = leafPath.substr(descendOn.size() + 1);
-        if (node->getCategory() == MatchExpression::MatchCategory::kLeaf &&
-            node->matchType() != MatchExpression::TYPE_OPERATOR) {
+        if (node->getCategory() == MatchExpression::MatchCategory::kLeaf) {
             auto leafNode = static_cast<LeafMatchExpression*>(node);
             leafNode->setPath(newPath).transitional_ignore();
         } else if (node->getCategory() == MatchExpression::MatchCategory::kArrayMatching) {
diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp
index 80c4a6ff89a..69e7c10602f 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -445,6 +445,24 @@ TEST(PipelineOptimizationTest, LookupShouldAbsorbUnwindMatch) {
     assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe);
 }
 
+TEST(PipelineOptimizationTest, LookupShouldAbsorbUnwindAndTypeMatch) {
+    string inputPipe =
+        "[{$lookup: {from: 'lookupColl', as: 'asField', localField: 'y', foreignField: "
+        "'z'}}, "
+        "{$unwind: '$asField'}, "
+        "{$match: {'asField.subfield': {$type: [2]}}}]";
+    string outputPipe =
+        "[{$lookup: {from: 'lookupColl', as: 'asField', localField: 'y', foreignField: 'z', "
+        "            unwinding: {preserveNullAndEmptyArrays: false}, "
+        "            matching: {subfield: {$type: [2]}}}}]";
+    string serializedPipe =
+        "[{$lookup: {from: 'lookupColl', as: 'asField', localField: 'y', foreignField: "
+        "'z'}}, "
+        "{$unwind: {path: '$asField'}}, "
+        "{$match: {'asField.subfield': {$type: [2]}}}]";
+    assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe);
+}
+
 TEST(PipelineOptimizationTest, LookupWithPipelineSyntaxShouldAbsorbUnwindMatch) {
     string inputPipe =
         "[{$lookup: {from: 'lookupColl', as: 'asField', pipeline: []}}, "