SERVER-114508 Allow localField override in JOO (#44495)

GitOrigin-RevId: de696c3697c672467836ae046e6c70ef1e7e4432

jstests/noPassthrough/query/joins/local_field_override.js

@@ -0,0 +1,77 @@
+/**
+ * Verifies that we correcly process overrding local fields by foreign documents.
+ * @tags: [
+ *   requires_fcv_83,
+ * ]
+ */
+
+import {assertArrayEq} from "jstests/aggregation/extras/utils.js";
+
+const docs = [
+    {_id: "first", a: 1, b: 1},
+    {_id: "second", a: 1, b: 2},
+];
+
+const config = {
+    setParameter: {
+        internalEnableJoinOptimization: true,
+    },
+};
+
+const conn = MongoRunner.runMongod(config);
+
+const db = conn.getDB(jsTestName());
+
+db.coll.drop();
+assert.commandWorked(db.coll.insertMany(docs));
+
+const pipeline = [
+    {$lookup: {from: "coll", localField: "_id", foreignField: "_id", as: "_id"}},
+    {$unwind: "$_id"},
+    {$lookup: {from: "coll", localField: "a", foreignField: "b", as: "a"}},
+    {$unwind: "$a"},
+    {$lookup: {from: "coll", localField: "b", foreignField: "b", as: "b"}},
+    {$unwind: "$b"},
+];
+
+const actual = db.coll.aggregate(pipeline).toArray();
+MongoRunner.stopMongod(conn);
+
+const expected = [
+    {
+        "_id": {
+            "_id": "first",
+            "a": 1,
+            "b": 1,
+        },
+        "a": {
+            "_id": "first",
+            "a": 1,
+            "b": 1,
+        },
+        "b": {
+            "_id": "first",
+            "a": 1,
+            "b": 1,
+        },
+    },
+    {
+        "_id": {
+            "_id": "second",
+            "a": 1,
+            "b": 2,
+        },
+        "a": {
+            "_id": "first",
+            "a": 1,
+            "b": 1,
+        },
+        "b": {
+            "_id": "second",
+            "a": 1,
+            "b": 2,
+        },
+    },
+];
+
+assertArrayEq({actual, expected});

src/mongo/db/query/compiler/optimizer/join/agg_join_model.cpp

@@ -192,13 +192,13 @@ StatusWith<AggJoinModel> AggJoinModel::constructJoinModel(const Pipeline& pipeli
                 return Status(ErrorCodes::BadValue, "Graph is too big: too many nodes");
             }
 
-            pathResolver.addNode(*foreignNodeId, lookup->getAsField());
-
             if (lookup->hasLocalFieldForeignFieldJoin()) {
                 // The order of resolving the paths are important here: localPathId shouln't be
                 // resolved to the foreign collection even if it is prefixed by the foreign
                 // collection's embedPath.
                 auto localPathId = pathResolver.resolve(*lookup->getLocalField());
+
+                pathResolver.addNode(*foreignNodeId, lookup->getAsField());
                 auto foreignPathId =
                     pathResolver.addPath(*foreignNodeId, *lookup->getForeignField());
 
@@ -208,6 +208,8 @@ StatusWith<AggJoinModel> AggJoinModel::constructJoinModel(const Pipeline& pipeli
                     // Cannot add an edge for existing nodes.
                     return Status(ErrorCodes::BadValue, "Graph is too big: too many edges");
                 }
+            } else {
+                pathResolver.addNode(*foreignNodeId, lookup->getAsField());
             }
 
             // TODO SERVER-111164: add edges from $expr's

src/mongo/db/query/compiler/optimizer/join/agg_join_model_test.cpp

@@ -411,4 +411,23 @@ TEST_F(PipelineAnalyzerTest, LongPrefix) {
     auto& joinModel = swJoinModel.getValue();
     goldenCtx.outStream() << joinModel.toString(true) << std::endl;
 }
+
+TEST_F(PipelineAnalyzerTest, LocalFieldOverride) {
+    unittest::GoldenTestContext goldenCtx(&goldenTestConfig);
+    const auto query = R"([
+            {$lookup: {from: "A", localField: "a", foreignField: "b", as: "a"}},
+            {$unwind: "$a"},
+            {$lookup: {from: "B", localField: "b", foreignField: "b", as: "b"}},
+            {$unwind: "$b"}
+        ])";
+
+    auto pipeline = makePipeline(query, {"A", "B"});
+
+    ASSERT_TRUE(AggJoinModel::pipelineEligibleForJoinReordering(*pipeline));
+
+    auto swJoinModel = AggJoinModel::constructJoinModel(*pipeline);
+    ASSERT_OK(swJoinModel);
+    auto& joinModel = swJoinModel.getValue();
+    goldenCtx.outStream() << joinModel.toString(true) << std::endl;
+}
 }  // namespace mongo::join_ordering

src/mongo/db/test_output/query/join/agg_join_model/pipeline_analyzer_test/local_field_override.txt

@@ -0,0 +1,98 @@
+{
+    "graph": {
+        "nodes": [
+            {
+                "collectionName": "test.pipeline_test",
+                "accessPath": {
+                    "filter": {}
+                },
+                "embedPath": ""
+            },
+            {
+                "collectionName": "test.A",
+                "accessPath": {
+                    "filter": {}
+                },
+                "embedPath": "a"
+            },
+            {
+                "collectionName": "test.B",
+                "accessPath": {
+                    "filter": {}
+                },
+                "embedPath": "b"
+            }
+        ],
+        "edges": [
+            {
+                "predicates": [
+                    {
+                        "op": "eq",
+                        "left": {"$numberInt":"0"},
+                        "right": {"$numberInt":"1"}
+                    }
+                ],
+                "left": "0000000000000000000000000000000000000000000000000000000000000001",
+                "right": "0000000000000000000000000000000000000000000000000000000000000010"
+            },
+            {
+                "predicates": [
+                    {
+                        "op": "eq",
+                        "left": {"$numberInt":"2"},
+                        "right": {"$numberInt":"3"}
+                    }
+                ],
+                "left": "0000000000000000000000000000000000000000000000000000000000000001",
+                "right": "0000000000000000000000000000000000000000000000000000000000000100"
+            }
+        ]
+    },
+    "resolvedPaths": [
+        {
+            "nodeId": {"$numberInt":"0"},
+            "fieldName": "a"
+        },
+        {
+            "nodeId": {"$numberInt":"1"},
+            "fieldName": "b"
+        },
+        {
+            "nodeId": {"$numberInt":"0"},
+            "fieldName": "b"
+        },
+        {
+            "nodeId": {"$numberInt":"2"},
+            "fieldName": "b"
+        }
+    ],
+    "prefix": [
+        {
+            "$lookup": {
+                "from": "A",
+                "as": "a",
+                "localField": "a",
+                "foreignField": "b"
+            }
+        },
+        {
+            "$unwind": {
+                "path": "$a"
+            }
+        },
+        {
+            "$lookup": {
+                "from": "B",
+                "as": "b",
+                "localField": "b",
+                "foreignField": "b"
+            }
+        },
+        {
+            "$unwind": {
+                "path": "$b"
+            }
+        }
+    ],
+    "suffix": []
+}