SERVER-54325 Fix SBE behavior with projections on dotted paths

3 files changed, 38 insertions, 14 deletions

diff --git a/jstests/aggregation/sources/merge/requires_unique_index.js b/jstests/aggregation/sources/merge/requires_unique_index.js
index dbc5788428c..1565381c6c3 100644
--- a/jstests/aggregation/sources/merge/requires_unique_index.js
+++ b/jstests/aggregation/sources/merge/requires_unique_index.js
@@ -5,9 +5,6 @@
 // Note that this test does *not* use the drop shell helper but instead runs the drop command
 // manually. This is to avoid implicit creation and sharding of the $merge target collections in the
 // passthrough suites.
-// @tags: [
-//   sbe_incompatible,
-// ]
 (function() {
 "use strict";
 
diff --git a/jstests/aggregation/sources/project/project_dotted_paths.js b/jstests/aggregation/sources/project/project_dotted_paths.js
index f5ca3f7461b..cf730861e5e 100644
--- a/jstests/aggregation/sources/project/project_dotted_paths.js
+++ b/jstests/aggregation/sources/project/project_dotted_paths.js
@@ -1,8 +1,5 @@
 // Test that projection of dotted paths which happens in the "agg" layer works correctly. See
 // SERVER-26066 for details.
-// @tags: [
-//   sbe_incompatible,
-// ]
 (function() {
 const coll = db.project_dotted_paths;
 coll.drop();
diff --git a/src/mongo/db/query/sbe_stage_builder_projection.cpp b/src/mongo/db/query/sbe_stage_builder_projection.cpp
index 4654a67e0b0..896b36fba27 100644
--- a/src/mongo/db/query/sbe_stage_builder_projection.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_projection.cpp
@@ -135,6 +135,9 @@ struct ProjectionTraversalVisitorContext {
         // Vector containing expressions for each of the projections at the current level. There is
         // an eval for each of the fields in the current nested level.
         std::vector<ProjectEval> evals;
+
+        // Whether or not any subtree of this level has a computed field.
+        bool subtreeContainsComputedField = false;
     };
 
     const auto& topFrontField() const {
@@ -259,7 +262,9 @@ public:
 
     void visit(const projection_ast::ProjectionElemMatchASTNode* node) final {}
 
-    void visit(const projection_ast::ExpressionASTNode* node) final {}
+    void visit(const projection_ast::ExpressionASTNode* node) final {
+        _context->topLevel().subtreeContainsComputedField = true;
+    }
 
     void visit(const projection_ast::MatchExpressionASTNode* node) final {}
 
@@ -424,12 +429,12 @@ public:
             prepareFieldEvals(_context, node);
 
         // Finally, inject an mkobj stage to generate a document for the current nested level. For
-        // inclusion projection also add constant filter stage on top to filter out input values for
-        // nested traversal if they're not documents.
+        // inclusion projection also add a filter stage on top to filter out input values for
+        // nested traversal if they don't result in documents.
         auto childLevelInputSlot = _context->topLevel().inputSlot;
         auto childLevelResultSlot = _context->slotIdGenerator->generate();
         if (_context->projectType == projection_ast::ProjectType::kInclusion) {
-            childLevelStage = sbe::makeS<sbe::FilterStage<true>>(
+            auto mkBsonStage =
                 sbe::makeS<sbe::MakeBsonObjStage>(std::move(childLevelStage),
                                                   childLevelResultSlot,
                                                   childLevelInputSlot,
@@ -439,9 +444,29 @@ public:
                                                   std::move(projectSlots),
                                                   true,
                                                   false,
-                                                  _context->planNodeId),
-                makeFunction("isObject"sv, sbe::makeE<sbe::EVariable>(childLevelInputSlot)),
-                _context->planNodeId);
+                                                  _context->planNodeId);
+
+            if (_context->topLevel().subtreeContainsComputedField) {
+                // Projections of computed fields should always be applied to elements of an array,
+                // even if the elements aren't objects. For example:
+                // projection: {a: {b: "x"}}
+                // document: {a: [1,2,3]}
+                // result: {a: [{b: "x"}, {b: "x"}, {b: "x"}, {b: "x"}]}
+
+                childLevelStage = std::move(mkBsonStage);
+            } else {
+                // There are no computed fields, only inclusions. So anything that's not a document
+                // will get projected out. Example:
+                // projection: {a: {b: 1}}
+                // document: {a: [1, {b: 2}, 3]}
+                // result: {a: [{b: 2}]}
+
+                childLevelStage = sbe::makeS<sbe::FilterStage<true>>(
+                    std::move(mkBsonStage),
+                    makeFunction("isObject"sv, sbe::makeE<sbe::EVariable>(childLevelInputSlot)),
+                    _context->planNodeId);
+            }
+
         } else {
             childLevelStage =
                 sbe::makeS<sbe::MakeBsonObjStage>(std::move(childLevelStage),
@@ -458,7 +483,12 @@ public:
 
         // We are done with the child level. Now we need to extract corresponding field from parent
         // level, traverse it and assign value to 'childLevelInputSlot'.
-        _context->popLevel();
+        {
+            const bool containsComputedField = _context->topLevel().subtreeContainsComputedField;
+            _context->popLevel();
+            _context->topLevel().subtreeContainsComputedField =
+                _context->topLevel().subtreeContainsComputedField || containsComputedField;
+        }
 
         auto parentLevelInputSlot = _context->topLevel().inputSlot;
         auto parentLevelStage{std::move(_context->topLevel().evalStage)};