Add multikey info for non-WC fields in planner

2 files changed, 49 insertions, 19 deletions

diff --git a/src/mongo/db/query/planner_wildcard_helpers.cpp b/src/mongo/db/query/planner_wildcard_helpers.cpp
index b728b76c363..fe4c2fed9da 100644
--- a/src/mongo/db/query/planner_wildcard_helpers.cpp
+++ b/src/mongo/db/query/planner_wildcard_helpers.cpp
@@ -175,35 +175,30 @@ FieldRef pathWithoutSpecifiedComponents(const FieldRef& path,
 }
 
 /**
- * Returns a MultikeyPaths which indicates which components of 'indexedPath' are multikey, by
- * looking up multikeyness in 'multikeyPathSet'.
+ * Returns a MultikeyPaths which indicates which components of of the index key pattern are
+ * multikey, by looking up multikeyness in 'multikeyPathSet'. 'indexedPath' is used as the new key
+ * for the wildcard element.
  */
 MultikeyPaths buildMultiKeyPathsForExpandedWildcardIndexEntry(
     const FieldRef& indexedPath,
     const std::set<FieldRef>& multikeyPathSet,
     const BSONObj& keyPattern) {
-    FieldRef pathToLookup;
     MultikeyPaths multikeyPaths = {};
 
-    // TODO: what if the non-wildcard entries do represent arrays? Then they should be multikey.
-    // For each non-wildcard index entry, indicate that they are not multikey.
-    // Note this assumes only one wildcard entry.
+    // For each key in the index key pattern, determine which components are multikey.
     for (auto& elem : keyPattern) {
-        if (!elem.fieldNameStringData().endsWith("$**")) {
-            MultikeyComponents comps;
-            multikeyPaths.insert(multikeyPaths.end(), comps);
-        }
-    }
-
-    // Check each part of the wildcard entry 'indexedPath' to determine if it's multikey.
-    MultikeyComponents comps;
-    for (size_t i = 0; i < indexedPath.numParts(); ++i) {
-        pathToLookup.appendPart(indexedPath.getPart(i));
-        if (fieldNameOrArrayIndexPathSetContains(multikeyPathSet, comps, pathToLookup)) {
-            comps.insert(i);
+        const FieldRef& currField =
+            elem.fieldNameStringData().endsWith("$**") ? indexedPath : FieldRef(elem.fieldName());
+        FieldRef pathToLookup;
+        MultikeyComponents comps;
+        for (size_t i = 0; i < currField.numParts(); ++i) {
+            pathToLookup.appendPart(currField.getPart(i));
+            if (fieldNameOrArrayIndexPathSetContains(multikeyPathSet, comps, pathToLookup)) {
+                comps.insert(i);
+            }
         }
+        multikeyPaths.insert(multikeyPaths.end(), comps);
     }
-    multikeyPaths.insert(multikeyPaths.end(), comps);
 
     return multikeyPaths;
 }
diff --git a/src/mongo/db/query/planner_wildcard_helpers_test.cpp b/src/mongo/db/query/planner_wildcard_helpers_test.cpp
index 8bcaf0e19a2..96fa36b7bbc 100644
--- a/src/mongo/db/query/planner_wildcard_helpers_test.cpp
+++ b/src/mongo/db/query/planner_wildcard_helpers_test.cpp
@@ -153,6 +153,41 @@ TEST_F(PlannerWildcardHelpersTest, ExpandCompoundWildcardIndexEntryNoMatch) {
     indexEntryKeyPatternsMatch(&expectedKeyPatterns, &out);
 }
 
+TEST_F(PlannerWildcardHelpersTest, ExpandEnsureMultikeySetForAllCompoundFields) {
+    std::vector<IndexEntry> out;
+    stdx::unordered_set<std::string> fields{"a", "b"};
+    const auto indexEntry = makeIndexEntry(BSON("a" << 1 << "$**" << 1),
+                                           {},
+                                           {FieldRef("a"), FieldRef("b"), FieldRef("c")},
+                                           {fromjson("{wildcardProjection: {a: 0}}")});
+    wildcard_planning::expandWildcardIndexEntry(indexEntry.first, fields, &out);
+
+    ASSERT_EQ(out.size(), 1u);
+    ASSERT_TRUE(out[0].multikey);
+    ASSERT_EQ(out[0].multikeyPaths.size(), 2);
+    ASSERT(out[0].multikeyPaths[0] == MultikeyComponents{0u});  // a is a multikey path
+    ASSERT(out[0].multikeyPaths[1] == MultikeyComponents{0u});  // and so is b
+    ASSERT_BSONOBJ_EQ(out[0].keyPattern, {fromjson("{a: 1, b: 1}")});
+}
+
+TEST_F(PlannerWildcardHelpersTest, ExpandEnsureMultikeySetForAllCompoundFieldsDotted) {
+    std::vector<IndexEntry> out;
+    stdx::unordered_set<std::string> fields{"a.b", "c.d.e"};
+    const auto indexEntry = makeIndexEntry(
+        BSON("a.b" << 1 << "$**" << 1),
+        {},
+        {FieldRef("a"), FieldRef("a.b"), FieldRef("b"), FieldRef("c"), FieldRef("c.d.e")},
+        {fromjson("{wildcardProjection: {a: 0}}")});
+    wildcard_planning::expandWildcardIndexEntry(indexEntry.first, fields, &out);
+
+    ASSERT_EQ(out.size(), 1u);
+    ASSERT_TRUE(out[0].multikey);
+    ASSERT_EQ(out[0].multikeyPaths.size(), 2);
+    ASSERT((out[0].multikeyPaths[0] == MultikeyComponents{0u, 1u}));  // a and a.b are multikey
+    ASSERT((out[0].multikeyPaths[1] == MultikeyComponents{0u, 2u}));  // c and c.d.e are multikey
+    ASSERT_BSONOBJ_EQ(out[0].keyPattern, {fromjson("{'a.b': 1, 'c.d.e': 1}")});
+}
+
 /*************************************** end section ***************************************/
 
 // translateWildcardIndexBoundsAndTightness