SERVER-45399 Add basic union aggregation stage for unsharded collections

4 files changed, 623 insertions, 0 deletions

diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index e9a4e66b9c2..ae5f56ae340 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -323,6 +323,7 @@ pipelineEnv.Library(
         'document_source_sort.cpp',
         'document_source_sort_by_count.cpp',
         'document_source_tee_consumer.cpp',
+        'document_source_union_with.cpp',
         'document_source_unwind.cpp',
         'pipeline.cpp',
         'semantic_analysis.cpp',
@@ -447,6 +448,7 @@ env.CppUnitTest(
         'document_source_skip_test.cpp',
         'document_source_sort_by_count_test.cpp',
         'document_source_sort_test.cpp',
+        'document_source_union_with_test.cpp',
         'document_source_unwind_test.cpp',
         'expression_and_test.cpp',
         'expression_compare_test.cpp',
diff --git a/src/mongo/db/pipeline/document_source_union_with.cpp b/src/mongo/db/pipeline/document_source_union_with.cpp
new file mode 100644
index 00000000000..fa478bf653f
--- /dev/null
+++ b/src/mongo/db/pipeline/document_source_union_with.cpp
@@ -0,0 +1,187 @@
+/**
+ *    Copyright (C) 2019-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/pipeline/document_source_union_with.h"
+
+namespace mongo {
+
+REGISTER_DOCUMENT_SOURCE(unionWith,
+                         DocumentSourceUnionWith::LiteParsed::parse,
+                         DocumentSourceUnionWith::createFromBson);
+
+std::unique_ptr<DocumentSourceUnionWith::LiteParsed> DocumentSourceUnionWith::LiteParsed::parse(
+    const AggregationRequest& request, const BSONElement& spec) {
+    uassert(ErrorCodes::FailedToParse,
+            str::stream()
+                << "the $unionWith stage specification must be an object or string, but found "
+                << typeName(spec.type()),
+            spec.type() == BSONType::Object || spec.type() == BSONType::String);
+
+    NamespaceString unionNss;
+    stdx::unordered_set<NamespaceString> foreignNssSet;
+    boost::optional<LiteParsedPipeline> liteParsedPipeline;
+    if (spec.type() == BSONType::String) {
+        unionNss = NamespaceString(request.getNamespaceString().db(), spec.valueStringData());
+    } else {
+        unionNss =
+            NamespaceString(request.getNamespaceString().db(), spec["coll"].valueStringData());
+
+        // Recursively lite parse the nested pipeline, if one exists.
+        if (auto pipelineElem = spec["pipeline"]) {
+            auto pipeline =
+                uassertStatusOK(AggregationRequest::parsePipelineFromBSON(pipelineElem));
+            AggregationRequest foreignAggReq(unionNss, std::move(pipeline));
+            liteParsedPipeline = LiteParsedPipeline(foreignAggReq);
+
+            foreignNssSet.merge(liteParsedPipeline->getInvolvedNamespaces());
+        }
+    }
+
+    foreignNssSet.insert(unionNss);
+
+    return std::make_unique<DocumentSourceUnionWith::LiteParsed>(
+        std::move(unionNss), std::move(foreignNssSet), std::move(liteParsedPipeline));
+}
+
+boost::intrusive_ptr<DocumentSource> DocumentSourceUnionWith::createFromBson(
+    BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& expCtx) {
+    uassert(ErrorCodes::FailedToParse,
+            str::stream()
+                << "the $unionWith stage specification must be an object or string, but found "
+                << typeName(elem.type()),
+            elem.type() == BSONType::Object || elem.type() == BSONType::String);
+
+    NamespaceString unionNss;
+    std::vector<BSONObj> pipeline;
+    if (elem.type() == BSONType::String) {
+        unionNss = NamespaceString(expCtx->ns.db().toString(), elem.valueStringData());
+    } else {
+        bool sawColl = false;
+        bool sawPipeline = false;
+        for (auto&& arg : elem.embeddedObject()) {
+            auto fieldName = arg.fieldNameStringData();
+            if (fieldName == "coll") {
+                if (sawColl)
+                    uasserted(ErrorCodes::FailedToParse,
+                              str::stream() << "Redundant 'coll' argument given to $unionWith");
+                unionNss = NamespaceString(expCtx->ns.db().toString(), arg.valueStringData());
+                sawColl = true;
+            } else if (fieldName == "pipeline") {
+                if (sawPipeline)
+                    uasserted(ErrorCodes::FailedToParse,
+                              str::stream() << "Redundant 'pipeline' argument given to $unionWith");
+                pipeline = uassertStatusOK(AggregationRequest::parsePipelineFromBSON(arg));
+                sawPipeline = true;
+            } else
+                uasserted(ErrorCodes::FailedToParse,
+                          str::stream()
+                              << "Unknown argument given to $unionWith stage: " << fieldName);
+        }
+        if (!sawColl)
+            uasserted(ErrorCodes::FailedToParse,
+                      str::stream() << "No 'coll' argument given to $unionWith");
+    }
+    return make_intrusive<DocumentSourceUnionWith>(
+        expCtx,
+        uassertStatusOK(
+            Pipeline::parse(std::move(pipeline), expCtx->copyWith(std::move(unionNss)))));
+}
+
+DocumentSource::GetNextResult DocumentSourceUnionWith::doGetNext() {
+    if (!_sourceExhausted) {
+        auto nextInput = pSource->getNext();
+        if (!nextInput.isEOF()) {
+            return nextInput;
+        }
+        _sourceExhausted = true;
+        // All documents from the base collection have been returned, switch to iterating the sub-
+        // pipeline by falling through below.
+    }
+
+    if (!_cursorAttached) {
+        auto ctx = _pipeline->getContext();
+        _pipeline =
+            pExpCtx->mongoProcessInterface->attachCursorSourceToPipeline(ctx, _pipeline.release());
+        _cursorAttached = true;
+    }
+
+    if (auto res = _pipeline->getNext())
+        return std::move(*res);
+
+    return GetNextResult::makeEOF();
+}
+
+DocumentSource::GetModPathsReturn DocumentSourceUnionWith::getModifiedPaths() const {
+    // Since we might have a document arrive from the foreign pipeline with the same path as a
+    // document in the main pipeline. Without introspecting the sub-pipeline, we must report that
+    // all paths have been modified.
+    return {GetModPathsReturn::Type::kAllPaths, {}, {}};
+}
+
+void DocumentSourceUnionWith::doDispose() {
+    _pipeline->dispose(pExpCtx->opCtx);
+    _pipeline.reset();
+}
+
+void DocumentSourceUnionWith::serializeToArray(
+    std::vector<Value>& array, boost::optional<ExplainOptions::Verbosity> explain) const {
+    BSONArrayBuilder bab;
+    for (auto&& stage : _pipeline->serialize())
+        bab << stage;
+    Document doc = DOC(getSourceName() << DOC("coll" << _pipeline->getContext()->ns.coll()
+                                                     << "pipeline" << bab.arr()));
+    array.push_back(Value(doc));
+}
+
+DepsTracker::State DocumentSourceUnionWith::getDependencies(DepsTracker* deps) const {
+    return DepsTracker::State::SEE_NEXT;
+}
+
+void DocumentSourceUnionWith::detachFromOperationContext() {
+    // We have a pipeline we're going to be executing across multiple calls to getNext(), so we
+    // use Pipeline::detachFromOperationContext() to take care of updating the Pipeline's
+    // ExpressionContext.
+    _pipeline->detachFromOperationContext();
+}
+
+void DocumentSourceUnionWith::reattachToOperationContext(OperationContext* opCtx) {
+    // We have a pipeline we're going to be executing across multiple calls to getNext(), so we
+    // use Pipeline::reattachToOperationContext() to take care of updating the Pipeline's
+    // ExpressionContext.
+    _pipeline->reattachToOperationContext(opCtx);
+}
+
+void DocumentSourceUnionWith::addInvolvedCollections(
+    stdx::unordered_set<NamespaceString>* collectionNames) const {
+    collectionNames->insert(_pipeline->getContext()->ns);
+    collectionNames->merge(_pipeline->getInvolvedCollections());
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_union_with.h b/src/mongo/db/pipeline/document_source_union_with.h
new file mode 100644
index 00000000000..7181a7d2ab7
--- /dev/null
+++ b/src/mongo/db/pipeline/document_source_union_with.h
@@ -0,0 +1,134 @@
+/**
+ *    Copyright (C) 2019-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include <boost/optional.hpp>
+
+#include "mongo/db/pipeline/document_source.h"
+#include "mongo/db/pipeline/lite_parsed_pipeline.h"
+
+namespace mongo {
+
+class DocumentSourceUnionWith final : public DocumentSource {
+public:
+    static constexpr StringData kStageName = "$unionWith"_sd;
+
+    static boost::intrusive_ptr<DocumentSource> createFromBson(
+        BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& expCtx);
+
+    class LiteParsed final : public LiteParsedDocumentSource {
+    public:
+        static std::unique_ptr<LiteParsed> parse(const AggregationRequest& request,
+                                                 const BSONElement& spec);
+
+        LiteParsed(NamespaceString withNss,
+                   stdx::unordered_set<NamespaceString> foreignNssSet,
+                   boost::optional<LiteParsedPipeline> liteParsedPipeline)
+            : _withNss{std::move(withNss)},
+              _foreignNssSet(std::move(foreignNssSet)),
+              _liteParsedPipeline(std::move(liteParsedPipeline)) {}
+
+        stdx::unordered_set<NamespaceString> getInvolvedNamespaces() const final {
+            return {_foreignNssSet};
+        }
+
+        PrivilegeVector requiredPrivileges(bool) const final {
+            return {};
+        }
+
+        bool allowShardedForeignCollection(NamespaceString) const final {
+            return true;
+        }
+
+        bool allowedToPassthroughFromMongos() const final {
+            return true;
+        }
+
+    private:
+        const NamespaceString _withNss;
+        const stdx::unordered_set<NamespaceString> _foreignNssSet;
+        const boost::optional<LiteParsedPipeline> _liteParsedPipeline;
+    };
+
+    DocumentSourceUnionWith(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                            std::unique_ptr<Pipeline, PipelineDeleter> pipeline)
+        : DocumentSource(kStageName, expCtx), _pipeline(std::move(pipeline)) {}
+
+    const char* getSourceName() const final {
+        return kStageName.rawData();
+    }
+
+    void serializeToArray(
+        std::vector<Value>& array,
+        boost::optional<ExplainOptions::Verbosity> explain = boost::none) const final;
+
+    GetModPathsReturn getModifiedPaths() const final;
+
+    StageConstraints constraints(Pipeline::SplitState) const final {
+        return StageConstraints(StreamType::kStreaming,
+                                PositionRequirement::kNone,
+                                HostTypeRequirement::kNone,
+                                DiskUseRequirement::kNoDiskUse,
+                                FacetRequirement::kAllowed,
+                                TransactionRequirement::kNotAllowed,
+                                LookupRequirement::kNotAllowed);
+    }
+
+    DepsTracker::State getDependencies(DepsTracker* deps) const final;
+
+    boost::optional<DistributedPlanLogic> distributedPlanLogic() final {
+        // {shardsStage, mergingStage, sortPattern}
+        return DistributedPlanLogic{nullptr, this, boost::none};
+    }
+
+    void addInvolvedCollections(stdx::unordered_set<NamespaceString>* collectionNames) const final;
+
+    void detachFromOperationContext() final;
+
+    void reattachToOperationContext(OperationContext* opCtx) final;
+
+protected:
+    GetNextResult doGetNext() final;
+    void doDispose() final;
+
+private:
+    /**
+     * Should not be called; use serializeToArray instead.
+     */
+    Value serialize(boost::optional<ExplainOptions::Verbosity> explain = boost::none) const final {
+        MONGO_UNREACHABLE;
+    }
+
+    std::unique_ptr<Pipeline, PipelineDeleter> _pipeline;
+    bool _sourceExhausted = false;
+    bool _cursorAttached = false;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_union_with_test.cpp b/src/mongo/db/pipeline/document_source_union_with_test.cpp
new file mode 100644
index 00000000000..d45e9ce4f87
--- /dev/null
+++ b/src/mongo/db/pipeline/document_source_union_with_test.cpp
@@ -0,0 +1,300 @@
+/**
+ *    Copyright (C) 2018-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include <array>
+#include <deque>
+#include <list>
+#include <set>
+#include <vector>
+
+#include "mongo/db/exec/document_value/document.h"
+#include "mongo/db/exec/document_value/document_comparator.h"
+#include "mongo/db/exec/document_value/document_value_test_util.h"
+#include "mongo/db/pipeline/aggregation_context_fixture.h"
+#include "mongo/db/pipeline/document_source.h"
+#include "mongo/db/pipeline/document_source_add_fields.h"
+#include "mongo/db/pipeline/document_source_match.h"
+#include "mongo/db/pipeline/document_source_mock.h"
+#include "mongo/db/pipeline/document_source_union_with.h"
+#include "mongo/db/pipeline/pipeline.h"
+#include "mongo/db/pipeline/stub_mongo_process_interface_lookup_single_document.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/intrusive_counter.h"
+
+namespace mongo {
+namespace {
+
+using MockMongoInterface = StubMongoProcessInterfaceLookupSingleDocument;
+
+// This provides access to getExpCtx(), but we'll use a different name for this test suite.
+using DocumentSourceUnionWithTest = AggregationContextFixture;
+
+TEST_F(DocumentSourceUnionWithTest, BasicSerialUnions) {
+    const auto docs = std::array{Document{{"a", 1}}, Document{{"b", 1}}, Document{{"c", 1}}};
+    const auto mock = DocumentSourceMock::createForTest(docs[0]);
+    const auto mockDequeOne = std::deque<DocumentSource::GetNextResult>{Document{docs[1]}};
+    const auto mockDequeTwo = std::deque<DocumentSource::GetNextResult>{Document{docs[2]}};
+    const auto mockCtxOne = getExpCtx()->copyWith({});
+    mockCtxOne->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeOne);
+    const auto mockCtxTwo = getExpCtx()->copyWith({});
+    mockCtxTwo->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeTwo);
+    auto unionWithOne = DocumentSourceUnionWith(
+        mockCtxOne,
+        uassertStatusOK(
+            Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx())));
+    auto unionWithTwo = DocumentSourceUnionWith(
+        mockCtxTwo,
+        uassertStatusOK(
+            Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx())));
+    unionWithOne.setSource(mock.get());
+    unionWithTwo.setSource(&unionWithOne);
+
+    auto results = DocumentComparator().makeUnorderedDocumentSet();
+    for (auto& doc [[maybe_unused]] : docs) {
+        auto next = unionWithTwo.getNext();
+        ASSERT_TRUE(next.isAdvanced());
+        const auto [ignored, inserted] = results.insert(next.releaseDocument());
+        ASSERT_TRUE(inserted);
+    }
+    for (const auto& doc : docs)
+        ASSERT_TRUE(results.find(doc) != results.end());
+
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+}
+
+TEST_F(DocumentSourceUnionWithTest, BasicNestedUnions) {
+    const auto docs = std::array{Document{{"a", 1}}, Document{{"b", 1}}, Document{{"c", 1}}};
+    const auto mock = DocumentSourceMock::createForTest(docs[0]);
+    const auto mockDequeOne = std::deque<DocumentSource::GetNextResult>{Document{docs[1]}};
+    const auto mockDequeTwo = std::deque<DocumentSource::GetNextResult>{Document{docs[2]}};
+    const auto mockCtxOne = getExpCtx()->copyWith({});
+    mockCtxOne->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeOne);
+    const auto mockCtxTwo = getExpCtx()->copyWith({});
+    mockCtxTwo->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeTwo);
+    auto unionWithOne = make_intrusive<DocumentSourceUnionWith>(
+        mockCtxOne,
+        uassertStatusOK(
+            Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx())));
+    auto unionWithTwo = DocumentSourceUnionWith(
+        mockCtxTwo,
+        uassertStatusOK(Pipeline::create(
+            std::list<boost::intrusive_ptr<DocumentSource>>{unionWithOne}, getExpCtx())));
+    unionWithTwo.setSource(mock.get());
+
+    auto results = DocumentComparator().makeUnorderedDocumentSet();
+    for (auto& doc [[maybe_unused]] : docs) {
+        auto next = unionWithTwo.getNext();
+        ASSERT_TRUE(next.isAdvanced());
+        const auto [ignored, inserted] = results.insert(next.releaseDocument());
+        ASSERT_TRUE(inserted);
+    }
+    for (const auto& doc : docs)
+        ASSERT_TRUE(results.find(doc) != results.end());
+
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+}
+
+TEST_F(DocumentSourceUnionWithTest, UnionsWithNonEmptySubPipelines) {
+    const auto inputDocs = std::array{Document{{"a", 1}}, Document{{"b", 1}}, Document{{"c", 1}}};
+    const auto outputDocs = std::array{Document{{"a", 1}}, Document{{"c", 1}, {"d", 1}}};
+    const auto mock = DocumentSourceMock::createForTest(inputDocs[0]);
+    const auto mockDequeOne = std::deque<DocumentSource::GetNextResult>{Document{inputDocs[1]}};
+    const auto mockDequeTwo = std::deque<DocumentSource::GetNextResult>{Document{inputDocs[2]}};
+    const auto mockCtxOne = getExpCtx()->copyWith({});
+    mockCtxOne->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeOne);
+    const auto mockCtxTwo = getExpCtx()->copyWith({});
+    mockCtxTwo->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeTwo);
+    const auto filter = DocumentSourceMatch::create(BSON("d" << 1), mockCtxOne);
+    const auto proj = DocumentSourceAddFields::create(BSON("d" << 1), mockCtxTwo);
+    auto unionWithOne = DocumentSourceUnionWith(
+        mockCtxOne,
+        uassertStatusOK(Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{filter},
+                                         getExpCtx())));
+    auto unionWithTwo = DocumentSourceUnionWith(
+        mockCtxTwo,
+        uassertStatusOK(
+            Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{proj}, getExpCtx())));
+    unionWithOne.setSource(mock.get());
+    unionWithTwo.setSource(&unionWithOne);
+
+    auto results = DocumentComparator().makeUnorderedDocumentSet();
+    for (auto& doc [[maybe_unused]] : outputDocs) {
+        auto next = unionWithTwo.getNext();
+        ASSERT_TRUE(next.isAdvanced());
+        const auto [ignored, inserted] = results.insert(next.releaseDocument());
+        ASSERT_TRUE(inserted);
+    }
+    for (const auto& doc : outputDocs)
+        ASSERT_TRUE(results.find(doc) != results.end());
+
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+}
+
+TEST_F(DocumentSourceUnionWithTest, SerializeAndParseWithPipeline) {
+    auto bson =
+        BSON("$unionWith" << BSON("coll"
+                                  << "foo"
+                                  << "pipeline"
+                                  << BSON_ARRAY(
+                                         BSON("$addFields" << BSON("a" << BSON("$const" << 3))))));
+    auto unionWith = DocumentSourceUnionWith::createFromBson(bson.firstElement(), getExpCtx());
+    ASSERT(unionWith->getSourceName() == DocumentSourceUnionWith::kStageName);
+    std::vector<Value> serializedArray;
+    unionWith->serializeToArray(serializedArray);
+    auto serializedBson = serializedArray[0].getDocument().toBson();
+    ASSERT_BSONOBJ_EQ(serializedBson, bson);
+    unionWith = DocumentSourceUnionWith::createFromBson(serializedBson.firstElement(), getExpCtx());
+    ASSERT(unionWith != nullptr);
+    ASSERT(unionWith->getSourceName() == DocumentSourceUnionWith::kStageName);
+}
+
+TEST_F(DocumentSourceUnionWithTest, SerializeAndParseWithoutPipeline) {
+    auto bson = BSON("$unionWith"
+                     << "foo");
+    auto desugaredBson = BSON("$unionWith" << BSON("coll"
+                                                   << "foo"
+                                                   << "pipeline" << BSONArray()));
+    auto unionWith = DocumentSourceUnionWith::createFromBson(bson.firstElement(), getExpCtx());
+    ASSERT(unionWith->getSourceName() == DocumentSourceUnionWith::kStageName);
+    std::vector<Value> serializedArray;
+    unionWith->serializeToArray(serializedArray);
+    auto serializedBson = serializedArray[0].getDocument().toBson();
+    ASSERT_BSONOBJ_EQ(serializedBson, desugaredBson);
+    unionWith = DocumentSourceUnionWith::createFromBson(serializedBson.firstElement(), getExpCtx());
+    ASSERT(unionWith != nullptr);
+    ASSERT(unionWith->getSourceName() == DocumentSourceUnionWith::kStageName);
+}
+
+TEST_F(DocumentSourceUnionWithTest, SerializeAndParseWithoutPipelineExtraSubobject) {
+    auto bson = BSON("$unionWith" << BSON("coll"
+                                          << "foo"));
+    auto desugaredBson = BSON("$unionWith" << BSON("coll"
+                                                   << "foo"
+                                                   << "pipeline" << BSONArray()));
+    auto unionWith = DocumentSourceUnionWith::createFromBson(bson.firstElement(), getExpCtx());
+    ASSERT(unionWith->getSourceName() == DocumentSourceUnionWith::kStageName);
+    std::vector<Value> serializedArray;
+    unionWith->serializeToArray(serializedArray);
+    auto serializedBson = serializedArray[0].getDocument().toBson();
+    ASSERT_BSONOBJ_EQ(serializedBson, desugaredBson);
+    unionWith = DocumentSourceUnionWith::createFromBson(serializedBson.firstElement(), getExpCtx());
+    ASSERT(unionWith != nullptr);
+    ASSERT(unionWith->getSourceName() == DocumentSourceUnionWith::kStageName);
+}
+
+TEST_F(DocumentSourceUnionWithTest, ParseErrors) {
+    ASSERT_THROWS_CODE(DocumentSourceUnionWith::createFromBson(
+                           BSON("$unionWith" << false).firstElement(), getExpCtx()),
+                       AssertionException,
+                       ErrorCodes::FailedToParse);
+    ASSERT_THROWS_CODE(DocumentSourceUnionWith::createFromBson(BSON("$unionWith" << BSON("coll"
+                                                                                         << "foo"
+                                                                                         << "coll"
+                                                                                         << "bar"))
+                                                                   .firstElement(),
+                                                               getExpCtx()),
+                       AssertionException,
+                       ErrorCodes::FailedToParse);
+    ASSERT_THROWS_CODE(
+        DocumentSourceUnionWith::createFromBson(
+            BSON("$unionWith" << BSON("coll"
+                                      << "foo"
+                                      << "pipeline"
+                                      << BSON_ARRAY(BSON("$addFields" << BSON("a" << 3))) << "coll"
+                                      << "bar"))
+                .firstElement(),
+            getExpCtx()),
+        AssertionException,
+        ErrorCodes::FailedToParse);
+    ASSERT_THROWS_CODE(
+        DocumentSourceUnionWith::createFromBson(
+            BSON("$unionWith" << BSON("coll"
+                                      << "foo"
+                                      << "pipeline"
+                                      << BSON_ARRAY(BSON("$addFields" << BSON("a" << 3))) << "myDog"
+                                      << "bar"))
+                .firstElement(),
+            getExpCtx()),
+        AssertionException,
+        ErrorCodes::FailedToParse);
+    ASSERT_THROWS_CODE(
+        DocumentSourceUnionWith::createFromBson(
+            BSON("$unionWith" << BSON("coll"
+                                      << "foo"
+                                      << "pipeline"
+                                      << BSON_ARRAY(BSON("$petMyDog" << BSON("myDog" << 3)))))
+                .firstElement(),
+            getExpCtx()),
+        AssertionException,
+        16436);
+}
+
+TEST_F(DocumentSourceUnionWithTest, PropagatePauses) {
+    const auto mock =
+        DocumentSourceMock::createForTest({Document(),
+                                           DocumentSource::GetNextResult::makePauseExecution(),
+                                           Document(),
+                                           DocumentSource::GetNextResult::makePauseExecution()});
+    const auto mockDequeOne = std::deque<DocumentSource::GetNextResult>{};
+    const auto mockDequeTwo = std::deque<DocumentSource::GetNextResult>{};
+    const auto mockCtxOne = getExpCtx()->copyWith({});
+    mockCtxOne->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeOne);
+    const auto mockCtxTwo = getExpCtx()->copyWith({});
+    mockCtxTwo->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeTwo);
+    auto unionWithOne = DocumentSourceUnionWith(
+        mockCtxOne,
+        uassertStatusOK(
+            Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx())));
+    auto unionWithTwo = DocumentSourceUnionWith(
+        mockCtxTwo,
+        uassertStatusOK(
+            Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx())));
+    unionWithOne.setSource(mock.get());
+    unionWithTwo.setSource(&unionWithOne);
+
+    ASSERT_TRUE(unionWithTwo.getNext().isAdvanced());
+    ASSERT_TRUE(unionWithTwo.getNext().isPaused());
+    ASSERT_TRUE(unionWithTwo.getNext().isAdvanced());
+    ASSERT_TRUE(unionWithTwo.getNext().isPaused());
+
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
+}
+
+}  // namespace
+}  // namespace mongo