SERVER-67050 Improve SBE NestedLoopJoin

1 files changed, 148 insertions, 65 deletions

diff --git a/src/mongo/db/query/sbe_stage_builder_lookup.cpp b/src/mongo/db/query/sbe_stage_builder_lookup.cpp
index aa4ae7fbde3..a70c4bd0346 100644
--- a/src/mongo/db/query/sbe_stage_builder_lookup.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_lookup.cpp
@@ -322,6 +322,25 @@ std::pair<SlotId /* keyValueSlot */, std::unique_ptr<sbe::PlanStage>> buildForei
     return {keyValueSlot, std::move(currentStage)};
 }
 
+std::pair<SlotId /* keyValuesSetSlot */, EvalStage> replaceEmptySetWithNullArray(
+    EvalStage innerStage,
+    SlotId innerRecordSlot,
+    SlotIdGenerator& slotIdGenerator,
+    const PlanNodeId nodeId) {
+    auto [arrayWithNullTag, arrayWithNullVal] = makeNewArray();
+    auto arrayWithNull = makeConstant(arrayWithNullTag, arrayWithNullVal);
+    value::Array* arrayWithNullView = getArrayView(arrayWithNullVal);
+    arrayWithNullView->push_back(TypeTags::Null, 0);
+    auto nonEmptySetSlot = slotIdGenerator.generate();
+    return {nonEmptySetSlot,
+            makeProject(std::move(innerStage),
+                        nodeId,
+                        nonEmptySetSlot,
+                        makeE<EIf>(makeFunction("isArrayEmpty", makeVariable(innerRecordSlot)),
+                                   std::move(arrayWithNull),
+                                   makeVariable(innerRecordSlot)))};
+}
+
 // Creates stages for traversing path 'fp' in the record from 'inputSlot'. Puts the set of key
 // values into 'keyValuesSetSlot. For example, if the record in the 'inputSlot' is:
 //     {a: [{b:[1,[2,3]]}, {b:4}, {b:1}, {b:2}]},
@@ -359,35 +378,21 @@ std::pair<SlotId /* keyValuesSetSlot */, std::unique_ptr<sbe::PlanStage>> buildK
     // that contains a single 'null' value. The set of foreign key values also can be empty but it
     // should produce no matches so we leave it empty.
     if (joinSide == JoinSide::Local) {
-        auto [arrayWithNullTag, arrayWithNullVal] = makeNewArray();
-        std::unique_ptr<EExpression> arrayWithNull =
-            makeConstant(arrayWithNullTag, arrayWithNullVal);
-        value::Array* arrayWithNullView = getArrayView(arrayWithNullVal);
-        arrayWithNullView->push_back(TypeTags::Null, 0);
-
-        std::unique_ptr<EExpression> isNonEmptySetExpr =
-            makeBinaryOp(EPrimBinary::greater,
-                         makeFunction("getArraySize", makeVariable(keyValuesSetSlot)),
-                         makeConstant(TypeTags::NumberInt32, 0));
-
-        SlotId nonEmptySetSlot = slotIdGenerator.generate();
-        packedKeyValuesStage = makeProject(std::move(packedKeyValuesStage),
-                                           nodeId,
-                                           nonEmptySetSlot,
-                                           makeE<EIf>(std::move(isNonEmptySetExpr),
-                                                      makeVariable(keyValuesSetSlot),
-                                                      std::move(arrayWithNull)));
-        keyValuesSetSlot = nonEmptySetSlot;
+        std::tie(keyValuesSetSlot, packedKeyValuesStage) = replaceEmptySetWithNullArray(
+            std::move(packedKeyValuesStage),  // NOLINT(bugprone-use-after-move)
+            keyValuesSetSlot,
+            slotIdGenerator,
+            nodeId);
     }
 
     // Attach the set of key values to the original local record.
-    std::unique_ptr<sbe::PlanStage> nljLocalWithKeyValuesSet =
-        makeS<LoopJoinStage>(std::move(inputStage),
-                             packedKeyValuesStage.extractStage(nodeId),
-                             makeSV(recordSlot) /* outerProjects */,
-                             makeSV(recordSlot) /* outerCorrelated */,
-                             nullptr /* predicate */,
-                             nodeId);
+    std::unique_ptr<sbe::PlanStage> nljLocalWithKeyValuesSet = makeS<LoopJoinStage>(
+        std::move(inputStage),
+        packedKeyValuesStage.extractStage(nodeId),  // NOLINT(bugprone-use-after-move)
+        makeSV(recordSlot) /* outerProjects */,
+        makeSV(recordSlot) /* outerCorrelated */,
+        nullptr /* predicate */,
+        nodeId);
 
     return {keyValuesSetSlot, std::move(nljLocalWithKeyValuesSet)};
 }
@@ -457,51 +462,124 @@ std::pair<SlotId /* resultSlot */, std::unique_ptr<sbe::PlanStage>> buildForeign
         makeLimitSkip(std::move(unionStage), nodeId, 1 /* limit */).extractStage(nodeId));
 }
 
-// Creates stages that extract key values from the given foreign record, compares them to the local
-// key values and groups the matching records into an array.
+/**
+ * Build keys set for NLJ foreign side using traverseF expression. Creates stages that extract key
+ * values from the given foreign record, compares them to the local key values and groups the
+ * matching records into an array.
+ *
+ * The traverseF expression will iterate each key value, including terminal arrays as
+ * a whole value, and compare it against local key set 'localKeySlot'. For example,
+ * if the record in the 'foreignRecordSlot' is:
+ *     {a: [{b:[1,[2,3]]}, {b:4}, {b:1}, {b:2}]},
+ * path "a.b" will be iterated as: 1, [2,3], [1, [2, 3]], 4, 1, 2.
+ * Scalars inside arrays on the path are skipped, that is, if the record in the 'foreignRecordSlot'
+ * is:  {a: [42, {b:{c:1}}, {b: [41,42,{c:2}]}, {b:42}, {b:{c:3}}]},
+ * path "a.b.c" will be iterated as: 1, 2, null, 3.
+ * Replaces other missing terminals with 'null', that is, if the record in the 'foreignRecordSlot'
+ * is:  {a: [{b:1}, {b:[]}, {no_b:42}, {b:2}]},
+ * path "a.b" will be iterated as: 1, [], null, 2.
+ *
+ * Here is an example plan for the NLJ inner side:
+ * limit 1
+ * union [unionOutputSlot] [
+ *   branch0[projOutputSlot]
+ *     project [projOutputSlot = getElement(groupSlot, 0)]
+ *     group [] [groupSlot = addToArrayCapped(foreignRecordSlot, 104857600)]
+ *     filter {traverseF (
+ *       let [
+ *           l11.0 = fillEmpty (getField (foreignRecordSlot, "a"), null)
+ *       ]
+ *       in
+ *           if typeMatch (l11.0, 24)
+ *           then l11.0
+ *           else Nothing
+ *       , lambda(l3.0) {
+ *           if fillEmpty (isObject (l3.0), true)
+ *           then traverseF (
+ *             fillEmpty (getField (l3.0, "b"), null), lambda(l2.0) {isMember (l2.0, localKeySlot)},
+ *             true),
+ *           else false
+ *        }, false)}
+ *     scan foreignRecordSlot recordIdSlot none none none none [] @uuid true false
+ * branch1[emptySlot] project [emptySlot = []] limit 1 coscan
+ * ]
+ */
 std::pair<SlotId /* matched docs */, std::unique_ptr<sbe::PlanStage>> buildForeignMatches(
-    SlotId localKeysSetSlot,
+    SlotId localKeySlot,
     std::unique_ptr<sbe::PlanStage> foreignStage,
     SlotId foreignRecordSlot,
     const FieldPath& foreignFieldName,
     boost::optional<SlotId> collatorSlot,
     const PlanNodeId nodeId,
     SlotIdGenerator& slotIdGenerator,
+    FrameIdGenerator& frameIdGenerator,
     bool allowDiskUse) {
-    // 'buildForeignKeysStream()' does not take a child stage. To connect the extraction of key
-    // values it does to the 'foreignStage' we'll use an 'nlj', inner branch of which will do the
-    // key value extraction and filtering. It's possible that the same foreign record contains
-    // multiple matching values and to avoid duplication of the record we must limit the inner
-    // branch to '1'.
-    auto [foreignValueSlot, currentRootStage] =
-        buildForeignKeysStream(foreignRecordSlot, foreignFieldName, nodeId, slotIdGenerator);
-
-    currentRootStage =
-        makeLimitTree(makeS<FilterStage<false /*IsConst*/>>(
-                          std::move(currentRootStage),
-                          collatorSlot ? makeFunction("collIsMember",
-                                                      makeVariable(*collatorSlot),
-                                                      makeVariable(foreignValueSlot),
-                                                      makeVariable(localKeysSetSlot))
-                                       : makeFunction("isMember",
-                                                      makeVariable(foreignValueSlot),
-                                                      makeVariable(localKeysSetSlot)),
-                          nodeId),
-                      nodeId,
-                      1 /* take the foreign record once, even if multiple key values match */);
-
-    currentRootStage = makeS<LoopJoinStage>(std::move(foreignStage) /* outer */,
-                                            std::move(currentRootStage) /* inner */,
-                                            makeSV(foreignRecordSlot) /* outerProjects */,
-                                            makeSV(foreignRecordSlot) /* outerCorrelated */,
-                                            nullptr,
-                                            nodeId);
+    auto frameId = frameIdGenerator.generate();
+    auto lambdaArg = makeVariable(frameId, 0);
+    auto filter = collatorSlot
+        ? makeFunction("collIsMember"_sd,
+                       makeVariable(*collatorSlot),
+                       lambdaArg->clone(),
+                       makeVariable(localKeySlot))
+        : makeFunction("isMember"_sd, lambdaArg->clone(), makeVariable(localKeySlot));
+
+    // Recursively create traverseF expressions to iterate elements in 'foreignRecordSlot' with path
+    // 'foreignFieldName', and check if key is in set 'localKeySlot'.
+    //
+    // If a non-terminal field is an array, we will ignore any element that is not an object inside
+    // the array.
+    const int32_t foreignPathLength = foreignFieldName.getPathLength();
+    for (int32_t i = foreignPathLength - 1; i >= 0; --i) {
+        auto arrayLambda = makeE<ELocalLambda>(frameId, std::move(filter));
+        frameId = frameIdGenerator.generate();
+        lambdaArg = i == 0 ? makeVariable(foreignRecordSlot) : makeVariable(frameId, 0);
+
+        auto getFieldOrNull = makeFillEmptyNull(makeFunction(
+            "getField"_sd, lambdaArg->clone(), makeConstant(foreignFieldName.getFieldName(i))));
+
+        // Non object/array field will be converted into Nothing, passing along recursive traverseF
+        // and will be treated as null to compared against local key set.
+        if (i != foreignPathLength - 1) {
+            getFieldOrNull = makeLocalBind(
+                &frameIdGenerator,
+                [&](sbe::EVariable var) {
+                    return sbe::makeE<sbe::EIf>(
+                        makeFunction("typeMatch"_sd,
+                                     var.clone(),
+                                     makeConstant(sbe::value::TypeTags::NumberInt64,
+                                                  sbe::value::bitcastFrom<int64_t>(
+                                                      getBSONTypeMask(BSONType::Array) |
+                                                      getBSONTypeMask(BSONType::Object)))),
+                        var.clone(),
+                        makeConstant(sbe::value::TypeTags::Nothing, 0));
+                },
+                std::move(getFieldOrNull));
+        }
+
+        filter = makeFunction(
+            "traverseF"_sd,
+            std::move(getFieldOrNull),
+            std::move(arrayLambda),
+            makeConstant(TypeTags::Boolean, i == foreignPathLength - 1) /*compareArray*/);
+
+        if (i > 0) {
+            // Ignoring the nulls produced by missing field in array.
+            filter =
+                sbe::makeE<sbe::EIf>(makeFunction("fillEmpty"_sd,
+                                                  makeFunction("isObject"_sd, lambdaArg->clone()),
+                                                  makeConstant(TypeTags::Boolean, true)),
+                                     std::move(filter),
+                                     makeConstant(TypeTags::Boolean, false));
+        }
+    }
 
     // Group the matched foreign documents into a list.
     // It creates a union stage internally so that when there's no matching foreign records, an
     // empty array will be returned.
     return buildForeignMatchedArray(
-        EvalStage{std::move(currentRootStage), makeSV(foreignRecordSlot)},
+        makeFilter<false /*IsConst*/>(EvalStage{std::move(foreignStage), makeSV(foreignRecordSlot)},
+                                      std::move(filter),
+                                      nodeId),
         foreignRecordSlot,
         nodeId,
         slotIdGenerator,
@@ -515,10 +593,11 @@ std::pair<SlotId /* matched docs */, std::unique_ptr<sbe::PlanStage>> buildNljLo
     const FieldPath& localFieldName,
     std::unique_ptr<sbe::PlanStage> foreignStage,
     SlotId foreignRecordSlot,
-    StringData foreignFieldName,
+    const FieldPath& foreignFieldName,
     boost::optional<SlotId> collatorSlot,
     const PlanNodeId nodeId,
-    SlotIdGenerator& slotIdGenerator) {
+    SlotIdGenerator& slotIdGenerator,
+    FrameIdGenerator& frameIdGenerator) {
     CurOp::get(state.opCtx)->debug().nestedLoopJoin += 1;
 
     // Build the outer branch that produces the set of local key values.
@@ -540,6 +619,7 @@ std::pair<SlotId /* matched docs */, std::unique_ptr<sbe::PlanStage>> buildNljLo
                                                                     collatorSlot,
                                                                     nodeId,
                                                                     slotIdGenerator,
+                                                                    frameIdGenerator,
                                                                     state.allowDiskUse);
 
     // 'innerRootStage' should not participate in trial run tracking as the number of reads that
@@ -556,7 +636,6 @@ std::pair<SlotId /* matched docs */, std::unique_ptr<sbe::PlanStage>> buildNljLo
                              makeSV(localKeySlot) /* outerCorrelated */,
                              nullptr /* predicate */,
                              nodeId);
-
     return {matchedRecordsSlot, std::move(nlj)};
 }
 
@@ -864,6 +943,7 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage(
                                                                      collatorSlot,
                                                                      nodeId,
                                                                      slotIdGenerator,
+                                                                     frameIdGenerator,
                                                                      state.allowDiskUse);
 
     // 'foreignGroupStage' should not participate in trial run tracking as the number of reads
@@ -953,7 +1033,8 @@ std::pair<SlotId /*matched docs*/, std::unique_ptr<sbe::PlanStage>> buildLookupS
     const FieldPath& foreignFieldName,
     boost::optional<SlotId> collatorSlot,
     const PlanNodeId nodeId,
-    SlotIdGenerator& slotIdGenerator) {
+    SlotIdGenerator& slotIdGenerator,
+    FrameIdGenerator& frameIdGenerator) {
     switch (lookupStrategy) {
         case EqLookupNode::LookupStrategy::kNestedLoopJoin:
             return buildNljLookupStage(state,
@@ -962,10 +1043,11 @@ std::pair<SlotId /*matched docs*/, std::unique_ptr<sbe::PlanStage>> buildLookupS
                                        localFieldName,
                                        std::move(foreignStage),
                                        foreignRecordSlot,
-                                       foreignFieldName.fullPath(),
+                                       foreignFieldName,
                                        collatorSlot,
                                        nodeId,
-                                       slotIdGenerator);
+                                       slotIdGenerator,
+                                       frameIdGenerator);
         case EqLookupNode::LookupStrategy::kHashJoin:
             return buildHashJoinLookupStage(state,
                                             std::move(localStage),
@@ -1138,7 +1220,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
                                         eqLookupNode->joinFieldForeign,
                                         collatorSlot,
                                         eqLookupNode->nodeId(),
-                                        _slotIdGenerator);
+                                        _slotIdGenerator,
+                                        _frameIdGenerator);
             }
             default:
                 MONGO_UNREACHABLE_TASSERT(5842605);