SERVER-35330 expand allPaths indexes in planner

4 files changed, 205 insertions, 23 deletions

diff --git a/src/mongo/db/index_names.cpp b/src/mongo/db/index_names.cpp
index 557a7c636bb..7773bec86a5 100644
--- a/src/mongo/db/index_names.cpp
+++ b/src/mongo/db/index_names.cpp
@@ -43,16 +43,24 @@ const string IndexNames::HASHED = "hashed";
 const string IndexNames::BTREE = "";
 const string IndexNames::ALLPATHS = "allPaths";
 
+const StringMap<IndexType> kIndexNameToType = {
+    {IndexNames::GEO_2D, INDEX_2D},
+    {IndexNames::GEO_HAYSTACK, INDEX_HAYSTACK},
+    {IndexNames::GEO_2DSPHERE, INDEX_2DSPHERE},
+    {IndexNames::TEXT, INDEX_TEXT},
+    {IndexNames::HASHED, INDEX_HASHED},
+    {IndexNames::ALLPATHS, INDEX_ALLPATHS},
+};
+
 // static
 string IndexNames::findPluginName(const BSONObj& keyPattern) {
     BSONObjIterator i(keyPattern);
     while (i.more()) {
         BSONElement e = i.next();
-        string fieldName(e.fieldName());
+        StringData fieldName(e.fieldNameStringData());
         if (String == e.type()) {
             return e.String();
-        } else if ((fieldName == "$**") ||
-                   (fieldName.size() > 4 && fieldName.substr(e.fieldNameSize() - 5) == ".$**")) {
+        } else if ((fieldName == "$**") || fieldName.endsWith(".$**")) {
             return IndexNames::ALLPATHS;
         } else
             continue;
@@ -76,20 +84,12 @@ bool IndexNames::isKnownName(const string& name) {
 }
 
 // static
-IndexType IndexNames::nameToType(const string& accessMethod) {
-    if (IndexNames::GEO_2D == accessMethod) {
-        return INDEX_2D;
-    } else if (IndexNames::GEO_HAYSTACK == accessMethod) {
-        return INDEX_HAYSTACK;
-    } else if (IndexNames::GEO_2DSPHERE == accessMethod) {
-        return INDEX_2DSPHERE;
-    } else if (IndexNames::TEXT == accessMethod) {
-        return INDEX_TEXT;
-    } else if (IndexNames::HASHED == accessMethod) {
-        return INDEX_HASHED;
-    } else {
+IndexType IndexNames::nameToType(StringData accessMethod) {
+    auto typeIt = kIndexNameToType.find(accessMethod);
+    if (typeIt == kIndexNameToType.end()) {
         return INDEX_BTREE;
     }
+    return typeIt->second;
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/index_names.h b/src/mongo/db/index_names.h
index 1680d5401aa..e02ea6f1fb8 100644
--- a/src/mongo/db/index_names.h
+++ b/src/mongo/db/index_names.h
@@ -30,6 +30,8 @@
 
 #include <string>
 
+#include "mongo/base/string_data.h"
+
 namespace mongo {
 
 class BSONObj;
@@ -46,6 +48,7 @@ class BSONObj;
  *    the data isn't what we expect.
  */
 enum IndexType {
+    INDEX_ALLPATHS,
     INDEX_BTREE,
     INDEX_2D,
     INDEX_HAYSTACK,
@@ -89,7 +92,7 @@ public:
     /**
      * Convert an index name to an IndexType.
      */
-    static IndexType nameToType(const std::string& accessMethod);
+    static IndexType nameToType(StringData accesMethod);
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/planner_ixselect.cpp b/src/mongo/db/query/planner_ixselect.cpp
index dfabd73f04d..5cb53ec978b 100644
--- a/src/mongo/db/query/planner_ixselect.cpp
+++ b/src/mongo/db/query/planner_ixselect.cpp
@@ -54,6 +54,43 @@ std::size_t numPathComponents(StringData path) {
     return FieldRef{path}.numParts();
 }
 
+
+/**
+ * Given a single allPaths index, and a set of fields which are being queried, create 'mock'
+ * IndexEntry for each of the appropriate fields.
+ */
+void expandIndex(const IndexEntry& allPathsIndex,
+                 const stdx::unordered_set<std::string>& fields,
+                 vector<IndexEntry>* out) {
+    invariant(out);
+
+    // TODO SERVER-35902: For now indexAllPaths indices cannot supply a projection.
+
+    out->reserve(out->size() + fields.size());
+    for (auto&& fieldName : fields) {
+        // TODO: This needs to change in SERVER-35902.
+        if (fieldName == "_id") {
+            continue;
+        }
+
+        IndexEntry entry(BSON(fieldName << 1),
+                         IndexNames::ALLPATHS,
+                         false,  // multikey (TODO SERVER-36109)
+                         {},     // multikey paths
+                         true,   // sparse
+                         false,  // unique
+                         // TODO: SERVER-35333: for plan caching to work, each IndexEntry must have
+                         // a unique name. We violate that requirement here by giving each
+                         // "expanded" index the same name. This must be fixed.
+                         allPathsIndex.name,
+                         allPathsIndex.filterExpr,
+                         allPathsIndex.infoObj,
+                         allPathsIndex.collator);
+
+
+        out->push_back(std::move(entry));
+    }
+}
 }  // namespace
 
 bool QueryPlannerIXSelect::notEqualsNullCanUseIndex(const IndexEntry& index,
@@ -231,15 +268,31 @@ void QueryPlannerIXSelect::getFields(const MatchExpression* node,
 }
 
 // static
-void QueryPlannerIXSelect::findRelevantIndices(const stdx::unordered_set<string>& fields,
-                                               const vector<IndexEntry>& allIndices,
-                                               vector<IndexEntry>* out) {
-    for (size_t i = 0; i < allIndices.size(); ++i) {
-        BSONObjIterator it(allIndices[i].keyPattern);
-        verify(it.more());
+void QueryPlannerIXSelect::findRelevantIndices(const stdx::unordered_set<std::string>& fields,
+                                               const std::vector<IndexEntry>& allIndices,
+                                               std::vector<IndexEntry>* out) {
+    for (auto&& entry : allIndices) {
+        if (entry.type == INDEX_ALLPATHS) {
+            // Should only have one field of the form {"&**" : 1}.
+            // TODO SERVER-35902: This code assumes the allPaths index only indexes the root,
+            // and not a subpath (e.g. {"a.b.$**": 1}).
+            uassert(ErrorCodes::NotImplemented,
+                    "Not yet implemented: expanding allPaths indexes for an index which indexes an "
+                    "entire subpath",
+                    entry.keyPattern.firstElement().fieldNameStringData() == "$**");
+            uassert(ErrorCodes::NotImplemented,
+                    "Do not support starPathsTempName yet",
+                    entry.infoObj["starPathsTempName"].eoo());
+
+            invariant(entry.keyPattern.nFields() == 1);
+            expandIndex(entry, fields, out);
+            continue;
+        }
+
+        BSONObjIterator it(entry.keyPattern);
         BSONElement elt = it.next();
         if (fields.end() != fields.find(elt.fieldName())) {
-            out->push_back(allIndices[i]);
+            out->push_back(entry);
         }
     }
 }
diff --git a/src/mongo/db/query/planner_ixselect_test.cpp b/src/mongo/db/query/planner_ixselect_test.cpp
index ea4836de09b..24fdb76c2d7 100644
--- a/src/mongo/db/query/planner_ixselect_test.cpp
+++ b/src/mongo/db/query/planner_ixselect_test.cpp
@@ -39,6 +39,7 @@
 #include "mongo/db/query/index_tag.h"
 #include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
+#include "mongo/util/assert_util.h"
 #include "mongo/util/text.h"
 #include <memory>
 
@@ -1250,4 +1251,129 @@ TEST(QueryPlannerIXSelectTest, HashedIndexShouldNotBeRelevantForNotEqualsNullPre
     std::set<size_t> expectedIndices = {};
     testRateIndices("{a: {$ne: null}}", "", kSimpleCollator, {entry}, "a,a", expectedIndices);
 }
+
+/*
+ * Will compare 'keyPatterns' with 'entries'. As part of comparing, it will sort both of them.
+ */
+bool indexEntryKeyPatternsMatch(vector<BSONObj>* keyPatterns, vector<IndexEntry>* entries) {
+    ASSERT_EQ(entries->size(), keyPatterns->size());
+
+    const auto cmpFn = [](const IndexEntry& a, const IndexEntry& b) {
+        return SimpleBSONObjComparator::kInstance.evaluate(a.keyPattern < b.keyPattern);
+    };
+
+    std::sort(entries->begin(), entries->end(), cmpFn);
+    std::sort(keyPatterns->begin(), keyPatterns->end(), [](const BSONObj& a, const BSONObj& b) {
+        return SimpleBSONObjComparator::kInstance.evaluate(a < b);
+    });
+
+    return std::equal(keyPatterns->begin(),
+                      keyPatterns->end(),
+                      entries->begin(),
+                      [](const BSONObj& keyPattern, const IndexEntry& ie) -> bool {
+                          return SimpleBSONObjComparator::kInstance.evaluate(keyPattern ==
+                                                                             ie.keyPattern);
+                      });
+}
+
+TEST(QueryPlannerIXSelectTest, ExpandAllPathsIndices) {
+    const auto indexEntry = makeIndexEntry(BSON("$**" << 1), {});
+
+    // Case where no fields are specified.
+    std::vector<IndexEntry> result;
+    QueryPlannerIXSelect::findRelevantIndices(stdx::unordered_set<string>(), {indexEntry}, &result);
+    ASSERT_TRUE(result.empty());
+
+    stdx::unordered_set<string> fields = {"fieldA", "fieldB"};
+
+    QueryPlannerIXSelect::findRelevantIndices(fields, {indexEntry}, &result);
+    ASSERT_EQ(result.size(), 2u);
+
+    std::vector<BSONObj> expectedKeyPatterns = {BSON("fieldA" << 1), BSON("fieldB" << 1)};
+    ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result));
+
+    const auto allPathsIndexWithSubpath = makeIndexEntry(BSON("a.b.$**" << 1), {});
+    // TODO SERVER-35902: Remove this.
+    ASSERT_THROWS_CODE(
+        QueryPlannerIXSelect::findRelevantIndices(fields, {allPathsIndexWithSubpath}, &result),
+        DBException,
+        ErrorCodes::NotImplemented);
+}
+
+TEST(QueryPlannerIXSelectTest, ExpandAllPathsIndicesInPresenceOfOtherIndices) {
+    auto allPathsIndexEntry = makeIndexEntry(BSON("$**" << 1), {});
+    auto aIndexEntry = makeIndexEntry(BSON("fieldA" << 1), {});
+    auto bIndexEntry = makeIndexEntry(BSON("fieldB" << 1), {});
+    auto abIndexEntry = makeIndexEntry(BSON("fieldA" << 1 << "fieldB" << 1), {});
+
+    const stdx::unordered_set<string> fields = {"fieldA", "fieldB", "fieldC"};
+    std::vector<IndexEntry> result;
+    std::vector<BSONObj> expectedKeyPatterns = {
+        BSON("fieldA" << 1), BSON("fieldA" << 1), BSON("fieldB" << 1), BSON("fieldC" << 1)};
+
+    QueryPlannerIXSelect::findRelevantIndices(fields, {aIndexEntry, allPathsIndexEntry}, &result);
+    ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result));
+    result.clear();
+
+    expectedKeyPatterns = {
+        BSON("fieldB" << 1), BSON("fieldA" << 1), BSON("fieldB" << 1), BSON("fieldC" << 1)};
+    QueryPlannerIXSelect::findRelevantIndices(fields, {bIndexEntry, allPathsIndexEntry}, &result);
+    ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result));
+    result.clear();
+
+    QueryPlannerIXSelect::findRelevantIndices(
+        fields, {aIndexEntry, allPathsIndexEntry, bIndexEntry}, &result);
+    expectedKeyPatterns = {BSON("fieldA" << 1),
+                           BSON("fieldA" << 1),
+                           BSON("fieldB" << 1),
+                           BSON("fieldC" << 1),
+                           BSON("fieldB" << 1)};
+    ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result));
+    result.clear();
+
+    QueryPlannerIXSelect::findRelevantIndices(fields, {allPathsIndexEntry, abIndexEntry}, &result);
+    expectedKeyPatterns = {BSON("fieldA" << 1),
+                           BSON("fieldB" << 1),
+                           BSON("fieldC" << 1),
+                           BSON("fieldA" << 1 << "fieldB" << 1)};
+    ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result));
+}
+
+TEST(QueryPlannerIXSelectTest, AllPathsIndicesExcludeIdField) {
+    auto indexEntry = makeIndexEntry(BSON("$**" << 1), {});
+
+    stdx::unordered_set<string> fields = {"_id", "abc", "def"};
+    std::vector<IndexEntry> result;
+
+    QueryPlannerIXSelect::findRelevantIndices(fields, {indexEntry}, &result);
+    std::vector<BSONObj> expectedKeyPatterns = {BSON("abc" << 1), BSON("def" << 1)};
+    ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result));
+}
+
+TEST(QueryPlannerIXSelectTest, AllPathsIndicesExpandedEntryHasCorrectProperties) {
+    auto allPathsIndexEntry = makeIndexEntry(BSON("$**" << 1), {});
+
+    stdx::unordered_set<string> fields = {"abc", "def"};
+    std::vector<IndexEntry> result;
+
+    QueryPlannerIXSelect::findRelevantIndices(fields, {allPathsIndexEntry}, &result);
+    std::vector<BSONObj> expectedKeyPatterns = {BSON("abc" << 1), BSON("def" << 1)};
+    ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result));
+
+    for (auto&& ie : result) {
+        // TODO SERVER-36109. For now we assume allPaths indexes are never multikey.
+        ASSERT_FALSE(ie.multikey);
+
+        // AllPaths indices are always sparse.
+        ASSERT_TRUE(ie.sparse);
+
+        // They are never unique.
+        ASSERT_FALSE(ie.unique);
+
+        // TODO SERVER-35333: Check that the name of the generated IndexEntry is different from the
+        // original IndexEntry.
+        ASSERT_EQ(ie.name, allPathsIndexEntry.name);
+    }
+}
+
 }  // namespace