SERVER-69260 Improve projection of FieldPathExpressions to slots and eliminate redundant project stages in SBE group translation

4 files changed, 143 insertions, 204 deletions

diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index 26716319fdd..d229d256000 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -2168,139 +2168,33 @@ struct FieldPathAndCondPreVisitor : public SelectiveConstExpressionVisitorBase {
     // To avoid overloaded-virtual warnings.
     using SelectiveConstExpressionVisitorBase::visit;
 
-    FieldPathAndCondPreVisitor(const F& fn, int32_t& nestedCondLevel)
-        : _fn(fn), _nestedCondLevel(nestedCondLevel) {}
+    explicit FieldPathAndCondPreVisitor(const F& fn) : _fn(fn) {}
 
     void visit(const ExpressionFieldPath* expr) final {
-        _fn(expr, _nestedCondLevel);
-    }
-
-    void visit(const ExpressionCond* expr) final {
-        ++_nestedCondLevel;
-    }
-
-    void visit(const ExpressionSwitch* expr) final {
-        ++_nestedCondLevel;
-    }
-
-    void visit(const ExpressionIfNull* expr) final {
-        ++_nestedCondLevel;
-    }
-
-    void visit(const ExpressionAnd* expr) final {
-        ++_nestedCondLevel;
-    }
-
-    void visit(const ExpressionOr* expr) final {
-        ++_nestedCondLevel;
+        _fn(expr);
     }
 
     F _fn;
-    // Tracks the number of conditional expressions like $cond or $ifNull that are above us in the
-    // tree.
-    int32_t& _nestedCondLevel;
-};
-
-struct CondPostVisitor : public SelectiveConstExpressionVisitorBase {
-    // To avoid overloaded-virtual warnings.
-    using SelectiveConstExpressionVisitorBase::visit;
-
-    CondPostVisitor(int32_t& nestedCondLevel) : _nestedCondLevel(nestedCondLevel) {}
-
-    void visit(const ExpressionCond* expr) final {
-        --_nestedCondLevel;
-    }
-
-    void visit(const ExpressionSwitch* expr) final {
-        --_nestedCondLevel;
-    }
-
-    void visit(const ExpressionIfNull* expr) final {
-        --_nestedCondLevel;
-    }
-
-    void visit(const ExpressionAnd* expr) final {
-        --_nestedCondLevel;
-    }
-
-    void visit(const ExpressionOr* expr) final {
-        --_nestedCondLevel;
-    }
-
-    int32_t& _nestedCondLevel;
 };
 
 /**
  * Walks through the 'expr' expression tree and whenever finds an 'ExpressionFieldPath', calls
  * the 'fn' function. Type requirement for 'fn' is it must have a const 'ExpressionFieldPath'
- * pointer parameter and 'nestedCondLevel' parameter.
+ * pointer parameter.
  */
 template <typename F>
 void walkAndActOnFieldPaths(Expression* expr, const F& fn) {
-    int32_t nestedCondLevel = 0;
-    FieldPathAndCondPreVisitor<F> preVisitor(fn, nestedCondLevel);
-    CondPostVisitor postVisitor(nestedCondLevel);
-    ExpressionWalker walker(&preVisitor, nullptr /*inVisitor*/, &postVisitor);
+    FieldPathAndCondPreVisitor<F> preVisitor(fn);
+    ExpressionWalker walker(&preVisitor, nullptr /*inVisitor*/, nullptr /*postVisitor*/);
     expression_walker::walk(expr, &walker);
 }
 
-/**
- * If there are adjacent $group stages in a pipeline and two $group stages are pushed down together,
- * the first $group becomes a child GROUP node and the second $group becomes a parent GROUP node in
- * a query solution tree. In the case that all field paths are top-level fields for the parent GROUP
- * node, we can skip the mkbson stage of the child GROUP node and instead, the child GROUP node
- * returns top-level fields and their associated slots. If a field path is found in child outputs,
- * we should replace getField() by the associated slot because there's no object on which we can
- * call getField().
- *
- * Also deduplicates field lookups for a field path which is accessed multiple times.
- */
-EvalStage optimizeFieldPaths(StageBuilderState& state,
-                             const boost::intrusive_ptr<Expression>& expr,
-                             EvalStage stage,
-                             const PlanStageSlots& outputs,
-                             PlanNodeId nodeId) {
-    using namespace fmt::literals;
-    auto rootSlot = outputs.getIfExists(PlanStageSlots::kResult);
-
-    walkAndActOnFieldPaths(expr.get(), [&](const ExpressionFieldPath* fieldExpr, int32_t) {
-        // We optimize neither a field path for the top-level document itself nor a field path that
-        // refers to a variable instead of calling getField().
-        if (fieldExpr->getFieldPath().getPathLength() == 1 || fieldExpr->isVariableReference()) {
-            return;
-        }
-
-        auto fieldPathStr = fieldExpr->getFieldPath().fullPath();
-
-        if (!state.preGeneratedExprs.contains(fieldPathStr)) {
-            auto rootExpr = rootSlot.has_value() ? EvalExpr{*rootSlot} : EvalExpr{};
-            auto expr = generateExpression(state, fieldExpr, std::move(rootExpr), &outputs);
-
-            auto [slot, projectStage] = projectEvalExpr(
-                std::move(expr), std::move(stage), nodeId, state.slotIdGenerator, state.slotVarMap);
-
-            state.preGeneratedExprs.emplace(fieldPathStr, slot);
-            stage = std::move(projectStage);
-        }
-    });
-
-    return stage;
-}
-
-EvalExprStagePair generateGroupByKeyImpl(StageBuilderState& state,
-                                         const boost::intrusive_ptr<Expression>& idExpr,
-                                         const PlanStageSlots& outputs,
-                                         const boost::optional<sbe::value::SlotId>& rootSlot,
-                                         EvalStage stage,
-                                         PlanNodeId nodeId,
-                                         sbe::value::SlotIdGenerator* slotIdGenerator) {
-    // Optimize field paths before generating the expression.
-    stage = optimizeFieldPaths(state, idExpr, std::move(stage), outputs, nodeId);
-
+EvalExpr generateGroupByKeyImpl(StageBuilderState& state,
+                                const boost::intrusive_ptr<Expression>& idExpr,
+                                const PlanStageSlots& outputs,
+                                const boost::optional<sbe::value::SlotId>& rootSlot) {
     auto rootExpr = rootSlot.has_value() ? EvalExpr{*rootSlot} : EvalExpr{};
-    auto expr = generateExpression(state, idExpr.get(), std::move(rootExpr), &outputs);
-
-    return {std::move(expr), std::move(stage)};
+    return generateExpression(state, idExpr.get(), std::move(rootExpr), &outputs);
 }
 
 std::tuple<sbe::value::SlotVector, EvalStage, std::unique_ptr<sbe::EExpression>> generateGroupByKey(
@@ -2317,11 +2211,10 @@ std::tuple<sbe::value::SlotVector, EvalStage, std::unique_ptr<sbe::EExpression>>
         sbe::EExpression::Vector exprs;
 
         for (auto&& [fieldName, fieldExpr] : idExprObj->getChildExpressions()) {
-            auto [groupByEvalExpr, groupByEvalStage] = generateGroupByKeyImpl(
-                state, fieldExpr, outputs, rootSlot, std::move(stage), nodeId, slotIdGenerator);
+            auto groupByEvalExpr = generateGroupByKeyImpl(state, fieldExpr, outputs, rootSlot);
 
             auto [slot, projectStage] = projectEvalExpr(std::move(groupByEvalExpr),
-                                                        std::move(groupByEvalStage),
+                                                        std::move(stage),
                                                         nodeId,
                                                         slotIdGenerator,
                                                         state.slotVarMap);
@@ -2357,34 +2250,26 @@ std::tuple<sbe::value::SlotVector, EvalStage, std::unique_ptr<sbe::EExpression>>
         return {slots, std::move(stage), sbe::makeE<sbe::EFunction>("newObj"_sd, std::move(exprs))};
     }
 
-    auto [groupByEvalExpr, groupByEvalStage] = generateGroupByKeyImpl(
-        state, idExpr, outputs, rootSlot, std::move(stage), nodeId, slotIdGenerator);
+    auto groupByEvalExpr = generateGroupByKeyImpl(state, idExpr, outputs, rootSlot);
 
     // The group-by field may end up being 'Nothing' and in that case _id: null will be
     // returned. Calling 'makeFillEmptyNull' for the group-by field takes care of that.
     auto fillEmptyNullExpr = makeFillEmptyNull(groupByEvalExpr.extractExpr(state.slotVarMap));
-    auto [slot, projectStage] = projectEvalExpr(std::move(fillEmptyNullExpr),
-                                                std::move(groupByEvalStage),
-                                                nodeId,
-                                                slotIdGenerator,
-                                                state.slotVarMap);
+    auto [slot, projectStage] = projectEvalExpr(
+        std::move(fillEmptyNullExpr), std::move(stage), nodeId, slotIdGenerator, state.slotVarMap);
     stage = std::move(projectStage);
 
     return {sbe::value::SlotVector{slot}, std::move(stage), nullptr};
 }
 
-std::tuple<sbe::value::SlotVector, EvalStage> generateAccumulator(
+sbe::value::SlotVector generateAccumulator(
     StageBuilderState& state,
     const AccumulationStatement& accStmt,
-    EvalStage stage,
     const PlanStageSlots& outputs,
-    PlanNodeId nodeId,
     sbe::value::SlotIdGenerator* slotIdGenerator,
     sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>>& accSlotToExprMap) {
     auto rootSlot = outputs.getIfExists(PlanStageSlots::kResult);
 
-    // Input fields may need field traversal.
-    stage = optimizeFieldPaths(state, accStmt.expr.argument, std::move(stage), outputs, nodeId);
     auto rootExpr = rootSlot.has_value() ? EvalExpr{*rootSlot} : EvalExpr{};
     auto argExpr =
         generateExpression(state, accStmt.expr.argument.get(), std::move(rootExpr), &outputs);
@@ -2403,7 +2288,7 @@ std::tuple<sbe::value::SlotVector, EvalStage> generateAccumulator(
         accSlotToExprMap.emplace(slot, std::move(accExpr));
     }
 
-    return {std::move(aggSlots), std::move(stage)};
+    return aggSlots;
 }
 
 std::tuple<std::vector<std::string>, sbe::value::SlotVector, EvalStage> generateGroupFinalStage(
@@ -2491,26 +2376,26 @@ sbe::value::SlotVector dedupGroupBySlots(const sbe::value::SlotVector& groupBySl
  * must return 'BSONObject'. The returned 'BSONObject' will always have an "_id" field for the group
  * key and zero or more field(s) for accumulators.
  *
- * For example, a QSN tree: GroupNode(nodeId=2) over a VirtualScanNode(nodeId=1), we would have the
- * following translated sbe::PlanStage tree. In this example, we assume that the $group pipeline
+ * For example, a QSN tree: GroupNode(nodeId=2) over a CollectionScanNode(nodeId=1), we would have
+ * the following translated sbe::PlanStage tree. In this example, we assume that the $group pipeline
  * spec is {"_id": "$a", "x": {"$min": "$b"}, "y": {"$first": "$b"}}.
  *
- * [2] mkbson s14 [_id = s9, x = s14, y = s13] true false
- * [2] project [s14 = fillEmpty (s11, null)]
- * [2] group [s9] [s12 = min (if (! exists (s9) || typeMatch (s9, 0x00000440), Nothing, s9)),
- *                 s13 = first (fillEmpty (s10, null))]
- * [2] project [s11 = getField (s7, "b")]
- * [2] project [s10 = getField (s7, "b")]
- * [2] project [s9 = fillEmpty (s8, null)]
- * [2] project [s8 = getField (s7, "a")]
- * [1] project [s7 = getElement (s5, 0)]
- * [1] unwind s5 s6 s4 false
- * [1] project [s4 = [[{"a" : 1, "b" : 1}], [{"a" : 1, "b" : 2}], [{"a" : 2, "b" : 3}]]]
- * [1] limit 1 [1] coscan
+ * [2] mkbson s12 [_id = s8, x = s11, y = s10] true false
+ * [2] project [s11 = (s9 ?: null)]
+ * [2] group [s8] [s9 = min(
+ *   let [
+ *      l1.0 = s5
+ *  ]
+ *  in
+ *      if (typeMatch(l1.0, 1088ll) ?: true)
+ *      then Nothing
+ *      else l1.0
+ * ), s10 = first((s5 ?: null))]
+ * [2] project [s8 = (s4 ?: null)]
+ * [1] scan s6 s7 none none none none [s4 = a, s5 = b] @<collUuid> true false
  */
 std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder::buildGroup(
     const QuerySolutionNode* root, const PlanStageReqs& reqs) {
-    using namespace fmt::literals;
     tassert(6023414, "buildGroup() does not support kSortKey", !reqs.hasSortKeys());
 
     auto groupNode = static_cast<const GroupNode*>(root);
@@ -2530,24 +2415,32 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
 
     auto childReqs = reqs.copy().set(kResult).clearAllFields();
 
-    // If the group node doesn't need the whole document, then we take all the top-level fields
-    // referenced by the group node and we add them to 'childReqs'.
-    if (!groupNode->needWholeDocument) {
-        childReqs.setFields(getTopLevelFields(groupNode->requiredFields));
-    }
-
-    // If the child is a GROUP and we can get everything we need from top-level field slots, then
-    // we can avoid unnecessary materialization and not request the kResult slot from the child.
-    if (childNode->getType() == StageType::STAGE_GROUP && !groupNode->needWholeDocument &&
-        !containsPoisonTopLevelField(groupNode->requiredFields)) {
-        childReqs.clear(kResult);
+    // If the group node references any top level fields, we take all of them and add them to
+    // 'childReqs'. Note that this happens regardless of whether we need the whole document because
+    // it can be the case that this stage references '$$ROOT' as well as some top level fields.
+    auto topLevelFields = getTopLevelFields(groupNode->requiredFields);
+    if (!topLevelFields.empty()) {
+        childReqs.setFields(topLevelFields);
     }
 
     if (!groupNode->needWholeDocument) {
+        // Tracks whether we need to request kResult. One such case is lookup of the '$$POISON'
+        // field.
+        bool rootDocIsNeeded = containsPoisonTopLevelField(groupNode->requiredFields);
+        auto referencesRoot = [&](const ExpressionFieldPath* fieldExpr) {
+            rootDocIsNeeded = rootDocIsNeeded || fieldExpr->isROOT();
+        };
+
+        // Walk over all field paths involved in this $group stage.
+        walkAndActOnFieldPaths(idExpr.get(), referencesRoot);
+        for (const auto& accStmt : accStmts) {
+            walkAndActOnFieldPaths(accStmt.expr.argument.get(), referencesRoot);
+        }
+
         // If the group node doesn't have any dependency (e.g. $count) or if the dependency can be
         // satisfied by the child node (e.g. covered index scan), we can clear the kResult
         // requirement for the child.
-        if (groupNode->requiredFields.empty()) {
+        if (groupNode->requiredFields.empty() || !rootDocIsNeeded) {
             childReqs.clear(kResult);
         } else if (childNode->getType() == StageType::STAGE_PROJECTION_COVERED) {
             auto pn = static_cast<const ProjectionNodeCovered*>(childNode);
@@ -2565,10 +2458,66 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
 
     // Builds the child and gets the child result slot.
     auto [childStage, childOutputs] = build(childNode, childReqs);
+    auto maybeRootSlot = childOutputs.getIfExists(kResult);
+    auto rootExpr = maybeRootSlot.has_value() ? EvalExpr{*maybeRootSlot} : EvalExpr{};
+    auto* childOutputsPtr = &childOutputs;
+
+    // Set of field paths referenced by group. Useful for de-duplicating fields and clearing the
+    // slots corresponding to fields in 'childOutputs' so that they are not mistakenly referenced by
+    // parent stages.
+    StringSet groupFieldSet;
+
+    // Slot to EExpression map that tracks path traversal expressions. Note that this only contains
+    // expressions corresponding to paths which require traversals (that is, if there exists a
+    // top level field slot corresponding to a field, we take care not to add it to 'fpMap' to
+    // avoid rebinding a slot).
+    sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> fpMap;
+
+    // Lambda which populates 'fpMap' and 'childOutputs' with an expression and/or a slot,
+    // respectively, corresponding to the value of 'fieldExpr'.
+    auto accumulateFieldPaths = [&](const ExpressionFieldPath* fieldExpr) {
+        // We optimize neither a field path for the top-level document itself nor a field path
+        // that refers to a variable instead.
+        if (fieldExpr->getFieldPath().getPathLength() == 1 || fieldExpr->isVariableReference()) {
+            return;
+        }
+
+        // Don't generate an expression if we have one already.
+        const std::string fp = fieldExpr->getFieldPathWithoutCurrentPrefix().fullPath();
+        if (groupFieldSet.count(fp)) {
+            return;
+        }
+
+        // Mark 'fp' as being seen and either find a slot corresponding to it or generate an
+        // expression for it and bind it to a slot.
+        groupFieldSet.insert(fp);
+        sbe::value::SlotId slot = [&]() -> sbe::value::SlotId {
+            // Special case: top level fields which already have a slot.
+            if (fieldExpr->getFieldPath().getPathLength() == 2) {
+                return childOutputsPtr->get({PlanStageSlots::kField, StringData(fp)});
+            } else {
+                // General case: we need to generate a path traversal expression.
+                auto result = stage_builder::generateExpression(
+                    _state, fieldExpr, rootExpr.clone(), childOutputsPtr);
+
+                if (result.hasSlot()) {
+                    return *result.getSlot();
+                } else {
+                    auto newSlot = _slotIdGenerator.generate();
+                    fpMap.emplace(newSlot, result.extractExpr(_state.slotVarMap));
+                    return newSlot;
+                }
+            }
+        }();
 
-    tassert(6075900,
-            "Expected no optimized expressions but got: {}"_format(_state.preGeneratedExprs.size()),
-            _state.preGeneratedExprs.empty());
+        childOutputsPtr->set(std::make_pair(PlanStageSlots::kPathExpr, std::move(fp)), slot);
+    };
+
+    // Walk over all field paths involved in this $group stage.
+    walkAndActOnFieldPaths(idExpr.get(), accumulateFieldPaths);
+    for (const auto& accStmt : accStmts) {
+        walkAndActOnFieldPaths(accStmt.expr.argument.get(), accumulateFieldPaths);
+    }
 
     // Translates the group-by expression and wraps it with 'fillEmpty(..., null)' because the
     // missing field value for _id should be mapped to 'Null'.
@@ -2576,24 +2525,21 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     auto childEvalStage =
         EvalStage{std::move(childStage), getSlotsToForward(forwardingReqs, childOutputs)};
 
+    if (!fpMap.empty()) {
+        childEvalStage = makeProject(std::move(childEvalStage), std::move(fpMap), nodeId);
+    }
+
     auto [groupBySlots, groupByEvalStage, idDocExpr] = generateGroupByKey(
         _state, idExpr, childOutputs, std::move(childEvalStage), nodeId, &_slotIdGenerator);
 
     // Translates accumulators which are executed inside the group stage and gets slots for
     // accumulators.
-    stage_builder::EvalStage accProjEvalStage = std::move(groupByEvalStage);
+    stage_builder::EvalStage currentStage = std::move(groupByEvalStage);
     sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> accSlotToExprMap;
     std::vector<sbe::value::SlotVector> aggSlotsVec;
     for (const auto& accStmt : accStmts) {
-        auto [aggSlots, tempEvalStage] = generateAccumulator(_state,
-                                                             accStmt,
-                                                             std::move(accProjEvalStage),
-                                                             childOutputs,
-                                                             nodeId,
-                                                             &_slotIdGenerator,
-                                                             accSlotToExprMap);
-        aggSlotsVec.emplace_back(std::move(aggSlots));
-        accProjEvalStage = std::move(tempEvalStage);
+        aggSlotsVec.emplace_back(generateAccumulator(
+            _state, accStmt, childOutputs, &_slotIdGenerator, accSlotToExprMap));
     }
 
     // There might be duplicated expressions and slots. Dedup them before creating a HashAgg
@@ -2601,7 +2547,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     // reasonable because duplicated expressions would not contribute to grouping.
     auto dedupedGroupBySlots = dedupGroupBySlots(groupBySlots);
     // Builds a group stage with accumulator expressions and group-by slot(s).
-    auto groupEvalStage = makeHashAgg(std::move(accProjEvalStage),
+    auto groupEvalStage = makeHashAgg(std::move(currentStage),
                                       dedupedGroupBySlots,
                                       std::move(accSlotToExprMap),
                                       _state.data->env->getSlotIfExists("collator"_sd),
@@ -2640,8 +2586,11 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
             "slots",
             finalSlots.size() == 1 + accStmts.size());
 
-    // Cleans up optimized expressions.
-    _state.preGeneratedExprs.clear();
+    // Clear all fields needed by this group stage from 'childOutputs' to avoid references to
+    // ExpressionFieldPath values that are no longer visible.
+    for (const auto& groupField : groupFieldSet) {
+        childOutputs.clear({PlanStageSlots::kPathExpr, StringData(groupField)});
+    }
 
     auto fieldNamesSet = StringDataSet{fieldNames.begin(), fieldNames.end()};
     auto [fields, additionalFields] =
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index c743342102b..fd7d95c14f5 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -123,10 +123,15 @@ public:
     //    given field path. This raw key value can be used for sorting / comparison, but it is not
     //    always equal to the actual value of the field path (for example, if the key is coming from
     //    an index that has a non-simple collation).
+    // 4) kPathExpr slots represent the value obtained from evaluating an 'ExpressionFieldPath'.
+    //    Typically, this is requested by stages that wish to avoid generating duplicate
+    //    expressions for path traversal (for example, $group stages which reference the same
+    //    field path across multiple accumulators).
     enum class Type {
         kMeta,
         kField,
         kSortKey,
+        kPathExpr,
     };
 
     using Name = std::pair<Type, StringData>;
@@ -135,6 +140,7 @@ public:
     static constexpr auto kMeta = Type::kMeta;
     static constexpr auto kField = Type::kField;
     static constexpr auto kSortKey = Type::kSortKey;
+    static constexpr auto kPathExpr = Type::kPathExpr;
 
     static constexpr Name kResult = {kMeta, "result"_sd};
     static constexpr Name kRecordId = {kMeta, "recordId"_sd};
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index fa6b538af16..837282c485e 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -2075,27 +2075,6 @@ public:
             sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(expExpr)));
     }
     void visit(const ExpressionFieldPath* expr) final {
-        // There's a chance that we've already generated a SBE plan stage tree for this field path,
-        // in which case we avoid regeneration of the same plan stage tree.
-        if (!_context->state.preGeneratedExprs.empty()) {
-            if (auto it = _context->state.preGeneratedExprs.find(expr->getFieldPath().fullPath());
-                it != _context->state.preGeneratedExprs.end()) {
-                tassert(6089301,
-                        "Expressions for top-level documents / variables must not be pre-generated",
-                        expr->getFieldPath().getPathLength() != 1 && !expr->isVariableReference());
-                if (auto optionalSlot = it->second.getSlot(); optionalSlot) {
-                    _context->pushExpr(*optionalSlot);
-                } else if (it->second.hasABT()) {
-                    _context->pushExpr(it->second.extractABT(_context->state.slotVarMap));
-                } else {
-                    auto preGeneratedExpr = it->second.extractExpr(_context->state.slotVarMap);
-                    _context->pushExpr(preGeneratedExpr->clone());
-                    it->second = std::move(preGeneratedExpr);
-                }
-                return;
-            }
-        }
-
         EvalExpr inputExpr;
         boost::optional<sbe::value::SlotId> topLevelFieldSlot;
         bool expectsDocumentInputOnly = false;
@@ -2104,14 +2083,24 @@ public:
             : boost::none;
 
         if (!Variables::isUserDefinedVariable(expr->getVariableId())) {
+            const auto* slots = _context->slots;
             if (expr->getVariableId() == Variables::kRootId) {
                 // Set inputExpr to refer to the root document.
                 inputExpr = _context->rootExpr.clone();
                 expectsDocumentInputOnly = true;
-                // Check if a slot is available for the top-level field referred to by 'expr'.
-                if (expr->getFieldPath().getPathLength() > 1 && _context->slots) {
+
+                if (slots && fp) {
+                    // Check if we already have a slot containing an expression corresponding
+                    // to 'expr'.
+                    auto fpe = std::make_pair(PlanStageSlots::kPathExpr, fp->fullPath());
+                    if (slots->has(fpe)) {
+                        _context->pushExpr(slots->get(fpe));
+                        return;
+                    }
+
+                    // Obtain a slot for the top-level field referred to by 'expr', if one exists.
                     auto topLevelField = std::make_pair(PlanStageSlots::kField, fp->front());
-                    topLevelFieldSlot = _context->slots->getIfExists(topLevelField);
+                    topLevelFieldSlot = slots->getIfExists(topLevelField);
                 }
             } else if (expr->getVariableId() == Variables::kRemoveId) {
                 // For the field paths that begin with "$$REMOVE", we always produce Nothing,
diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.h b/src/mongo/db/query/sbe_stage_builder_helpers.h
index 3b5cf77c310..74b27315ae0 100644
--- a/src/mongo/db/query/sbe_stage_builder_helpers.h
+++ b/src/mongo/db/query/sbe_stage_builder_helpers.h
@@ -665,11 +665,6 @@ struct StageBuilderState {
     // A flag to indicate the user allows disk use for spilling.
     bool allowDiskUse;
 
-    // This map is used to plumb through pre-generated field expressions ('EvalExpr')
-    // corresponding to field paths to 'generateExpression' to avoid repeated expression generation.
-    // Key is expected to represent field paths in form CURRENT.<field_name>[.<field_name>]*.
-    stdx::unordered_map<std::string /*field path*/, EvalExpr> preGeneratedExprs;
-
     // Holds the mapping between the custom ABT variable names and the slot id they are referencing.
     optimizer::SlotVarMap slotVarMap;
 };