SERVER-51303 Fix lookup match absorbtion optimization for $typer4.4.2-rc1r4.4.2

3 files changed, 237 insertions, 52 deletions

diff --git a/jstests/aggregation/sources/lookup/lookup_absorb_match.js b/jstests/aggregation/sources/lookup/lookup_absorb_match.js
index c39b8defff2..6b21b18a2f3 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 @@ let testDB = db.getSiblingDB("lookup_absorb_match");
 testDB.dropDatabase();
 
 let locations = testDB.getCollection("locations");
-assert.commandWorked(locations.insert({_id: "doghouse", coordinates: [25.0, 60.0]}));
-assert.commandWorked(locations.insert({_id: "bullpen", coordinates: [-25.0, -60.0]}));
+assert.commandWorked(locations.insert({
+    _id: "doghouse",
+    coordinates: [25.0, 60.0],
+    extra: {breeds: ["terrier", "dachshund", "bulldog"]}
+}));
+assert.commandWorked(locations.insert({
+    _id: "bullpen",
+    coordinates: [-25.0, -60.0],
+    extra: {breeds: "Scottish Highland", feeling: "bullish"}
+}));
 
 let animals = testDB.getCollection("animals");
 assert.commandWorked(animals.insert({_id: "dog", locationId: "doghouse"}));
@@ -25,32 +33,33 @@ assert.commandWorked(animals.insert({_id: "bull", locationId: "bullpen"}));
 let result = testDB.animals
                      .aggregate([
                          {
-                           $lookup: {
-                               from: "locations",
-                               localField: "locationId",
-                               foreignField: "_id",
-                               as: "location"
-                           }
+                             $lookup: {
+                                 from: "locations",
+                                 localField: "locationId",
+                                 foreignField: "_id",
+                                 as: "location"
+                             }
                          },
                          {$unwind: "$location"},
                          {
-                           $match: {
-                               "location.coordinates": {
-                                   $geoWithin: {
-                                       $geometry: {
-                                           type: "MultiPolygon",
-                                           coordinates: [[[
-                                               [20.0, 70.0],
-                                               [30.0, 70.0],
-                                               [30.0, 50.0],
-                                               [20.0, 50.0],
-                                               [20.0, 70.0]
-                                           ]]]
-                                       }
-                                   }
-                               }
-                           }
-                         }
+                             $match: {
+                                 "location.coordinates": {
+                                     $geoWithin: {
+                                         $geometry: {
+                                             type: "MultiPolygon",
+                                             coordinates: [[[
+                                                 [20.0, 70.0],
+                                                 [30.0, 70.0],
+                                                 [30.0, 50.0],
+                                                 [20.0, 50.0],
+                                                 [20.0, 70.0]
+                                             ]]]
+                                         }
+                                     }
+                                 }
+                             }
+                         },
+                         {$project: {"location.extra": false}}
                      ])
                      .toArray();
 let expected =
@@ -61,35 +70,194 @@ assert.eq(result, expected);
 result = testDB.animals
                  .aggregate([
                      {
-                       $lookup: {
-                           from: "locations",
-                           localField: "locationId",
-                           foreignField: "_id",
-                           as: "location"
-                       }
+                         $lookup: {
+                             from: "locations",
+                             localField: "locationId",
+                             foreignField: "_id",
+                             as: "location"
+                         }
                      },
                      {$unwind: "$location"},
                      {
-                       $match: {
-                           "location.coordinates": {
-                               $geoIntersects: {
-                                   $geometry: {
-                                       type: "MultiPolygon",
-                                       coordinates: [[[
-                                           [-20.0, -70.0],
-                                           [-30.0, -70.0],
-                                           [-30.0, -50.0],
-                                           [-20.0, -50.0],
-                                           [-20.0, -70.0]
-                                       ]]]
-                                   }
-                               }
-                           }
-                       }
-                     }
+                         $match: {
+                             "location.coordinates": {
+                                 $geoIntersects: {
+                                     $geometry: {
+                                         type: "MultiPolygon",
+                                         coordinates: [[[
+                                             [-20.0, -70.0],
+                                             [-30.0, -70.0],
+                                             [-30.0, -50.0],
+                                             [-20.0, -50.0],
+                                             [-20.0, -70.0]
+                                         ]]]
+                                     }
+                                 }
+                             }
+                         }
+                     },
+                     {$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 $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 7da9d5a3f3c..f5dc049df46 100644
--- a/src/mongo/db/pipeline/document_source_match.cpp
+++ b/src/mongo/db/pipeline/document_source_match.cpp
@@ -446,8 +446,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);
         } 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 2fc38946417..44d48015332 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -519,6 +519,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: []}}, "