SERVER-39764 Fix plan caching for negation of equality to array

5 files changed, 210 insertions, 14 deletions

diff --git a/jstests/noPassthroughWithMongod/ne_array_indexability.js b/jstests/noPassthroughWithMongod/ne_array_indexability.js
new file mode 100644
index 00000000000..7e869e6beb1
--- /dev/null
+++ b/jstests/noPassthroughWithMongod/ne_array_indexability.js
@@ -0,0 +1,47 @@
+/**
+ * Test that $ne: [] queries are cached correctly. See SERVER-39764.
+ */
+(function() {
+    "use strict";
+    load("jstests/libs/analyze_plan.js");  // For planHasStage().
+
+    const coll = db.ne_array_indexability;
+    coll.drop();
+
+    coll.createIndex({"obj": 1});
+    coll.createIndex({"obj": 1, "abc": 1});
+
+    assert.commandWorked(coll.insert({obj: "hi there"}));
+
+    // Check that 'queryToCache' and 'queryToRunAfterCaching' don't share a cache entry, and get
+    // different plans.
+    function runTest(queryToCache, queryToRunAfterCaching) {
+        assert.eq(coll.find(queryToCache).itcount(), 1);
+        assert.eq(coll.find(queryToCache).itcount(), 1);
+
+        // Be sure the query has a plan in the cache.
+        const res =
+            coll.runCommand('planCacheListPlans', {query: queryToCache, sort: {}, projection: {}});
+        assert.commandWorked(res);
+        assert.gt(res.plans.length, 0);
+
+        const explForFirstQuery = coll.find(queryToCache).explain();
+        assert.eq(planHasStage(db, explForFirstQuery.queryPlanner.winningPlan, "IXSCAN"),
+                  true,
+                  explForFirstQuery);
+
+        assert.eq(coll.find(queryToRunAfterCaching).itcount(), 1);
+
+        const explForSecondQuery = coll.find(queryToRunAfterCaching).explain();
+        assert.eq(planHasStage(db, explForSecondQuery.queryPlanner.winningPlan, "COLLSCAN"),
+                  true,
+                  explForSecondQuery);
+    }
+
+    runTest({'obj': {$ne: 'def'}}, {'obj': {$ne: [[1]]}});
+
+    // Clear the cache.
+    assert.commandWorked(coll.runCommand('planCacheClear'));
+
+    runTest({'obj': {$nin: ['abc', 'def']}}, {'obj': {$nin: [[1], 'abc']}});
+})();
diff --git a/src/mongo/db/query/plan_cache.cpp b/src/mongo/db/query/plan_cache.cpp
index 851b2118921..29f8d603976 100644
--- a/src/mongo/db/query/plan_cache.cpp
+++ b/src/mongo/db/query/plan_cache.cpp
@@ -47,6 +47,7 @@
 #include "mongo/db/matcher/expression_geo.h"
 #include "mongo/db/query/collation/collator_interface.h"
 #include "mongo/db/query/plan_ranker.h"
+#include "mongo/db/query/planner_ixselect.h"
 #include "mongo/db/query/query_knobs.h"
 #include "mongo/db/query/query_solution.h"
 #include "mongo/util/assert_util.h"
@@ -635,14 +636,22 @@ void PlanCache::encodeKeyForMatch(const MatchExpression* tree, StringBuilder* ke
     }
 
     // Encode indexability.
-    const IndexToDiscriminatorMap& discriminators =
-        _indexabilityState.getDiscriminators(tree->path());
-    if (!discriminators.empty()) {
-        *keyBuilder << kEncodeDiscriminatorsBegin;
-        // For each discriminator on this path, append the character '0' or '1'.
-        for (auto&& indexAndDiscriminatorPair : discriminators) {
-            *keyBuilder << indexAndDiscriminatorPair.second.isMatchCompatibleWithIndex(tree);
+    if (!tree->path().empty()) {
+        const IndexToDiscriminatorMap& discriminators =
+            _indexabilityState.getDiscriminators(tree->path());
+        if (!discriminators.empty()) {
+            *keyBuilder << kEncodeDiscriminatorsBegin;
+            // For each discriminator on this path, append the character '0' or '1'.
+            for (auto&& indexAndDiscriminatorPair : discriminators) {
+                *keyBuilder << indexAndDiscriminatorPair.second.isMatchCompatibleWithIndex(tree);
+            }
+            *keyBuilder << kEncodeDiscriminatorsEnd;
         }
+    } else if (tree->matchType() == MatchExpression::MatchType::NOT) {
+        // If the node is not compatible with any type of index, add a single '0' discriminator
+        // here. Otherwise add a '1'.
+        *keyBuilder << kEncodeDiscriminatorsBegin;
+        *keyBuilder << QueryPlannerIXSelect::logicalNodeMayBeSupportedByAnIndex(tree);
         *keyBuilder << kEncodeDiscriminatorsEnd;
     }
 
diff --git a/src/mongo/db/query/plan_cache_test.cpp b/src/mongo/db/query/plan_cache_test.cpp
index 0e4026ac801..c79341deaba 100644
--- a/src/mongo/db/query/plan_cache_test.cpp
+++ b/src/mongo/db/query/plan_cache_test.cpp
@@ -1515,14 +1515,14 @@ TEST(PlanCacheTest, ComputeKeyMatchInDependsOnPresenceOfRegexAndFlags) {
                    "ina_re/ims/");
 
     // Test that $not-$in-$regex similarly records the presence and flags of any regexes.
-    testComputeKey("{a: {$not: {$in: [1, 'foo']}}}", "{}", "{}", "nt[ina]");
-    testComputeKey("{a: {$not: {$in: [1, /foo/]}}}", "{}", "{}", "nt[ina_re]");
+    testComputeKey("{a: {$not: {$in: [1, 'foo']}}}", "{}", "{}", "nt<1>[ina]");
+    testComputeKey("{a: {$not: {$in: [1, /foo/]}}}", "{}", "{}", "nt<1>[ina_re]");
     testComputeKey(
-        "{a: {$not: {$in: [1, /foo/i, /bar/i, /baz/msi]}}}", "{}", "{}", "nt[ina_re/ims/]");
+        "{a: {$not: {$in: [1, /foo/i, /bar/i, /baz/msi]}}}", "{}", "{}", "nt<1>[ina_re/ims/]");
 
     // Test that a $not-$in containing a single regex is unwrapped to $not-$regex.
-    testComputeKey("{a: {$not: {$in: [/foo/]}}}", "{}", "{}", "nt[rea]");
-    testComputeKey("{a: {$not: {$in: [/foo/i]}}}", "{}", "{}", "nt[rea/i/]");
+    testComputeKey("{a: {$not: {$in: [/foo/]}}}", "{}", "{}", "nt<1>[rea]");
+    testComputeKey("{a: {$not: {$in: [/foo/i]}}}", "{}", "{}", "nt<1>[rea/i/]");
 }
 
 // When a sparse index is present, computeKey() should generate different keys depending on
@@ -1631,4 +1631,132 @@ TEST(PlanCacheTest, ComputeKeyCollationIndex) {
               planCache.computeKey(*inContainsStringHasCollation));
 }
 
+/**
+ * Check that the given plan cache key has the discriminators in the given order.
+ */
+bool hasDiscriminators(const std::string& key, const std::vector<std::string>& discriminators) {
+    size_t pos = 0;
+    for (auto&& discriminator : discriminators) {
+        size_t foundPos = key.find(discriminator, pos);
+        if (foundPos == std::string::npos) {
+            return false;
+        }
+        pos = foundPos + discriminator.size();
+    }
+
+    // Return whether the key has any more discriminators.
+    return key.find("<", pos) == std::string::npos;
+}
+
+TEST(PlanCacheTest, ComputeKeyNotEqualsArray) {
+    PlanCache planCache;
+    unique_ptr<CanonicalQuery> cqNeArray(canonicalize("{a: {$ne: [1]}}"));
+    unique_ptr<CanonicalQuery> cqNeScalar(canonicalize("{a: {$ne: 123}}"));
+
+    const PlanCacheKey noIndexNeArrayKey = planCache.computeKey(*cqNeArray);
+    const PlanCacheKey noIndexNeScalarKey = planCache.computeKey(*cqNeScalar);
+    ASSERT_TRUE(hasDiscriminators(noIndexNeArrayKey, {"<0>"}));
+    ASSERT_TRUE(hasDiscriminators(noIndexNeScalarKey, {"<1>"}));
+    ASSERT_NE(noIndexNeScalarKey, noIndexNeArrayKey);
+
+    // Create a normal btree index. It will have a discriminator.
+    IndexEntry entry(BSON("a" << 1),
+                     false,    // multikey
+                     false,    // sparse
+                     false,    // unique
+                     "",       // name
+                     nullptr,  // filterExpr
+                     BSONObj());
+    planCache.notifyOfIndexEntries({entry});
+
+    const PlanCacheKey withIndexNeArrayKey = planCache.computeKey(*cqNeArray);
+    const PlanCacheKey withIndexNeScalarKey = planCache.computeKey(*cqNeScalar);
+
+    ASSERT_NE(noIndexNeArrayKey, withIndexNeArrayKey);
+
+    // There will be one discriminator for the $not and another for the leaf node ({$eq: 123}).
+    // Both will be <1>.
+    ASSERT_TRUE(hasDiscriminators(withIndexNeScalarKey, {"<1>", "<1>"}));
+
+    // There will be one discriminator for the $not and another for the leaf node ({$eq: [1]}).
+    // Since the index can support equality to an array, the second discriminator will have a value
+    // of '1'.
+    ASSERT_TRUE(hasDiscriminators(withIndexNeArrayKey, {"<0>", "<1>"}));
+}
+
+TEST(PlanCacheTest, ComputeKeyNinArray) {
+    PlanCache planCache;
+    unique_ptr<CanonicalQuery> cqNinArray(canonicalize("{a: {$nin: [123, [1]]}}"));
+    unique_ptr<CanonicalQuery> cqNinScalar(canonicalize("{a: {$nin: [123, 456]}}"));
+
+    const PlanCacheKey noIndexNinArrayKey = planCache.computeKey(*cqNinArray);
+    const PlanCacheKey noIndexNinScalarKey = planCache.computeKey(*cqNinScalar);
+    ASSERT_TRUE(hasDiscriminators(noIndexNinArrayKey, {"<0>"}));
+    ASSERT_TRUE(hasDiscriminators(noIndexNinScalarKey, {"<1>"}));
+    ASSERT_NE(noIndexNinScalarKey, noIndexNinArrayKey);
+
+    // Create a normal btree index. It will have a discriminator.
+    IndexEntry entry(BSON("a" << 1),
+                     false,    // multikey
+                     false,    // sparse
+                     false,    // unique
+                     "",       // name
+                     nullptr,  // filterExpr
+                     BSONObj());
+    planCache.notifyOfIndexEntries({entry});
+
+    const PlanCacheKey withIndexNinArrayKey = planCache.computeKey(*cqNinArray);
+    const PlanCacheKey withIndexNinScalarKey = planCache.computeKey(*cqNinScalar);
+
+    // The index discriminators should have changed.
+    ASSERT_NE(noIndexNinArrayKey, withIndexNinArrayKey);
+
+    // Both cache keys have two discriminators: one for the $not and one for the $in. The query
+    // which has an array inside the $in won't be able to use the index, and should have a <0>
+    // associated with the $not.
+    ASSERT_TRUE(hasDiscriminators(withIndexNinScalarKey, {"<1>", "<1>"}));
+    ASSERT_TRUE(hasDiscriminators(withIndexNinArrayKey, {"<0>", "<1>"}));
+}
+
+TEST(PlanCacheTest, ComputeNeArrayInOrStagePreserversOrder) {
+    PlanCache planCache;
+    unique_ptr<CanonicalQuery> cqNeA(canonicalize("{$or: [{a: {$ne: 5}}, {a: {$ne: [12]}}]}"));
+    unique_ptr<CanonicalQuery> cqNeB(canonicalize("{$or: [{a: {$ne: [12]}}, {a: {$ne: 5}}]}"));
+
+    const PlanCacheKey keyA = planCache.computeKey(*cqNeA);
+    const PlanCacheKey keyB = planCache.computeKey(*cqNeB);
+    ASSERT_NE(keyA, keyB);
+
+    // Create a normal btree index. It will have a discriminator.
+    IndexEntry entry(BSON("a" << 1),
+                     false,    // multikey
+                     false,    // sparse
+                     false,    // unique
+                     "",       // name
+                     nullptr,  // filterExpr
+                     BSONObj());
+    planCache.notifyOfIndexEntries({entry});
+
+    const PlanCacheKey keyAWithIndex = planCache.computeKey(*cqNeA);
+    const PlanCacheKey keyBWithIndex = planCache.computeKey(*cqNeB);
+
+    ASSERT_NE(keyAWithIndex, keyBWithIndex);
+    ASSERT_TRUE(hasDiscriminators(keyAWithIndex,
+                                  {
+                                      "<1>",  // $ne
+                                      "<1>",  // 5
+
+                                      "<0>",  // $ne
+                                      "<1>"   // [12]
+                                  }));
+    ASSERT_TRUE(hasDiscriminators(keyBWithIndex,
+                                  {
+                                      "<0>",  // $ne
+                                      "<1>",  // [12]
+
+                                      "<1>",  // $ne
+                                      "<1>"   // 5
+                                  }));
+}
+
 }  // namespace
diff --git a/src/mongo/db/query/planner_ixselect.cpp b/src/mongo/db/query/planner_ixselect.cpp
index 57544c2f2cf..260b90fac0a 100644
--- a/src/mongo/db/query/planner_ixselect.cpp
+++ b/src/mongo/db/query/planner_ixselect.cpp
@@ -73,7 +73,7 @@ bool elemMatchValueCompatible(const BSONElement& elt,
 
 // Can't index negations of {$eq: <Array>} or {$in: [<Array>, ...]}. Note that we could
 // use the index in principle, though we would need to generate special bounds.
-bool isNotEqualsArrayOrNotInArray(const MatchExpression* me) {
+bool isEqualsArrayOrNotInArray(const MatchExpression* me) {
     const auto type = me->matchType();
     if (type == MatchExpression::EQ) {
         return static_cast<const ComparisonMatchExpression*>(me)->getData().type() ==
@@ -297,7 +297,7 @@ bool QueryPlannerIXSelect::compatible(const BSONElement& elt,
 
             // Can't index negations of {$eq: <Array>} or {$in: [<Array>, ...]}. Note that we could
             // use the index in principle, though we would need to generate special bounds.
-            if (isNotEqualsArrayOrNotInArray(child)) {
+            if (isEqualsArrayOrNotInArray(child)) {
                 return false;
             }
 
@@ -936,4 +936,9 @@ void QueryPlannerIXSelect::stripInvalidAssignmentsTo2dsphereIndices(
     }
 }
 
+bool QueryPlannerIXSelect::logicalNodeMayBeSupportedByAnIndex(const MatchExpression* queryExpr) {
+    return !(queryExpr->matchType() == MatchExpression::NOT &&
+             isEqualsArrayOrNotInArray(queryExpr->getChild(0)));
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/query/planner_ixselect.h b/src/mongo/db/query/planner_ixselect.h
index 8d6a827b7e6..439ede7a196 100644
--- a/src/mongo/db/query/planner_ixselect.h
+++ b/src/mongo/db/query/planner_ixselect.h
@@ -137,6 +137,13 @@ public:
      */
     static bool indexedFieldHasMultikeyComponents(StringData indexedField, const IndexEntry& index);
 
+    /**
+     * Some types of matches are not supported by any type of index. If this function returns
+     * false, then 'queryExpr' is definitely not supported for any type of index. If the function
+     * returns true then 'queryExpr' may (or may not) be supported by some index.
+     */
+    static bool logicalNodeMayBeSupportedByAnIndex(const MatchExpression* queryExpr);
+
 private:
     /**
      * Amend the RelevantTag lists for all predicates in the subtree rooted at 'node' to remove