diff options
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 |