SERVER-73253: Better path tracking when renaming nested/compound grouping fields -- updated

mattsimply · Evergreen Agent · commit 288a08f4c02b · 2023-06-27T02:55:54.000Z
diff --git a/jstests/aggregation/sources/group/grouped_match_push_down.js b/jstests/aggregation/sources/group/grouped_match_push_down.js
@@ -0,0 +1,72 @@
+// Tests optimization rule for pushing match stages over group stages
+
+(function() {
+"use strict";
+
+load("jstests/libs/analyze_plan.js");  // For 'aggPlanHasStage' and other explain helpers.
+
+const coll = db.grouped_match_push_down;
+coll.drop();
+assert.commandWorked(coll.insert({_id: 1, x: 10}));
+assert.commandWorked(coll.insert({_id: 2, x: 20}));
+assert.commandWorked(coll.insert({_id: 3, x: 30}));
+
+// Asserts end-to-end the optimization of group-project-match stage sequence which includes a rename
+// over a dotted path. It evaluates the correctness of the result, as well as whether the
+// optimization pushed the predicate down before the aggregation.
+function assertOptimizeMatchRenameAggregationPipelineWithDottedRename(
+    {pipeline, expectedStages, expectedResult} = {}) {
+    const explainOutput = coll.explain().aggregate(pipeline);
+
+    if (expectedStages) {
+        for (let expectedStage of expectedStages) {
+            assert(aggPlanHasStage(explainOutput, expectedStage), explainOutput);
+        }
+    }
+
+    if (expectedResult) {
+        const actualResult = coll.aggregate(pipeline).toArray();
+        assert.sameMembers(expectedResult, actualResult);
+    }
+}
+
+assert.commandWorked(coll.insert({_id: 20, d: 2}));
+// Assert that a sequence of stages group, project, match over a rename on a dotted path (depth 3)
+// will push the predicate before the group stage.
+assertOptimizeMatchRenameAggregationPipelineWithDottedRename({
+    pipeline: [
+        {$group: {_id: {c: '$d'}, c: {$sum: {$const: 1}}}},
+        {$project: {m: '$_id.c'}},
+        {$match: {m: {$eq: 2}}}
+    ],
+    expectedStages: ["GROUP"],
+    expectedResult: [{"_id": {"c": 2}, "m": 2}]
+});
+
+// Assert that the optimization over group, project, match over a renamed dotted path will not
+// push down the predicate over multiple projection stages.
+assertOptimizeMatchRenameAggregationPipelineWithDottedRename({
+    pipeline: [
+        {$group: {_id: {c: '$d'}, c: {$sum: {$const: 1}}}},
+        {$project: {m: '$_id.c'}},
+        {$project: {m2: '$m'}},
+        {$match: {m2: {$eq: 2}}}
+    ],
+    // the optimization will not push the filter, thus, the query will perform a COLLSCAN. The
+    // system does not have an index to use.
+    expectedStages: ["COLLSCAN"],
+    expectedResult: [{"_id": {"c": 2}, "m2": 2}]
+});
+
+// Assert that the optimization over group, project, match over a renamed dotted path will not
+// push down the predicate when the rename stage renames a dotted path with depth more than 3.
+assertOptimizeMatchRenameAggregationPipelineWithDottedRename({
+    pipeline: [
+        {$group: {_id: {c: {d: '$d'}}, c: {$sum: {$const: 1}}}},
+        {$project: {m: '$_id.c.d'}},
+        {$match: {m: {$eq: 2}}}
+    ],
+    expectedStages: ["COLLSCAN"],
+    expectedResult: [{"_id": {"c": {"d": 2}}, "m": 2}]
+});
+}());
diff --git a/src/mongo/db/exec/add_fields_projection_executor.h b/src/mongo/db/exec/add_fields_projection_executor.h
@@ -140,10 +140,12 @@ class AddFieldsProjectionExecutor : public ProjectionExecutor {
     DocumentSource::GetModPathsReturn getModifiedPaths() const final {
         OrderedPathSet computedPaths;
         StringMap<std::string> renamedPaths;
-        _root->reportComputedPaths(&computedPaths, &renamedPaths);
+        StringMap<std::string> complexRenamedPaths;
+        _root->reportComputedPaths(&computedPaths, &renamedPaths, &complexRenamedPaths);
         return {DocumentSource::GetModPathsReturn::Type::kFiniteSet,
                 std::move(computedPaths),
-                std::move(renamedPaths)};
+                std::move(renamedPaths),
+                std::move(complexRenamedPaths)};
     }
 
     /**
diff --git a/src/mongo/db/exec/inclusion_projection_executor.h b/src/mongo/db/exec/inclusion_projection_executor.h
@@ -304,11 +304,13 @@ class InclusionProjectionExecutor : public ProjectionExecutor {
 
         OrderedPathSet computedPaths;
         StringMap<std::string> renamedPaths;
-        _root->reportComputedPaths(&computedPaths, &renamedPaths);
+        StringMap<std::string> complexRenamedPaths;
+        _root->reportComputedPaths(&computedPaths, &renamedPaths, &complexRenamedPaths);
 
         return {DocumentSource::GetModPathsReturn::Type::kAllExcept,
                 std::move(preservedPaths),
-                std::move(renamedPaths)};
+                std::move(renamedPaths),
+                std::move(complexRenamedPaths)};
     }
 
     /**
diff --git a/src/mongo/db/exec/projection_node.cpp b/src/mongo/db/exec/projection_node.cpp
@@ -259,7 +259,8 @@ void ProjectionNode::reportProjectedPaths(OrderedPathSet* projectedPaths) const
 }
 
 void ProjectionNode::reportComputedPaths(OrderedPathSet* computedPaths,
-                                         StringMap<std::string>* renamedPaths) const {
+                                         StringMap<std::string>* renamedPaths,
+                                         StringMap<std::string>* complexRenamedPaths) const {
     for (auto&& computedPair : _expressions) {
         // The expression's path is the concatenation of the path to this node, plus the field name
         // associated with the expression.
@@ -270,9 +271,15 @@ void ProjectionNode::reportComputedPaths(OrderedPathSet* computedPaths,
         for (auto&& rename : exprComputedPaths.renames) {
             (*renamedPaths)[rename.first] = rename.second;
         }
+
+        if (complexRenamedPaths) {
+            for (auto&& complexRename : exprComputedPaths.complexRenames) {
+                (*complexRenamedPaths)[complexRename.first] = complexRename.second;
+            }
+        }
     }
     for (auto&& childPair : _children) {
-        childPair.second->reportComputedPaths(computedPaths, renamedPaths);
+        childPair.second->reportComputedPaths(computedPaths, renamedPaths, complexRenamedPaths);
     }
 }
 
diff --git a/src/mongo/db/exec/projection_node.h b/src/mongo/db/exec/projection_node.h
@@ -139,10 +139,14 @@ class ProjectionNode {
      *
      * Computed paths that are identified as the result of a simple rename are instead filled out in
      * 'renamedPaths'. Each entry in 'renamedPaths' maps from the path's new name to its old name
-     * prior to application of this projection.
+     * prior to application of this projection. 'complexRenamedPaths' is an optional parameter that
+     * acts exactly as the 'renamedPaths' map and includes renames whose old name includes dotted
+     * paths (Note: the dotted path renames are constrained to length 3). The paths that are
+     * included in 'complexRenamedPaths' are also included in 'computedPaths'.
      */
     void reportComputedPaths(OrderedPathSet* computedPaths,
-                             StringMap<std::string>* renamedPaths) const;
+                             StringMap<std::string>* renamedPaths,
+                             StringMap<std::string>* complexRenamedPaths = nullptr) const;
 
     const std::string& getPath() const {
         return _pathToNode;
diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
@@ -643,8 +643,14 @@ class DocumentSource : public RefCountable {
             kAllExcept,
         };
 
-        GetModPathsReturn(Type type, OrderedPathSet&& paths, StringMap<std::string>&& renames)
-            : type(type), paths(std::move(paths)), renames(std::move(renames)) {}
+        GetModPathsReturn(Type type,
+                          OrderedPathSet&& paths,
+                          StringMap<std::string>&& renames,
+                          StringMap<std::string>&& complexRenames = {})
+            : type(type),
+              paths(std::move(paths)),
+              renames(std::move(renames)),
+              complexRenames(std::move(complexRenames)) {}
 
         std::set<std::string> getNewNames() {
             std::set<std::string> newNames;
@@ -706,6 +712,11 @@ class DocumentSource : public RefCountable {
         // This stage should return kAllExcept, since it modifies all paths other than "a". It can
         // also fill out 'renames' with the mapping "b" => "c".
         StringMap<std::string> renames;
+
+        // Including space for returning renames which include dotted paths.
+        // i.e., "a.b" => c
+        //
+        StringMap<std::string> complexRenames;
     };
 
     /**
diff --git a/src/mongo/db/pipeline/document_source_group.cpp b/src/mongo/db/pipeline/document_source_group.cpp
@@ -45,6 +45,10 @@
 #include "mongo/util/assert_util.h"
 #include "mongo/util/intrusive_counter.h"
 
+#include "mongo/db/pipeline/document_source_match.h"
+#include "mongo/db/pipeline/document_source_single_document_transformation.h"
+#include "mongo/db/pipeline/semantic_analysis.h"
+
 namespace mongo {
 
 constexpr StringData DocumentSourceGroup::kStageName;
@@ -82,6 +86,88 @@ boost::intrusive_ptr<DocumentSource> DocumentSourceGroup::createFromBson(
     return createFromBsonWithMaxMemoryUsage(std::move(elem), expCtx, boost::none);
 }
 
+Pipeline::SourceContainer::iterator DocumentSourceGroup::doOptimizeAt(
+    Pipeline::SourceContainer::iterator itr, Pipeline::SourceContainer* container) {
+    invariant(*itr == this);
+
+    if (pushDotRenamedMatch(itr, container)) {
+        return itr;
+    }
+
+    return std::next(itr);
+}
+
+bool DocumentSourceGroup::pushDotRenamedMatch(Pipeline::SourceContainer::iterator itr,
+                                              Pipeline::SourceContainer* container) {
+    if (std::next(itr) == container->end() || std::next(std::next(itr)) == container->end()) {
+        return false;
+    }
+
+    // Keep separate iterators for each stage (projection, match).
+    auto prospectiveProjectionItr = std::next(itr);
+    auto prospectiveProjection =
+        dynamic_cast<DocumentSourceSingleDocumentTransformation*>(prospectiveProjectionItr->get());
+
+    auto prospectiveMatchItr = std::next(std::next(itr));
+    auto prospectiveMatch = dynamic_cast<DocumentSourceMatch*>(prospectiveMatchItr->get());
+
+    if (!prospectiveProjection || !prospectiveMatch) {
+        return false;
+    }
+
+    stdx::unordered_set<std::string> groupingFields;
+    StringMap<std::string> relevantRenames;
+
+    auto itsGroup = dynamic_cast<DocumentSourceGroup*>(itr->get());
+
+    auto idFields = itsGroup->getIdFields();
+    for (auto& idFieldsItr : idFields) {
+        groupingFields.insert(idFieldsItr.first);
+    }
+
+    GetModPathsReturn paths = prospectiveProjection->getModifiedPaths();
+
+    for (const auto& thisComplexRename : paths.complexRenames) {
+
+        // Check if the dotted renaming is done on a grouping field.
+        // This ensures that the top level is flat i.e., no arrays.
+        if (groupingFields.find(thisComplexRename.second) != groupingFields.end()) {
+            relevantRenames.insert(std::pair<std::string, std::string>(thisComplexRename.first,
+                                                                       thisComplexRename.second));
+        }
+    }
+
+    // Perform all changes on a copy of the match source.
+    boost::intrusive_ptr<DocumentSource> currentMatchCopyDocument =
+        prospectiveMatch->clone(prospectiveMatch->getContext());
+
+    auto currentMatchCopyDocumentMatch =
+        dynamic_cast<DocumentSourceMatch*>(currentMatchCopyDocument.get());
+
+    paths.renames = std::move(relevantRenames);
+
+    // Translate predicate statements based on the projection renames.
+    auto matchSplitForProject = currentMatchCopyDocumentMatch->splitMatchByModifiedFields(
+        currentMatchCopyDocumentMatch, paths);
+
+    if (matchSplitForProject.first) {
+        // Perform the swap of the projection and the match stages.
+        container->erase(prospectiveMatchItr);
+        container->insert(prospectiveProjectionItr, std::move(matchSplitForProject.first));
+
+        if (matchSplitForProject.second) {
+            // If there is a portion of the match stage predicate that is conflicting with the
+            // projection, re-insert it below the projection stage.
+            container->insert(std::next(prospectiveProjectionItr),
+                              std::move(matchSplitForProject.second));
+        }
+
+        return true;
+    }
+
+    return false;
+}
+
 boost::intrusive_ptr<DocumentSource> DocumentSourceGroup::createFromBsonWithMaxMemoryUsage(
     BSONElement elem,
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
diff --git a/src/mongo/db/pipeline/document_source_group.h b/src/mongo/db/pipeline/document_source_group.h
@@ -78,13 +78,39 @@ class DocumentSourceGroup final : public DocumentSourceGroupBase {
         const boost::intrusive_ptr<ExpressionContext>& expCtx,
         boost::optional<size_t> maxMemoryUsageBytes);
 
+    Pipeline::SourceContainer::iterator doOptimizeAt(Pipeline::SourceContainer::iterator itr,
+                                                     Pipeline::SourceContainer* container) override;
+
 protected:
     GetNextResult doGetNext() final;
 
     bool isSpecFieldReserved(StringData) final {
         return false;
     }
 
+    /**
+     * This optimization pushes a filter over a renamed grouping field
+     * before the group to improve performance.
+     *
+     * Specifically:
+     * $group { _id: {c: $x}, c: {aggregation}},
+     * $project { newVar: $_id.c }
+     * $match { newVar: "value"}
+     * ->
+     * $match { x: "value"}
+     * $group { _id: {c: $x}, c: {aggregation}},
+     * $project { newVar: $_id.c }
+     *
+     * Note: This optimization will not push over multiple grouping stages
+     * or multiple rename stages. Only the last set of group, project, match
+     * is taken into account. Furthermore, the optimization addresses specifically
+     * the defined sequence of operations to ensure the semantics of filters over arrays. Renaming
+     * dotted paths which include arrays change the evaluation of the filter statement and may lead
+     * to erroneous results.
+     */
+    bool pushDotRenamedMatch(Pipeline::SourceContainer::iterator itr,
+                             Pipeline::SourceContainer* container);
+
 private:
     explicit DocumentSourceGroup(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                                  boost::optional<size_t> maxMemoryUsageBytes = boost::none);
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
@@ -2624,6 +2624,15 @@ Expression::ComputedPaths ExpressionFieldPath::getComputedPaths(const std::strin
     if (_variable == renamingVar && _fieldPath.getPathLength() == 2u) {
         outputPaths.renames[exprFieldPath] = _fieldPath.tail().fullPath();
     } else {
+
+        // Add dotted renames also to complex renames, to be used prospectively in optimizations
+        // (e.g., pushDotRenamedMatch).
+        // We only include dotted rename paths of length 3, as current optimization are constrained
+        // to accepting only such paths to avoid semantic errors from array flattening.
+        if (_variable == renamingVar && _fieldPath.getPathLength() == 3u) {
+            outputPaths.complexRenames[exprFieldPath] = _fieldPath.tail().fullPath();
+        }
+
         outputPaths.paths.insert(exprFieldPath);
     }
 
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
@@ -212,8 +212,13 @@ class Expression : public RefCountable {
         OrderedPathSet paths;
 
         // Mappings from the old name of a path before applying this expression, to the new one
-        // after applying this expression.
+        // after applying this expression. This map includes solely simple field renaming
+        // expression.
         StringMap<std::string> renames;
+
+        // Mappings from the old name of a path before applying this expression. This map includes
+        // expressions which have dotted notation on the right side.
+        StringMap<std::string> complexRenames;
     };
 
     virtual ~Expression(){};
diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp

Original file line number	Diff line number	Diff line change
`@@ -259,7 +259,8 @@ void ProjectionNode::reportProjectedPaths(OrderedPathSet* projectedPaths) const`
`259`	`259`	`}`
`260`	`260`
`261`	`261`	`void ProjectionNode::reportComputedPaths(OrderedPathSet* computedPaths,`
`262`		`- StringMap<std::string>* renamedPaths) const {`
	`262`	`+ StringMap<std::string>* renamedPaths,`
	`263`	`+ StringMap<std::string>* complexRenamedPaths) const {`
`263`	`264`	`for (auto&& computedPair : _expressions) {`
`264`	`265`	`// The expression's path is the concatenation of the path to this node, plus the field name`
`265`	`266`	`// associated with the expression.`
`@@ -270,9 +271,15 @@ void ProjectionNode::reportComputedPaths(OrderedPathSet* computedPaths,`
`270`	`271`	`for (auto&& rename : exprComputedPaths.renames) {`
`271`	`272`	`(*renamedPaths)[rename.first] = rename.second;`
`272`	`273`	`}`
	`274`	`+`
	`275`	`+ if (complexRenamedPaths) {`
	`276`	`+ for (auto&& complexRename : exprComputedPaths.complexRenames) {`
	`277`	`+ (*complexRenamedPaths)[complexRename.first] = complexRename.second;`
	`278`	`+ }`
	`279`	`+ }`
`273`	`280`	`}`
`274`	`281`	`for (auto&& childPair : _children) {`
`275`		`- childPair.second->reportComputedPaths(computedPaths, renamedPaths);`
	`282`	`+ childPair.second->reportComputedPaths(computedPaths, renamedPaths, complexRenamedPaths);`
`276`	`283`	`}`
`277`	`284`	`}`
`278`	`285`