SERVER-83490 Do not push to SBE $addFields stage created as part of $setWindowFields

Author: fotiniAlvanaki

src/mongo/db/pipeline/document_source_add_fields.cpp

@@ -81,14 +81,17 @@ intrusive_ptr<DocumentSource> DocumentSourceAddFields::create(
 intrusive_ptr<DocumentSource> DocumentSourceAddFields::create(
     const FieldPath& fieldPath,
     const intrusive_ptr<Expression>& expr,
-    const intrusive_ptr<ExpressionContext>& expCtx) {
+    const intrusive_ptr<ExpressionContext>& expCtx,
+    bool createdBySetWindowFields) {
 
     const bool isIndependentOfAnyCollection = false;
-    return make_intrusive<DocumentSourceSingleDocumentTransformation>(
+    auto docSrc = make_intrusive<DocumentSourceSingleDocumentTransformation>(
         expCtx,
         projection_executor::AddFieldsProjectionExecutor::create(expCtx, fieldPath, expr),
         kStageName,
         isIndependentOfAnyCollection);
+    docSrc->setCreatedBySetWindowFields(createdBySetWindowFields);
+    return docSrc;
 }
 
 intrusive_ptr<DocumentSource> DocumentSourceAddFields::createFromBson(

src/mongo/db/pipeline/document_source_add_fields.h

@@ -69,7 +69,8 @@ class DocumentSourceAddFields final {
     static boost::intrusive_ptr<DocumentSource> create(
         const FieldPath& fieldPath,
         const boost::intrusive_ptr<Expression>& expr,
-        const boost::intrusive_ptr<ExpressionContext>& expCtx);
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        bool createdBySetWindowFields);
 
     /**
      * Parses a $addFields stage from the user-supplied BSON.

src/mongo/db/pipeline/document_source_set_window_fields.cpp

@@ -244,7 +244,7 @@ list<intrusive_ptr<DocumentSource>> document_source_set_window_fields::create(
     // $set
     if (complexPartitionBy) {
         result.push_back(
-            DocumentSourceAddFields::create(*simplePartitionBy, *complexPartitionBy, expCtx));
+            DocumentSourceAddFields::create(*simplePartitionBy, *complexPartitionBy, expCtx, true));
     }
 
     // $sort

src/mongo/db/pipeline/document_source_single_document_transformation.h

@@ -163,6 +163,14 @@ class DocumentSourceSingleDocumentTransformation final : public DocumentSource {
                                                                                          newName);
     }
 
+    void setCreatedBySetWindowFields(bool val) {
+        createdBySetWindowFields = val;
+    }
+
+    bool isCreatedBySetWindowFields() const {
+        return createdBySetWindowFields;
+    }
+
 protected:
     GetNextResult doGetNext() final;
     void doDispose() final;
@@ -182,6 +190,11 @@ class DocumentSourceSingleDocumentTransformation final : public DocumentSource {
     // Cached stage options in case this DocumentSource is disposed before serialized (e.g. explain
     // with a sort which will auto-dispose of the pipeline).
     Document _cachedStageOptions;
+
+    // We set the createdBySetWindowFields flag to prevent an addFields stage from being pushed to
+    // SBE. TODO (SERVER-75103) : Once setWindowFields has been pushed to SBE, this should be
+    // removed.
+    bool createdBySetWindowFields = false;
 };
 
 }  // namespace mongo

src/mongo/db/pipeline/pipeline_d.cpp

@@ -295,6 +295,13 @@ bool pushDownPipelineStageIfCompatible(
         return true;
     } else if (auto transformStage =
                    dynamic_cast<DocumentSourceSingleDocumentTransformation*>(stage.get())) {
+        // We do not push to SBE an addFields that has been created as part of a setWindowFields
+        // stage because it causes a performance regression. TODO (SERVER-75103) : Once
+        // setWindowFields has been pushed to SBE, this should be removed.
+        if (transformStage->isCreatedBySetWindowFields()) {
+            return false;
+        }
+
         if (!allowedStages.transform) {
             return false;
         }