diff --git a/jstests/noPassthrough/query/joins/local_field_override.js b/jstests/noPassthrough/query/joins/local_field_override.js new file mode 100644 index 00000000000..ce2fe7e769e --- /dev/null +++ b/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}); diff --git a/src/mongo/db/query/compiler/optimizer/join/agg_join_model.cpp b/src/mongo/db/query/compiler/optimizer/join/agg_join_model.cpp index 66e807ee034..8f40462e5d5 100644 --- a/src/mongo/db/query/compiler/optimizer/join/agg_join_model.cpp +++ b/src/mongo/db/query/compiler/optimizer/join/agg_join_model.cpp @@ -192,13 +192,13 @@ StatusWith 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::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 diff --git a/src/mongo/db/query/compiler/optimizer/join/agg_join_model_test.cpp b/src/mongo/db/query/compiler/optimizer/join/agg_join_model_test.cpp index 9ad07cedf84..bf446dd7759 100644 --- a/src/mongo/db/query/compiler/optimizer/join/agg_join_model_test.cpp +++ b/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 diff --git a/src/mongo/db/test_output/query/join/agg_join_model/pipeline_analyzer_test/local_field_override.txt b/src/mongo/db/test_output/query/join/agg_join_model/pipeline_analyzer_test/local_field_override.txt new file mode 100644 index 00000000000..cf8e8d5fd1d --- /dev/null +++ b/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": [] +}