diff options
author | Drew Paroski <drew.paroski@mongodb.com> | 2022-09-10 01:28:17 +0000 |
---|---|---|
committer | Evergreen Agent <no-reply@evergreen.mongodb.com> | 2022-10-20 22:37:10 +0000 |
commit | 7bad8c9b694f599e49b4699f2dd149ebe210d188 (patch) | |
tree | 20413245099638a8af0d420d453d9e259474f03f | |
parent | 9027627ca94b5dbdcd741fcb520c45ee2dcbd184 (diff) | |
download | mongo-7bad8c9b694f599e49b4699f2dd149ebe210d188.tar.gz |
SERVER-69648 [SBE] Improve visit(ExpressionConcatArrays*) and visit(ExpressionObject*)
-rw-r--r-- | src/mongo/db/exec/sbe/expressions/expression.cpp | 3 | ||||
-rw-r--r-- | src/mongo/db/exec/sbe/vm/vm.cpp | 25 | ||||
-rw-r--r-- | src/mongo/db/exec/sbe/vm/vm.h | 2 | ||||
-rw-r--r-- | src/mongo/db/query/sbe_stage_builder_expression.cpp | 127 |
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); |