SERVER-65863 Have optimized _id expression(s) pushed down properly

(cherry picked from commit 7f55c457738556cf96236ef9c7aa4b30bbe07390)

2 files changed, 24 insertions, 9 deletions

diff --git a/src/mongo/db/pipeline/document_source_group.cpp b/src/mongo/db/pipeline/document_source_group.cpp
index 4b79d54f7c2..c561c1102b8 100644
--- a/src/mongo/db/pipeline/document_source_group.cpp
+++ b/src/mongo/db/pipeline/document_source_group.cpp
@@ -459,9 +459,6 @@ intrusive_ptr<Expression> parseIdExpression(const intrusive_ptr<ExpressionContex
 }  // namespace
 
 void DocumentSourceGroup::setIdExpression(const boost::intrusive_ptr<Expression> idExpression) {
-    // Remember the whole _id expression so that we can translate it to a SBE stage tree as a whole.
-    _idExpression = idExpression;
-
     if (auto object = dynamic_cast<ExpressionObject*>(idExpression.get())) {
         auto& childExpressions = object->getChildExpressions();
         invariant(!childExpressions.empty());  // We expect to have converted an empty object into a
@@ -479,6 +476,28 @@ void DocumentSourceGroup::setIdExpression(const boost::intrusive_ptr<Expression>
     }
 }
 
+boost::intrusive_ptr<Expression> DocumentSourceGroup::getIdExpression() const {
+    // _idFieldNames is empty and _idExpressions has one element when the _id expression is not an
+    // object expression.
+    if (_idFieldNames.empty() && _idExpressions.size() == 1) {
+        return _idExpressions[0];
+    }
+
+    tassert(6586300,
+            "Field and its expression must be always paired in ExpressionObject",
+            _idFieldNames.size() > 0 && _idFieldNames.size() == _idExpressions.size());
+
+    // Each expression in '_idExpressions' may have been optimized and so, compose the object _id
+    // expression out of the optimized expressions.
+    std::vector<std::pair<std::string, boost::intrusive_ptr<Expression>>> fieldsAndExprs;
+    for (size_t i = 0; i < _idExpressions.size(); ++i) {
+        fieldsAndExprs.emplace_back(_idFieldNames[i], _idExpressions[i]);
+    }
+
+    return ExpressionObject::create(_idExpressions[0]->getExpressionContext(),
+                                    std::move(fieldsAndExprs));
+}
+
 intrusive_ptr<DocumentSource> DocumentSourceGroup::createFromBson(
     BSONElement elem, const intrusive_ptr<ExpressionContext>& expCtx) {
     uassert(15947, "a group's fields must be specified in an object", elem.type() == Object);
diff --git a/src/mongo/db/pipeline/document_source_group.h b/src/mongo/db/pipeline/document_source_group.h
index 72fd2832bf1..edffed6758b 100644
--- a/src/mongo/db/pipeline/document_source_group.h
+++ b/src/mongo/db/pipeline/document_source_group.h
@@ -145,9 +145,7 @@ public:
     /**
      * Returns the expression to use to determine the group id of each document.
      */
-    const boost::intrusive_ptr<Expression> getIdExpression() const {
-        return _idExpression;
-    }
+    boost::intrusive_ptr<Expression> getIdExpression() const;
 
     /**
      * Returns true if this $group stage represents a 'global' $group which is merging together
@@ -284,13 +282,11 @@ private:
 
     std::shared_ptr<Sorter<Value, Value>::File> _file;
 
-    // The expression for the '_id' field.
-    boost::intrusive_ptr<Expression> _idExpression;
     // If the expression for the '_id' field represents a non-empty object, we track its fields'
     // names in '_idFieldNames'.
     std::vector<std::string> _idFieldNames;
     // Expressions for the individual fields when '_id' produces a document in the order of
-    // '_idFieldNames' or the whole expression if '_id' produces a scalar value.
+    // '_idFieldNames' or the whole expression otherwise.
     std::vector<boost::intrusive_ptr<Expression>> _idExpressions;
 
     bool _initialized;