SERVER-31530 Certain queries using 2d index could incorrecly compute

boundaries delivering incorrect results. The root cause was in incorrectly
propagating index tags during the query plan construction.

3 files changed, 38 insertions, 8 deletions

diff --git a/jstests/core/geo_2d_trailing_fields.js b/jstests/core/geo_2d_trailing_fields.js
index 28ed715e858..60c1e207ca7 100644
--- a/jstests/core/geo_2d_trailing_fields.js
+++ b/jstests/core/geo_2d_trailing_fields.js
@@ -37,4 +37,13 @@
     assert.eq(0, coll.find({a: {$geoWithin: {$center: [[0, 0], 1]}}, "b.c": [2, 3]}).itcount());
     assert.eq(
         0, coll.find({a: {$geoWithin: {$center: [[0, 0], 1]}}, "b.c": {$type: "array"}}).itcount());
+
+    coll.drop();
+    assert.commandWorked(coll.createIndex({a: "2d", "b.c": 1}));
+    assert.writeOK(coll.insert({a: [0, 0], b: [{c: 1}, {c: 2}]}));
+
+    // Verify that non-near 2d queries correctly handle predicates which cannot be covered due to
+    // array semantics.
+    assert.eq(
+        1, coll.find({a: {$geoWithin: {$center: [[0, 0], 1]}}, b: {$elemMatch: {c: 1}}}).itcount());
 }());
diff --git a/src/mongo/db/query/plan_enumerator.cpp b/src/mongo/db/query/plan_enumerator.cpp
index c35e3ae10a9..7016d57b043 100644
--- a/src/mongo/db/query/plan_enumerator.cpp
+++ b/src/mongo/db/query/plan_enumerator.cpp
@@ -223,27 +223,30 @@ bool canAssignPredToIndex(const RelevantTag* rt,
  */
 void tagForSort(MatchExpression* tree) {
     if (!Indexability::nodeCanUseIndexOnOwnField(tree)) {
-        size_t myTagValue = IndexTag::kNoIndex;
+        const IndexTag* myIndexTag = nullptr;
         for (size_t i = 0; i < tree->numChildren(); ++i) {
             MatchExpression* child = tree->getChild(i);
             tagForSort(child);
             if (child->getTag() &&
                 child->getTag()->getType() == MatchExpression::TagData::Type::IndexTag) {
-                IndexTag* childTag = static_cast<IndexTag*>(child->getTag());
-                myTagValue = std::min(myTagValue, childTag->index);
+                auto childTag = static_cast<const IndexTag*>(child->getTag());
+                if (!myIndexTag || myIndexTag->index > childTag->index) {
+                    myIndexTag = childTag;
+                }
             } else if (child->getTag() &&
                        child->getTag()->getType() ==
                            MatchExpression::TagData::Type::OrPushdownTag) {
                 OrPushdownTag* childTag = static_cast<OrPushdownTag*>(child->getTag());
                 if (childTag->getIndexTag()) {
-                    const IndexTag* indexTag =
-                        static_cast<const IndexTag*>(childTag->getIndexTag());
-                    myTagValue = std::min(myTagValue, indexTag->index);
+                    auto indexTag = static_cast<const IndexTag*>(childTag->getIndexTag());
+                    if (!myIndexTag || myIndexTag->index > indexTag->index) {
+                        myIndexTag = indexTag;
+                    }
                 }
             }
         }
-        if (myTagValue != IndexTag::kNoIndex) {
-            tree->setTag(new IndexTag(myTagValue));
+        if (myIndexTag) {
+            tree->setTag(new IndexTag(*myIndexTag));
         }
     }
 }
diff --git a/src/mongo/db/query/query_planner_geo_test.cpp b/src/mongo/db/query/query_planner_geo_test.cpp
index 15fa04bdf3d..c4b883ca580 100644
--- a/src/mongo/db/query/query_planner_geo_test.cpp
+++ b/src/mongo/db/query/query_planner_geo_test.cpp
@@ -1766,4 +1766,22 @@ TEST_F(QueryPlannerTest, 2dsphereNonNearWithInternalExprEqOverTrailingFieldMulti
         "{a: [], b: [['MinKey','MaxKey',true,true]]}}}}}");
 }
 
+TEST_F(QueryPlannerTest, 2dWithinPredicateOverTrailingFieldElemMatchMultikey) {
+    params.options = QueryPlannerParams::NO_TABLE_SCAN;
+
+    const bool multikey = true;
+    addIndex(BSON("a"
+                  << "2d"
+                  << "b"
+                  << 1),
+             multikey);
+
+    runQuery(fromjson("{a: {$geoWithin: {$center: [[0, 0], 1]}}, b: {$elemMatch: {c: 1}}}"));
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: {$and: [{b: {$elemMatch: {c: 1}}}, {a: {$geoWithin: {$center: [ [ 0, 0 "
+        "], 1 ]}}}]}, node: {ixscan: {filter: null, pattern: {a: '2d', b: 1}}}}}");
+}
+
+
 }  // namespace