SERVER-31072 reorder $limit before $lookup in pipeline if no $unwind is present

4 files changed, 13 insertions, 5 deletions

diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index e2e321735c5..c30c3a9d6c7 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -203,6 +203,7 @@ StageConstraints DocumentSourceLookUp::constraints(Pipeline::SplitState) const {
                                  LookupRequirement::kAllowed);
 
     constraints.canSwapWithMatch = true;
+    constraints.canSwapWithLimitAndSample = !_unwindSrc;
     return constraints;
 }
 
diff --git a/src/mongo/db/pipeline/pipeline.cpp b/src/mongo/db/pipeline/pipeline.cpp
index bbe003e036e..4990e916245 100644
--- a/src/mongo/db/pipeline/pipeline.cpp
+++ b/src/mongo/db/pipeline/pipeline.cpp
@@ -70,11 +70,6 @@ Value appendExecStats(Value docSource, const CommonStats& stats) {
 }
 }  // namespace
 
-/**
- * Enabling the disablePipelineOptimization fail point will stop the aggregate command from
- * attempting to optimize the pipeline or the pipeline stages. Neither DocumentSource::optimizeAt()
- * nor DocumentSource::optimize() will be attempted.
- */
 MONGO_FAIL_POINT_DEFINE(disablePipelineOptimization);
 
 using boost::intrusive_ptr;
diff --git a/src/mongo/db/pipeline/pipeline.h b/src/mongo/db/pipeline/pipeline.h
index 7c32a51c197..aa08f8d98b8 100644
--- a/src/mongo/db/pipeline/pipeline.h
+++ b/src/mongo/db/pipeline/pipeline.h
@@ -57,6 +57,13 @@ class OperationContext;
 class PipelineDeleter;
 
 /**
+ * Enabling the disablePipelineOptimization fail point will stop the aggregate command from
+ * attempting to optimize the pipeline or the pipeline stages. Neither DocumentSource::optimizeAt()
+ * nor DocumentSource::optimize() will be attempted.
+ */
+extern FailPoint disablePipelineOptimization;
+
+/**
  * A Pipeline object represents a list of DocumentSources and is responsible for optimizing the
  * pipeline.
  */
diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index d97a7d7f47e..1590410d5d2 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -450,6 +450,11 @@ getSortAndGroupStagesFromPipeline(const Pipeline::SourceContainer& sources) {
 }
 
 boost::optional<long long> extractLimitForPushdown(Pipeline* pipeline) {
+    // If the disablePipelineOptimization failpoint is enabled, then do not attempt the limit
+    // pushdown optimization.
+    if (MONGO_unlikely(disablePipelineOptimization.shouldFail())) {
+        return boost::none;
+    }
     auto&& sources = pipeline->getSources();
     auto limit = DocumentSourceSort::extractLimitForPushdown(sources.begin(), &sources);
     if (limit) {