SERVER-69648 [SBE] Improve visit(ExpressionConcatArrays*) and visit(ExpressionObject*)

4 files changed, 63 insertions, 94 deletions

diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index 30d73ec7d7e..6a740db24fa 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -472,7 +472,8 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = {
     {"tan", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::tan, false}},
     {"tanh", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::tanh, false}},
     {"round", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::round, false}},
-    {"concat", BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::concat, false}},
+    {"concat", BuiltinFn{kAnyNumberOfArgs, vm::Builtin::concat, false}},
+    {"concatArrays", BuiltinFn{kAnyNumberOfArgs, vm::Builtin::concatArrays, false}},
     {"isMember", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::isMember, false}},
     {"collIsMember", BuiltinFn{[](size_t n) { return n == 3; }, vm::Builtin::collIsMember, false}},
     {"indexOfBytes",
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index 38a17246d7e..cc055ca7471 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -3254,6 +3254,29 @@ FastTuple<bool, value::TypeTags, value::Value> ByteCode::builtinConcat(ArityType
     return {true, strTag, strValue};
 }
 
+FastTuple<bool, value::TypeTags, value::Value> ByteCode::builtinConcatArrays(ArityType arity) {
+    auto [resTag, resVal] = value::makeNewArray();
+    value::ValueGuard resGuard{resTag, resVal};
+    auto resView = value::getArrayView(resVal);
+
+    for (ArityType idx = 0; idx < arity; ++idx) {
+        auto [_, tag, val] = getFromStack(idx);
+        if (!value::isArray(tag)) {
+            return {false, value::TypeTags::Nothing, 0};
+        }
+
+        for (auto ae = value::ArrayEnumerator{tag, val}; !ae.atEnd(); ae.advance()) {
+            auto [elTag, elVal] = ae.getViewOfValue();
+            auto [copyTag, copyVal] = value::copyValue(elTag, elVal);
+            resView->push_back(copyTag, copyVal);
+        }
+    }
+
+    resGuard.reset();
+
+    return {true, resTag, resVal};
+}
+
 std::pair<value::TypeTags, value::Value> ByteCode::genericIsMember(value::TypeTags lhsTag,
                                                                    value::Value lhsVal,
                                                                    value::TypeTags rhsTag,
@@ -4695,6 +4718,8 @@ FastTuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builtin
             return builtinRound(arity);
         case Builtin::concat:
             return builtinConcat(arity);
+        case Builtin::concatArrays:
+            return builtinConcatArrays(arity);
         case Builtin::isMember:
             return builtinIsMember(arity);
         case Builtin::collIsMember:
diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h
index 3b7b1501d99..907775d8ec9 100644
--- a/src/mongo/db/exec/sbe/vm/vm.h
+++ b/src/mongo/db/exec/sbe/vm/vm.h
@@ -570,6 +570,7 @@ enum class Builtin : uint8_t {
     toLower,
     coerceToString,
     concat,
+    concatArrays,
     acos,
     acosh,
     asin,
@@ -1145,6 +1146,7 @@ private:
     FastTuple<bool, value::TypeTags, value::Value> builtinTanh(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinRound(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinConcat(ArityType arity);
+    FastTuple<bool, value::TypeTags, value::Value> builtinConcatArrays(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinIsMember(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinCollIsMember(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinIndexOfBytes(ArityType arity);
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index e5c1a28a889..d96f7ea6a92 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -1126,99 +1126,46 @@ public:
             return;
         }
 
-        sbe::EExpression::Vector nullChecks;
-        std::vector<EvalStage> unionBranches;
-        std::vector<sbe::value::SlotVector> unionInputSlots;
-        sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projections;
-
-        nullChecks.reserve(numChildren);
-        unionBranches.reserve(numChildren);
-        unionInputSlots.reserve(numChildren);
-        for (size_t idx = 0; idx < numChildren; ++idx) {
-            auto outputSlot = _context->state.slotId();
-            projections.emplace(outputSlot, _context->popExpr());
-            unionBranches.emplace_back();
-            unionInputSlots.emplace_back(sbe::makeSV(outputSlot));
-            nullChecks.emplace_back(generateNullOrMissing(outputSlot));
+        auto binds = sbe::makeEs();
+        for (size_t i = 0; i < numChildren; ++i) {
+            binds.emplace_back(_context->popExpr());
         }
+        std::reverse(binds.begin(), binds.end());
 
-        // Build a project to capture our child expressions.
-        std::reverse(std::begin(unionInputSlots), std::end(unionInputSlots));
-        auto project = makeProject(
-            _context->extractCurrentEvalStage(), std::move(projections), _context->planNodeId);
+        auto frameId = _context->state.frameId();
+        auto args = sbe::makeEs();
 
-        // Build a union stage to consolidate array input branches into a stream.
-        auto unionOutputSlot = _context->state.slotId();
-        auto unionStage = makeUnion(std::move(unionBranches),
-                                    std::move(unionInputSlots),
-                                    sbe::makeSV(unionOutputSlot),
-                                    _context->planNodeId);
+        std::unique_ptr<sbe::EExpression> checkArgsForNull;
+        for (size_t i = 0; i < numChildren; ++i) {
+            sbe::EVariable argRef(frameId, i);
+            args.emplace_back(argRef.clone());
 
-        auto collatorSlot = _context->state.data->env->getSlotIfExists("collator"_sd);
+            checkArgsForNull = checkArgsForNull ? makeBinaryOp(sbe::EPrimBinary::logicOr,
+                                                               std::move(checkArgsForNull),
+                                                               generateNullOrMissing(argRef))
+                                                : generateNullOrMissing(argRef);
+        }
 
-        // Build a filter that will throw an 'EFail' if any element coming from the union is NOT
-        // an array.
-        auto filter = makeFilter<false, false>(
-            std::move(unionStage),
-            makeBinaryOp(sbe::EPrimBinary::logicOr,
-                         makeFunction("isArray", makeVariable(unionOutputSlot)),
-                         sbe::makeE<sbe::EFail>(ErrorCodes::Error{5153400},
-                                                "$concatArrays only supports arrays")),
-            _context->planNodeId);
-
-        // Build subtree to handle nulls. If an input is null, return null. Otherwise, unwind the
-        // input and concatenate it into an array using addToArray.
-        auto unwindEvalStage =
-            makeUnwind(std::move(filter), _context->state.slotIdGenerator, _context->planNodeId);
-        auto unwindSlot = unwindEvalStage.getOutSlots().front();
-
-        // Create a group stage to append all streamed elements into one array. This is the final
-        // output when the input consists entirely of arrays.
-        auto finalAddToArrayExpr = makeFunction("addToArray", makeVariable(unwindSlot));
-        auto finalGroupSlot = _context->state.slotId();
-        auto finalGroupStage =
-            makeHashAgg(std::move(unwindEvalStage),
-                        sbe::makeSV(),
-                        sbe::makeEM(finalGroupSlot, std::move(finalAddToArrayExpr)),
-                        collatorSlot,
-                        _context->state.allowDiskUse,
-                        _context->planNodeId);
-
-        // Returns true if any of our input expressions return null.
-        using iter_t = sbe::EExpression::Vector::iterator;
-        auto checkPartsForNull = std::accumulate(
-            std::move_iterator<iter_t>(nullChecks.begin() + 1),
-            std::move_iterator<iter_t>(nullChecks.end()),
-            std::move(nullChecks.front()),
-            [](auto&& acc, auto&& b) {
-                return makeBinaryOp(sbe::EPrimBinary::logicOr, std::move(acc), std::move(b));
-            });
+        auto nullOrFailExpr =
+            sbe::makeE<sbe::EIf>(std::move(checkArgsForNull),
+                                 makeConstant(sbe::value::TypeTags::Null, 0),
+                                 sbe::makeE<sbe::EFail>(ErrorCodes::Error{5153400},
+                                                        "$concatArrays only supports arrays"));
 
-        // Create a branch stage to select between the branch that produces one null if any elements
-        // in the original input were null or missing, or otherwise select the branch that unwinds
-        // and concatenates elements into the output array.
-        auto [nullSlot, nullStage] = [&] {
-            auto outputSlot = _context->state.slotId();
-            auto nullEvalStage = makeProject(
-                {}, _context->planNodeId, outputSlot, makeConstant(sbe::value::TypeTags::Null, 0));
-            return std::make_pair(outputSlot, std::move(nullEvalStage));
-        }();
+        auto resultExpr = makeLocalBind(
+            _context->state.frameIdGenerator,
+            [&](sbe::EVariable concatArraysRef) {
+                // We optimize for the case where all of the args are arrays. If concatArrays()
+                // returns Nothing, then we deal with checking if any of the args are null and
+                // either returning null or raising an error.
+                return sbe::makeE<sbe::EIf>(makeFunction("exists", concatArraysRef.clone()),
+                                            concatArraysRef.clone(),
+                                            std::move(nullOrFailExpr));
+            },
+            sbe::makeE<sbe::EFunction>("concatArrays"_sd, std::move(args)));
 
-        auto branchSlot = _context->state.slotId();
-        auto branchNullEvalStage = makeBranch(std::move(nullStage),
-                                              std::move(finalGroupStage),
-                                              std::move(checkPartsForNull),
-                                              sbe::makeSV(nullSlot),
-                                              sbe::makeSV(finalGroupSlot),
-                                              sbe::makeSV(branchSlot),
-                                              _context->planNodeId);
-
-        // Create nlj to connect outer project with inner branch that handles null input.
-        _context->pushExpr(branchSlot,
-                           makeLoopJoin(std::move(project),
-                                        std::move(branchNullEvalStage),
-                                        _context->planNodeId,
-                                        _context->getLexicalEnvironment()));
+        _context->pushExpr(
+            sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(resultExpr)));
     }
     void visit(const ExpressionCond* expr) final {
         visitConditionalExpression(expr);
@@ -2602,13 +2549,7 @@ public:
             exprs[--i] = makeConstant(rit->first);
         }
 
-        auto fieldSlot{_context->state.slotIdGenerator->generate()};
-        auto stage = makeProject(_context->extractCurrentEvalStage(),
-                                 _context->planNodeId,
-                                 fieldSlot,
-                                 sbe::makeE<sbe::EFunction>("newObj"_sd, std::move(exprs)));
-
-        _context->pushExpr(fieldSlot, std::move(stage));
+        _context->pushExpr(sbe::makeE<sbe::EFunction>("newObj"_sd, std::move(exprs)));
     }
     void visit(const ExpressionOr* expr) final {
         visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicOr);