diff options
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)}; |