SERVER-35905 Add logic to detect when $out is eligible for an exchange

1 files changed, 210 insertions, 0 deletions

diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp
index 65d3c4fecc0..280473848cc 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -45,10 +45,12 @@
 #include "mongo/db/pipeline/document_source_out.h"
 #include "mongo/db/pipeline/document_source_project.h"
 #include "mongo/db/pipeline/document_source_sort.h"
+#include "mongo/db/pipeline/document_source_test_optimizations.h"
 #include "mongo/db/pipeline/document_value_test_util.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
 #include "mongo/db/pipeline/field_path.h"
 #include "mongo/db/pipeline/pipeline.h"
+#include "mongo/db/pipeline/stub_mongo_process_interface.h"
 #include "mongo/db/query/collation/collator_interface_mock.h"
 #include "mongo/db/query/query_test_service_context.h"
 #include "mongo/db/repl/replication_coordinator_mock.h"
@@ -2246,11 +2248,24 @@ DEATH_TEST_F(PipelineMustRunOnMongoSTest,
     splitPipeline.shardsPipeline->requiredToRunOnMongos();
 }
 
+/**
+ * For the purpsoses of this test, assume every collection is unsharded. Stages may ask this during
+ * setup. For example, to compute its constraints, the $out stage needs to know if the output
+ * collection is sharded.
+ */
+class FakeMongoProcessInterface : public StubMongoProcessInterface {
+public:
+    bool isSharded(OperationContext* opCtx, const NamespaceString& ns) override {
+        return false;
+    }
+};
+
 TEST_F(PipelineMustRunOnMongoSTest, SplitMongoSMergePipelineAssertsIfShardStagePresent) {
     auto expCtx = getExpCtx();
 
     expCtx->allowDiskUse = true;
     expCtx->inMongos = true;
+    expCtx->mongoProcessInterface = std::make_shared<FakeMongoProcessInterface>();
 
     auto match = DocumentSourceMatch::create(fromjson("{x: 5}"), expCtx);
     auto split = DocumentSourceInternalSplitPipeline::create(expCtx, HostTypeRequirement::kNone);
@@ -2678,6 +2693,201 @@ TEST_F(PipelineDependenciesTest, ShouldNotRequireTextScoreIfAvailableButDefinite
 
 }  // namespace Dependencies
 
+namespace {
+TEST(PipelineRenameTracking, ReportsIdentityMapWhenEmpty) {
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest());
+    auto pipeline = unittest::assertGet(Pipeline::create({DocumentSourceMock::create()}, expCtx));
+    auto renames = pipeline->renamedPaths({"a", "b", "c.d"});
+    ASSERT(static_cast<bool>(renames));
+    auto nameMap = *renames;
+    ASSERT_EQ(nameMap.size(), 3UL);
+    ASSERT_EQ(nameMap["a"], "a");
+    ASSERT_EQ(nameMap["b"], "b");
+    ASSERT_EQ(nameMap["c.d"], "c.d");
+}
+
+class NoModifications : public DocumentSourceTestOptimizations {
+public:
+    NoModifications() : DocumentSourceTestOptimizations() {}
+    static boost::intrusive_ptr<NoModifications> create() {
+        return new NoModifications();
+    }
+
+    /**
+     * Returns a description which communicate that this stage modifies nothing.
+     */
+    GetModPathsReturn getModifiedPaths() const final {
+        return {GetModPathsReturn::Type::kFiniteSet, std::set<std::string>(), {}};
+    }
+};
+
+TEST(PipelineRenameTracking, ReportsIdentityWhenNoStageModifiesAnything) {
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest());
+    {
+        auto pipeline = unittest::assertGet(
+            Pipeline::create({DocumentSourceMock::create(), NoModifications::create()}, expCtx));
+        auto renames = pipeline->renamedPaths({"a", "b", "c.d"});
+        ASSERT(static_cast<bool>(renames));
+        auto nameMap = *renames;
+        ASSERT_EQ(nameMap.size(), 3UL);
+        ASSERT_EQ(nameMap["a"], "a");
+        ASSERT_EQ(nameMap["b"], "b");
+        ASSERT_EQ(nameMap["c.d"], "c.d");
+    }
+    {
+        auto pipeline = unittest::assertGet(Pipeline::create({DocumentSourceMock::create(),
+                                                              NoModifications::create(),
+                                                              NoModifications::create(),
+                                                              NoModifications::create()},
+                                                             expCtx));
+        auto renames = pipeline->renamedPaths({"a", "b", "c.d"});
+        ASSERT(static_cast<bool>(renames));
+        auto nameMap = *renames;
+        ASSERT_EQ(nameMap.size(), 3UL);
+        ASSERT_EQ(nameMap["a"], "a");
+        ASSERT_EQ(nameMap["b"], "b");
+        ASSERT_EQ(nameMap["c.d"], "c.d");
+    }
+}
+
+class NotSupported : public DocumentSourceTestOptimizations {
+public:
+    NotSupported() : DocumentSourceTestOptimizations() {}
+    static boost::intrusive_ptr<NotSupported> create() {
+        return new NotSupported();
+    }
+
+    /**
+     * Returns a description which communicate that this stage modifies nothing.
+     */
+    GetModPathsReturn getModifiedPaths() const final {
+        return {GetModPathsReturn::Type::kNotSupported, std::set<std::string>(), {}};
+    }
+};
+
+TEST(PipelineRenameTracking, DoesNotReportRenamesIfAStageDoesNotSupportTrackingThem) {
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest());
+    auto pipeline = unittest::assertGet(Pipeline::create({DocumentSourceMock::create(),
+                                                          NoModifications::create(),
+                                                          NotSupported::create(),
+                                                          NoModifications::create()},
+                                                         expCtx));
+    ASSERT_FALSE(static_cast<bool>(pipeline->renamedPaths({"a"})));
+    ASSERT_FALSE(static_cast<bool>(pipeline->renamedPaths({"a", "b"})));
+    ASSERT_FALSE(static_cast<bool>(pipeline->renamedPaths({"x", "yahoo", "c.d"})));
+}
+
+class RenamesAToB : public DocumentSourceTestOptimizations {
+public:
+    RenamesAToB() : DocumentSourceTestOptimizations() {}
+    static boost::intrusive_ptr<RenamesAToB> create() {
+        return new RenamesAToB();
+    }
+    GetModPathsReturn getModifiedPaths() const final {
+        return {GetModPathsReturn::Type::kFiniteSet, std::set<std::string>{}, {{"b", "a"}}};
+    }
+};
+
+TEST(PipelineRenameTracking, ReportsNewNamesWhenSingleStageRenames) {
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest());
+    auto pipeline = unittest::assertGet(
+        Pipeline::create({DocumentSourceMock::create(), RenamesAToB::create()}, expCtx));
+    {
+        auto renames = pipeline->renamedPaths({"b"});
+        ASSERT(static_cast<bool>(renames));
+        auto nameMap = *renames;
+        ASSERT_EQ(nameMap.size(), 1UL);
+        ASSERT_EQ(nameMap["b"], "a");
+    }
+    {
+        auto renames = pipeline->renamedPaths({"b", "c.d"});
+        ASSERT(static_cast<bool>(renames));
+        auto nameMap = *renames;
+        ASSERT_EQ(nameMap.size(), 2UL);
+        ASSERT_EQ(nameMap["b"], "a");
+        ASSERT_EQ(nameMap["c.d"], "c.d");
+    }
+    {
+        // This is strange; the mock stage reports to essentially duplicate the "a" field into "b".
+        // Because of this, both "b" and "a" should map to "a".
+        auto renames = pipeline->renamedPaths({"b", "a"});
+        ASSERT(static_cast<bool>(renames));
+        auto nameMap = *renames;
+        ASSERT_EQ(nameMap.size(), 2UL);
+        ASSERT_EQ(nameMap["b"], "a");
+        ASSERT_EQ(nameMap["a"], "a");
+    }
+}
+
+TEST(PipelineRenameTracking, ReportsIdentityMapWhenGivenEmptyIteratorRange) {
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest());
+    auto pipeline = unittest::assertGet(
+        Pipeline::create({DocumentSourceMock::create(), RenamesAToB::create()}, expCtx));
+    {
+        auto renames = Pipeline::renamedPaths(
+            pipeline->getSources().rbegin(), pipeline->getSources().rbegin(), {"b"});
+        ASSERT(static_cast<bool>(renames));
+        auto nameMap = *renames;
+        ASSERT_EQ(nameMap.size(), 1UL);
+        ASSERT_EQ(nameMap["b"], "b");
+    }
+    {
+        auto renames = Pipeline::renamedPaths(
+            pipeline->getSources().rbegin(), pipeline->getSources().rbegin(), {"b", "c.d"});
+        ASSERT(static_cast<bool>(renames));
+        auto nameMap = *renames;
+        ASSERT_EQ(nameMap.size(), 2UL);
+        ASSERT_EQ(nameMap["b"], "b");
+        ASSERT_EQ(nameMap["c.d"], "c.d");
+    }
+}
+
+class RenamesBToC : public DocumentSourceTestOptimizations {
+public:
+    RenamesBToC() : DocumentSourceTestOptimizations() {}
+    static boost::intrusive_ptr<RenamesBToC> create() {
+        return new RenamesBToC();
+    }
+    GetModPathsReturn getModifiedPaths() const final {
+        return {GetModPathsReturn::Type::kFiniteSet, std::set<std::string>{}, {{"c", "b"}}};
+    }
+};
+
+TEST(PipelineRenameTracking, ReportsNewNameAcrossMultipleRenames) {
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest());
+    auto pipeline = unittest::assertGet(Pipeline::create(
+        {DocumentSourceMock::create(), RenamesAToB::create(), RenamesBToC::create()}, expCtx));
+    auto renames = pipeline->renamedPaths({"c"});
+    ASSERT(static_cast<bool>(renames));
+    auto nameMap = *renames;
+    ASSERT_EQ(nameMap.size(), 1UL);
+    ASSERT_EQ(nameMap["c"], "a");
+}
+
+class RenamesBToA : public DocumentSourceTestOptimizations {
+public:
+    RenamesBToA() : DocumentSourceTestOptimizations() {}
+    static boost::intrusive_ptr<RenamesBToA> create() {
+        return new RenamesBToA();
+    }
+    GetModPathsReturn getModifiedPaths() const final {
+        return {GetModPathsReturn::Type::kFiniteSet, std::set<std::string>{}, {{"b", "a"}}};
+    }
+};
+
+TEST(PipelineRenameTracking, CanHandleBackAndForthRename) {
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest());
+    auto pipeline = unittest::assertGet(Pipeline::create(
+        {DocumentSourceMock::create(), RenamesAToB::create(), RenamesBToA::create()}, expCtx));
+    auto renames = pipeline->renamedPaths({"a"});
+    ASSERT(static_cast<bool>(renames));
+    auto nameMap = *renames;
+    ASSERT_EQ(nameMap.size(), 1UL);
+    ASSERT_EQ(nameMap["a"], "a");
+}
+
+}  // namespace
+
 class All : public Suite {
 public:
     All() : Suite("PipelineOptimizations") {}