diff options
Diffstat (limited to 'src/mongo/db/pipeline/document_source_merge.cpp')
| -rw-r--r-- | src/mongo/db/pipeline/document_source_merge.cpp | 305 |
1 files changed, 86 insertions, 219 deletions
diff --git a/src/mongo/db/pipeline/document_source_merge.cpp b/src/mongo/db/pipeline/document_source_merge.cpp index c46cfa8c641..96a1dd55547 100644 --- a/src/mongo/db/pipeline/document_source_merge.cpp +++ b/src/mongo/db/pipeline/document_source_merge.cpp @@ -57,11 +57,10 @@ namespace { using MergeStrategyDescriptor = DocumentSourceMerge::MergeStrategyDescriptor; using MergeMode = MergeStrategyDescriptor::MergeMode; using MergeStrategy = MergeStrategyDescriptor::MergeStrategy; -using BatchedCommandGenerator = MergeStrategyDescriptor::BatchedCommandGenerator; using MergeStrategyDescriptorsMap = std::map<const MergeMode, const MergeStrategyDescriptor>; using WhenMatched = MergeStrategyDescriptor::WhenMatched; using WhenNotMatched = MergeStrategyDescriptor::WhenNotMatched; -using BatchTransform = DocumentSourceMerge::BatchTransform; +using BatchTransform = std::function<void(DocumentSourceMerge::BatchedObjects&)>; using UpdateModification = write_ops::UpdateModification; using UpsertType = MongoProcessInterface::UpsertType; @@ -84,71 +83,21 @@ constexpr auto kPipelineDiscardMode = MergeMode{WhenMatched::kPipeline, WhenNotM const auto kDefaultPipelineLet = BSON("new" << "$$ROOT"); -BatchedCommandGenerator makeInsertCommandGenerator() { - return [](const auto& expCtx, const auto& ns) -> BatchedCommandRequest { - return DocumentSourceMerge::DocumentSourceWriter::makeInsertCommand( - ns, expCtx->bypassDocumentValidation); - }; -} - -BatchedCommandGenerator makeUpdateCommandGenerator() { - return [](const auto& expCtx, const auto& ns) -> BatchedCommandRequest { - write_ops::UpdateCommandRequest updateOp(ns); - updateOp.setWriteCommandRequestBase([&] { - write_ops::WriteCommandRequestBase wcb; - wcb.setOrdered(false); - wcb.setBypassDocumentValidation(expCtx->bypassDocumentValidation); - return wcb; - }()); - auto [constants, letParams] = - expCtx->variablesParseState.transitionalCompatibilitySerialize(expCtx->variables); - updateOp.setLegacyRuntimeConstants(std::move(constants)); - if (!letParams.isEmpty()) { - updateOp.setLet(std::move(letParams)); - } - return BatchedCommandRequest(std::move(updateOp)); - }; -} - /** - * Converts 'batch' into a vector of UpdateOpEntries. + * Creates a merge strategy which uses update semantics to perform a merge operation. If + * 'BatchTransform' function is provided, it will be called to transform batched objects before + * passing them to the 'update'. */ -std::vector<write_ops::UpdateOpEntry> constructUpdateEntries( - DocumentSourceMerge::DocumentSourceWriter::BatchedObjects&& batch, - UpsertType upsert, - bool multi) { - std::vector<write_ops::UpdateOpEntry> updateEntries; - for (auto&& obj : batch) { - write_ops::UpdateOpEntry entry; - auto&& [q, u, c] = obj; - entry.setQ(std::move(q)); - entry.setU(std::move(u)); - entry.setC(std::move(c)); - entry.setUpsert(upsert != UpsertType::kNone); - entry.setUpsertSupplied({{entry.getUpsert(), upsert == UpsertType::kInsertSuppliedDoc}}); - entry.setMulti(multi); - - updateEntries.push_back(std::move(entry)); - } - return updateEntries; -} +MergeStrategy makeUpdateStrategy(UpsertType upsert, BatchTransform transform) { + return [upsert, transform]( + const auto& expCtx, const auto& ns, const auto& wc, auto epoch, auto&& batch) { + if (transform) { + transform(batch); + } -/** - * Creates a merge strategy which uses update semantics to perform a merge operation. - */ -MergeStrategy makeUpdateStrategy() { - return [](const auto& expCtx, - const auto& ns, - const auto& wc, - auto epoch, - auto&& batch, - auto&& bcr, - UpsertType upsert) { constexpr auto multi = false; - auto updateCommand = bcr.extractUpdateRequest(); - updateCommand->setUpdates(constructUpdateEntries(std::move(batch), upsert, multi)); uassertStatusOK(expCtx->mongoProcessInterface->update( - expCtx, ns, std::move(updateCommand), wc, upsert, multi, epoch)); + expCtx, ns, std::move(batch), wc, upsert, multi, epoch)); }; } @@ -157,22 +106,20 @@ MergeStrategy makeUpdateStrategy() { * that each document in the batch has a matching document in the 'ns' collection (note that a * matching document may not be modified as a result of an update operation, yet it still will be * counted as matching). If at least one document doesn't have a match, this strategy returns an - * error. + * error. If 'BatchTransform' function is provided, it will be called to transform batched objects + * before passing them to the 'update'. */ -MergeStrategy makeStrictUpdateStrategy() { - return [](const auto& expCtx, - const auto& ns, - const auto& wc, - auto epoch, - auto&& batch, - auto&& bcr, - UpsertType upsert) { +MergeStrategy makeStrictUpdateStrategy(UpsertType upsert, BatchTransform transform) { + return [upsert, transform]( + const auto& expCtx, const auto& ns, const auto& wc, auto epoch, auto&& batch) { + if (transform) { + transform(batch); + } + const int64_t batchSize = batch.size(); constexpr auto multi = false; - auto updateCommand = bcr.extractUpdateRequest(); - updateCommand->setUpdates(constructUpdateEntries(std::move(batch), upsert, multi)); auto updateResult = uassertStatusOK(expCtx->mongoProcessInterface->update( - expCtx, ns, std::move(updateCommand), wc, upsert, multi, epoch)); + expCtx, ns, std::move(batch), wc, upsert, multi, epoch)); uassert(ErrorCodes::MergeStageNoMatchingDocument, "{} could not find a matching document in the target collection " "for at least one document in the source collection"_format(kStageName), @@ -184,34 +131,28 @@ MergeStrategy makeStrictUpdateStrategy() { * Creates a merge strategy which uses insert semantics to perform a merge operation. */ MergeStrategy makeInsertStrategy() { - return [](const auto& expCtx, - const auto& ns, - const auto& wc, - auto epoch, - auto&& batch, - auto&& bcr, - UpsertType upsertType) { + return [](const auto& expCtx, const auto& ns, const auto& wc, auto epoch, auto&& batch) { std::vector<BSONObj> objectsToInsert(batch.size()); // The batch stores replacement style updates, but for this "insert" style of $merge we'd // like to just insert the new document without attempting any sort of replacement. std::transform(batch.begin(), batch.end(), objectsToInsert.begin(), [](const auto& obj) { return std::get<UpdateModification>(obj).getUpdateReplacement(); }); - auto insertCommand = bcr.extractInsertRequest(); - insertCommand->setDocuments(std::move(objectsToInsert)); - uassertStatusOK( - expCtx->mongoProcessInterface->insert(expCtx, ns, std::move(insertCommand), wc, epoch)); + uassertStatusOK(expCtx->mongoProcessInterface->insert( + expCtx, ns, std::move(objectsToInsert), wc, epoch)); }; } /** - * Creates a batched object transformation function which wraps 'obj' into the given 'updateOp' - * operator. + * Creates a batched objects transformation function which wraps each element of the + * 'batch.modifications' array into the given 'updateOp' operator. */ BatchTransform makeUpdateTransform(const std::string& updateOp) { - return [updateOp](auto& obj) { - std::get<UpdateModification>(obj) = UpdateModification::parseFromClassicUpdate( - BSON(updateOp << std::get<UpdateModification>(obj).getUpdateReplacement())); + return [updateOp](auto& batch) { + for (auto&& obj : batch) { + std::get<UpdateModification>(obj) = UpdateModification::parseFromClassicUpdate( + BSON(updateOp << std::get<UpdateModification>(obj).getUpdateReplacement())); + } }; } @@ -230,95 +171,53 @@ const MergeStrategyDescriptorsMap& getDescriptors() { // be initialized first. By wrapping the map into a function we can guarantee that it won't be // initialized until the first use, which is when the program already started and all global // variables had been initialized. - static const auto mergeStrategyDescriptors = - MergeStrategyDescriptorsMap{// whenMatched: replace, whenNotMatched: insert - {kReplaceInsertMode, - {kReplaceInsertMode, - {ActionType::insert, ActionType::update}, - makeUpdateStrategy(), - {}, - UpsertType::kGenerateNewDoc, - makeUpdateCommandGenerator()}}, - // whenMatched: replace, whenNotMatched: fail - {kReplaceFailMode, - {kReplaceFailMode, - {ActionType::update}, - makeStrictUpdateStrategy(), - {}, - UpsertType::kNone, - makeUpdateCommandGenerator()}}, - // whenMatched: replace, whenNotMatched: discard - {kReplaceDiscardMode, - {kReplaceDiscardMode, - {ActionType::update}, - makeUpdateStrategy(), - {}, - UpsertType::kNone, - makeUpdateCommandGenerator()}}, - // whenMatched: merge, whenNotMatched: insert - {kMergeInsertMode, - {kMergeInsertMode, - {ActionType::insert, ActionType::update}, - makeUpdateStrategy(), - makeUpdateTransform("$set"), - UpsertType::kGenerateNewDoc, - makeUpdateCommandGenerator()}}, - // whenMatched: merge, whenNotMatched: fail - {kMergeFailMode, - {kMergeFailMode, - {ActionType::update}, - makeStrictUpdateStrategy(), - makeUpdateTransform("$set"), - UpsertType::kNone, - makeUpdateCommandGenerator()}}, - // whenMatched: merge, whenNotMatched: discard - {kMergeDiscardMode, - {kMergeDiscardMode, - {ActionType::update}, - makeUpdateStrategy(), - makeUpdateTransform("$set"), - UpsertType::kNone, - makeUpdateCommandGenerator()}}, - // whenMatched: keepExisting, whenNotMatched: insert - {kKeepExistingInsertMode, - {kKeepExistingInsertMode, - {ActionType::insert, ActionType::update}, - makeUpdateStrategy(), - makeUpdateTransform("$setOnInsert"), - UpsertType::kGenerateNewDoc, - makeUpdateCommandGenerator()}}, - // whenMatched: [pipeline], whenNotMatched: insert - {kPipelineInsertMode, - {kPipelineInsertMode, - {ActionType::insert, ActionType::update}, - makeUpdateStrategy(), - {}, - UpsertType::kInsertSuppliedDoc, - makeUpdateCommandGenerator()}}, - // whenMatched: [pipeline], whenNotMatched: fail - {kPipelineFailMode, - {kPipelineFailMode, - {ActionType::update}, - makeStrictUpdateStrategy(), - {}, - UpsertType::kNone, - makeUpdateCommandGenerator()}}, - // whenMatched: [pipeline], whenNotMatched: discard - {kPipelineDiscardMode, - {kPipelineDiscardMode, - {ActionType::update}, - makeUpdateStrategy(), - {}, - UpsertType::kNone, - makeUpdateCommandGenerator()}}, - // whenMatched: fail, whenNotMatched: insert - {kFailInsertMode, - {kFailInsertMode, - {ActionType::insert}, - makeInsertStrategy(), - {}, - UpsertType::kNone, - makeInsertCommandGenerator()}}}; + static const auto mergeStrategyDescriptors = MergeStrategyDescriptorsMap{ + // whenMatched: replace, whenNotMatched: insert + {kReplaceInsertMode, + {kReplaceInsertMode, + {ActionType::insert, ActionType::update}, + makeUpdateStrategy(UpsertType::kGenerateNewDoc, {})}}, + // whenMatched: replace, whenNotMatched: fail + {kReplaceFailMode, + {kReplaceFailMode, {ActionType::update}, makeStrictUpdateStrategy(UpsertType::kNone, {})}}, + // whenMatched: replace, whenNotMatched: discard + {kReplaceDiscardMode, + {kReplaceDiscardMode, {ActionType::update}, makeUpdateStrategy(UpsertType::kNone, {})}}, + // whenMatched: merge, whenNotMatched: insert + {kMergeInsertMode, + {kMergeInsertMode, + {ActionType::insert, ActionType::update}, + makeUpdateStrategy(UpsertType::kGenerateNewDoc, makeUpdateTransform("$set"))}}, + // whenMatched: merge, whenNotMatched: fail + {kMergeFailMode, + {kMergeFailMode, + {ActionType::update}, + makeStrictUpdateStrategy(UpsertType::kNone, makeUpdateTransform("$set"))}}, + // whenMatched: merge, whenNotMatched: discard + {kMergeDiscardMode, + {kMergeDiscardMode, + {ActionType::update}, + makeUpdateStrategy(UpsertType::kNone, makeUpdateTransform("$set"))}}, + // whenMatched: keepExisting, whenNotMatched: insert + {kKeepExistingInsertMode, + {kKeepExistingInsertMode, + {ActionType::insert, ActionType::update}, + makeUpdateStrategy(UpsertType::kGenerateNewDoc, makeUpdateTransform("$setOnInsert"))}}, + // whenMatched: [pipeline], whenNotMatched: insert + {kPipelineInsertMode, + {kPipelineInsertMode, + {ActionType::insert, ActionType::update}, + makeUpdateStrategy(UpsertType::kInsertSuppliedDoc, {})}}, + // whenMatched: [pipeline], whenNotMatched: fail + {kPipelineFailMode, + {kPipelineFailMode, + {ActionType::update}, + makeStrictUpdateStrategy(UpsertType::kNone, {})}}, + // whenMatched: [pipeline], whenNotMatched: discard + {kPipelineDiscardMode, + {kPipelineDiscardMode, {ActionType::update}, makeUpdateStrategy(UpsertType::kNone, {})}}, + // whenMatched: fail, whenNotMatched: insert + {kFailInsertMode, {kFailInsertMode, {ActionType::insert}, makeInsertStrategy()}}}; return mergeStrategyDescriptors; } @@ -605,7 +504,7 @@ boost::optional<DocumentSource::DistributedPlanLogic> DocumentSourceMerge::distr return DocumentSourceWriter::distributedPlanLogic(); } -Value DocumentSourceMerge::serialize(const SerializationOptions& opts) const { +Value DocumentSourceMerge::serialize(boost::optional<ExplainOptions::Verbosity> explain) const { DocumentSourceMergeSpec spec; spec.setTargetNss(_outputNs); spec.setLet([&]() -> boost::optional<BSONObj> { @@ -615,27 +514,11 @@ Value DocumentSourceMerge::serialize(const SerializationOptions& opts) const { BSONObjBuilder bob; for (auto&& [name, expr] : *_letVariables) { - bob << opts.serializeFieldPathFromString(name) << expr->serialize(opts); + bob << name << expr->serialize(static_cast<bool>(explain)); } return bob.obj(); }()); - spec.setWhenMatched(MergeWhenMatchedPolicy{ - _descriptor.mode.first, [&]() -> boost::optional<std::vector<BSONObj>> { - if (!_pipeline.has_value()) { - return boost::none; - } - auto expCtxWithLetVariables = pExpCtx->copyWith(getOutputNs()); - if (spec.getLet()) { - BSONObjBuilder cleanLetSpecBuilder; - for (auto&& [name, expr] : *_letVariables) { - cleanLetSpecBuilder.append(name, BSONObj{}); - } - expCtxWithLetVariables->variables.seedVariablesWithLetParameters( - expCtxWithLetVariables.get(), cleanLetSpecBuilder.obj()); - } - return Pipeline::parse(_pipeline.value(), expCtxWithLetVariables) - ->serializeToBson(opts); - }()}); + spec.setWhenMatched(MergeWhenMatchedPolicy{_descriptor.mode.first, _pipeline}); spec.setWhenNotMatched(_descriptor.mode.second); spec.setOn([&]() { std::vector<std::string> mergeOnFields; @@ -645,7 +528,7 @@ Value DocumentSourceMerge::serialize(const SerializationOptions& opts) const { return mergeOnFields; }()); spec.setTargetCollectionVersion(_targetCollectionVersion); - return Value(Document{{getSourceName(), spec.toBSON(opts)}}); + return Value(Document{{getSourceName(), spec.toBSON()}}); } std::pair<DocumentSourceMerge::BatchObject, int> DocumentSourceMerge::makeBatchObject( @@ -660,29 +543,17 @@ std::pair<DocumentSourceMerge::BatchObject, int> DocumentSourceMerge::makeBatchO auto mergeOnFields = extractMergeOnFieldsFromDoc(doc, _mergeOnFields); auto mod = makeBatchUpdateModification(doc); auto vars = resolveLetVariablesIfNeeded(doc); - BatchObject batchObject{std::move(mergeOnFields), std::move(mod), std::move(vars)}; - if (_descriptor.transform) { - _descriptor.transform(batchObject); - } - - tassert(6628901, "_writeSizeEstimator should be initialized", _writeSizeEstimator); - return {batchObject, - _writeSizeEstimator->estimateUpdateSizeBytes(batchObject, _descriptor.upsertType)}; + auto modSize = mod.objsize() + (vars ? vars->objsize() : 0); + return {{std::move(mergeOnFields), std::move(mod), std::move(vars)}, modSize}; } -void DocumentSourceMerge::spill(BatchedCommandRequest&& bcr, BatchedObjects&& batch) try { +void DocumentSourceMerge::spill(BatchedObjects&& batch) try { DocumentSourceWriteBlock writeBlock(pExpCtx->opCtx); auto targetEpoch = _targetCollectionVersion ? boost::optional<OID>(_targetCollectionVersion->epoch()) : boost::none; - _descriptor.strategy(pExpCtx, - _outputNs, - _writeConcern, - targetEpoch, - std::move(batch), - std::move(bcr), - _descriptor.upsertType); + _descriptor.strategy(pExpCtx, _outputNs, _writeConcern, targetEpoch, std::move(batch)); } catch (const ExceptionFor<ErrorCodes::ImmutableField>& ex) { uassertStatusOKWithContext(ex.toStatus(), "$merge failed to update the matching document, did you " @@ -706,10 +577,6 @@ void DocumentSourceMerge::spill(BatchedCommandRequest&& bcr, BatchedObjects&& ba } } -BatchedCommandRequest DocumentSourceMerge::initializeBatchedWriteRequest() const { - return _descriptor.batchedCommandGenerator(pExpCtx, _outputNs); -} - void DocumentSourceMerge::waitWhileFailPointEnabled() { CurOpFailpointHelpers::waitWhileFailPointEnabled( &hangWhileBuildingDocumentSourceMergeBatch, |