SERVER-53977 Allow rewrite to produce only empty documents in $unpackBucket if it is followed by $count

3 files changed, 93 insertions, 1 deletions

diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index bcec0c46033..059f669c39c 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -398,6 +398,7 @@ env.CppUnitTest(
         'document_source_union_with_test.cpp',
         'document_source_internal_unpack_bucket_test/extract_or_build_project_to_internalize_test.cpp',
         'document_source_internal_unpack_bucket_test/create_predicates_on_bucket_level_field_test.cpp',
+        'document_source_internal_unpack_bucket_test/group_reorder_test.cpp',
         'document_source_internal_unpack_bucket_test/internalize_project_test.cpp',
         'document_source_internal_unpack_bucket_test/sort_reorder_test.cpp',
         'document_source_internal_unpack_bucket_test/unpack_bucket_exec_test.cpp',
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 d17840698d9..2c064194a01 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
@@ -38,6 +38,7 @@
 #include "mongo/db/matcher/expression.h"
 #include "mongo/db/matcher/expression_algo.h"
 #include "mongo/db/matcher/expression_internal_expr_comparison.h"
+#include "mongo/db/pipeline/document_source_group.h"
 #include "mongo/db/pipeline/document_source_match.h"
 #include "mongo/db/pipeline/document_source_project.h"
 #include "mongo/db/pipeline/document_source_sample.h"
@@ -297,7 +298,7 @@ void BucketUnpacker::reset(BSONObj&& bucket) {
         // Includes a field when '_unpackerBehavior' is 'kInclude' and it's found in 'fieldSet' or
         // _unpackerBehavior is 'kExclude' and it's not found in 'fieldSet'.
         if (determineIncludeField(colName, _unpackerBehavior, _spec)) {
-            _fieldIters.push_back({colName.toString(), BSONObjIterator{elem.Obj()}});
+            _fieldIters.emplace_back(colName.toString(), BSONObjIterator{elem.Obj()});
         }
     }
 
@@ -767,6 +768,21 @@ Pipeline::SourceContainer::iterator DocumentSourceInternalUnpackBucket::doOptimi
     // Optimize the pipeline after the $unpackBucket.
     optimizeEndOfPipeline(itr, container);
 
+    {
+        // Check if the rest of the pipeline needs any fields. For example we might only be
+        // interested in $count.
+        auto deps = Pipeline::getDependenciesForContainer(
+            pExpCtx, Pipeline::SourceContainer{std::next(itr), container->end()}, boost::none);
+        if (deps.hasNoRequirements()) {
+            _bucketUnpacker.setBucketSpecAndBehavior({_bucketUnpacker.bucketSpec().timeField,
+                                                      _bucketUnpacker.bucketSpec().metaField,
+                                                      {}},
+                                                     BucketUnpacker::Behavior::kInclude);
+
+            // Keep going for next optimization.
+        }
+    }
+
     if (auto nextMatch = dynamic_cast<DocumentSourceMatch*>((*std::next(itr)).get())) {
         // Attempt to push predicates on the metaField past $_internalUnpackBucket.
         auto [metaMatch, remainingMatch] = splitMatchOnMetaAndRename(nextMatch);
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/group_reorder_test.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/group_reorder_test.cpp
new file mode 100644
index 00000000000..b07dee1982e
--- /dev/null
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/group_reorder_test.cpp
@@ -0,0 +1,75 @@
+/**
+ *    Copyright (C) 2020-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/aggregation_context_fixture.h"
+#include "mongo/db/query/util/make_data_structure.h"
+
+namespace mongo {
+namespace {
+
+using InternalUnpackBucketGroupReorder = AggregationContextFixture;
+
+TEST_F(InternalUnpackBucketGroupReorder, OptimizeForCount) {
+    auto unpackSpecObj = fromjson(
+        "{$_internalUnpackBucket: { include: ['a', 'b', 'c'], metaField: 'meta', timeField: 't'}}");
+    auto countSpecObj = fromjson("{$count: 'foo'}");
+
+    auto pipeline = Pipeline::parse(makeVector(unpackSpecObj, countSpecObj), getExpCtx());
+    pipeline->optimizePipeline();
+
+    auto serialized = pipeline->serializeToBson();
+    // $count gets rewritten to $group + $project.
+    ASSERT_EQ(3, serialized.size());
+
+    auto optimized =
+        fromjson("{$_internalUnpackBucket: { include: [], timeField: 't', metaField: 'meta'}}");
+    ASSERT_BSONOBJ_EQ(optimized, serialized[0]);
+}
+
+TEST_F(InternalUnpackBucketGroupReorder, OptimizeForCountNegative) {
+    auto unpackSpecObj = fromjson(
+        "{$_internalUnpackBucket: { include: ['a', 'b', 'c'], metaField: 'meta', timeField: 't'}}");
+    auto groupSpecObj = fromjson("{$group: {_id: '$a', s: {$sum: '$b'}}}");
+
+    auto pipeline = Pipeline::parse(makeVector(unpackSpecObj, groupSpecObj), getExpCtx());
+    pipeline->optimizePipeline();
+
+    auto serialized = pipeline->serializeToBson();
+    ASSERT_EQ(2, serialized.size());
+
+    // We do not get the reorder since we are grouping on a field.
+    auto optimized = fromjson(
+        "{$_internalUnpackBucket: { include: ['a', 'b', 'c'], timeField: 't', metaField: 'meta'}}");
+    ASSERT_BSONOBJ_EQ(optimized, serialized[0]);
+}
+
+}  // namespace
+}  // namespace mongo