From d1a117a1abc9ead34fc6399697055a46ca0c1df4 Mon Sep 17 00:00:00 2001
From: Nicholas Zolnierz <nicholas.zolnierz@mongodb.com>
Date: Thu, 15 Sep 2022 15:11:22 +0000
Subject: SERVER-69380 Avoid swapping meta projection with
 $_internalUnpackBucket if a projection has already been absorbed

---
 jstests/core/timeseries/timeseries_project.js      |  6 +++
 .../document_source_internal_unpack_bucket.cpp     |  2 +-
 .../optimize_pipeline_test.cpp                     | 49 ++++++++++++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/jstests/core/timeseries/timeseries_project.js b/jstests/core/timeseries/timeseries_project.js
index b9ba914130a..b5d615ebde5 100644
--- a/jstests/core/timeseries/timeseries_project.js
+++ b/jstests/core/timeseries/timeseries_project.js
@@ -4,6 +4,7 @@
  * @tags: [
  *   # We need a timeseries collection.
  *   requires_timeseries,
+ *   requires_fcv_62,
  * ]
  */
 (function() {
@@ -111,5 +112,10 @@ pipeline = [{$addFields: {a: "$x", b: "$a"}}, {$project: {_id: 0}}];
 tsDoc = tsColl.aggregate(pipeline).toArray();
 regDoc = regColl.aggregate(pipeline).toArray();
 assert.docEq(tsDoc, regDoc);
+
+pipeline = [{$project: {a: 1, _id: 0}}, {$project: {newMeta: "$x"}}];
+tsDoc = tsColl.aggregate(pipeline).toArray();
+regDoc = regColl.aggregate(pipeline).toArray();
+assert.docEq(tsDoc, regDoc);
 })();
 })();
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
index d4d38a2b2be..a950aeee4f4 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
@@ -514,7 +514,7 @@ bool DocumentSourceInternalUnpackBucket::pushDownComputedMetaProjection(
     if (std::next(itr) == container->end()) {
         return nextStageWasRemoved;
     }
-    if (!_bucketUnpacker.bucketSpec().metaField()) {
+    if (!_bucketUnpacker.getMetaField() || !_bucketUnpacker.includeMetaField()) {
         return nextStageWasRemoved;
     }
 
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/optimize_pipeline_test.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/optimize_pipeline_test.cpp
index cfe92c09f5f..2001cb2cab3 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/optimize_pipeline_test.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/optimize_pipeline_test.cpp
@@ -750,6 +750,55 @@ TEST_F(OptimizePipeline, InternalizeProjectAndPushdownAddFields) {
                       serialized[1]);
 }
 
+TEST_F(OptimizePipeline, DoNotSwapAddFieldsIfDependencyIsExcluded) {
+    {
+        auto unpackSpecObj = fromjson(
+            "{$_internalUnpackBucket: { exclude: [], timeField: 'time', metaField: 'myMeta', "
+            "bucketMaxSpanSeconds: 3600}}");
+        auto projectSpecObj = fromjson("{$project: {x: true, _id: false}}");
+        auto addFieldsSpec = fromjson("{$addFields: {newMeta: '$myMeta'}}");
+
+        auto pipeline =
+            Pipeline::parse(makeVector(unpackSpecObj, projectSpecObj, addFieldsSpec), getExpCtx());
+
+        pipeline->optimizePipeline();
+
+        // We should internalize the $project but _not_ push down the $addFields because it's field
+        // dependency has been excluded. Theoretically we could remove the $addFields for this
+        // trivial except but not always.
+        auto serialized = pipeline->serializeToBson();
+        ASSERT_EQ(2u, serialized.size());
+        ASSERT_BSONOBJ_EQ(fromjson("{$_internalUnpackBucket: { include: ['x'], timeField: 'time', "
+                                   "metaField: 'myMeta', bucketMaxSpanSeconds: 3600}}"),
+                          serialized[0]);
+        ASSERT_BSONOBJ_EQ(fromjson("{$addFields: {newMeta: '$myMeta'}}"), serialized[1]);
+    }
+
+    // Similar test except the dependency is on an excluded non-meta field.
+    {
+        auto unpackSpecObj = fromjson(
+            "{$_internalUnpackBucket: { exclude: [], timeField: 'time', metaField: 'myMeta', "
+            "bucketMaxSpanSeconds: 3600}}");
+        auto projectSpecObj = fromjson("{$project: {x: true, _id: false}}");
+        auto addFieldsSpec = fromjson("{$addFields: {newMeta: '$excluded'}}");
+
+        auto pipeline =
+            Pipeline::parse(makeVector(unpackSpecObj, projectSpecObj, addFieldsSpec), getExpCtx());
+
+        pipeline->optimizePipeline();
+
+        // We should internalize the $project but _not_ push down the $addFields because it's field
+        // dependency has been excluded. Theoretically we could remove the $addFields for this
+        // trivial except but not always.
+        auto serialized = pipeline->serializeToBson();
+        ASSERT_EQ(2u, serialized.size());
+        ASSERT_BSONOBJ_EQ(fromjson("{$_internalUnpackBucket: { include: ['x'], timeField: 'time', "
+                                   "metaField: 'myMeta', bucketMaxSpanSeconds: 3600}}"),
+                          serialized[0]);
+        ASSERT_BSONOBJ_EQ(fromjson("{$addFields: {newMeta: '$excluded'}}"), serialized[1]);
+    }
+}
+
 TEST_F(OptimizePipeline, PushdownSortAndAddFields) {
     auto unpackSpecObj = fromjson(
         "{$_internalUnpackBucket: { exclude: [], timeField: 'time', metaField: 'myMeta', "
-- 
cgit v1.2.1