SERVER-48614 Fix plan cache discriminators for partial wildcard indexes

Previously, discrimination based on the partial filter
expression was done for all paths included in the wildcard
projection. This could lead to a situation where two queries
were erroneously assigned the same plan cache key.

The fix is to ensure that for wildcard indexes, partial
index discriminators are instead registered only for those
paths mentioned in the partial filter expression. Unlike
other kinds of wildcard index discriminators (e.g. handling
concerns of null equality or collation), the paths in the
partial filter expression are known a priori. Therefore,
discrimination based on the partial filter can be done in
the same way for wildcard and non-wildcard indexes.

4 files changed, 215 insertions, 28 deletions

diff --git a/src/mongo/db/query/plan_cache_indexability.cpp b/src/mongo/db/query/plan_cache_indexability.cpp
index 4846cd85778..216b714dbed 100644
--- a/src/mongo/db/query/plan_cache_indexability.cpp
+++ b/src/mongo/db/query/plan_cache_indexability.cpp
@@ -119,7 +119,7 @@ void PlanCacheIndexabilityState::processWildcardIndex(const CoreIndexInfo& cii)
     invariant(cii.type == IndexType::INDEX_WILDCARD);
 
     _wildcardIndexDiscriminators.emplace_back(
-        cii.wildcardProjection->exec(), cii.identifier.catalogName, cii.filterExpr, cii.collator);
+        cii.wildcardProjection->exec(), cii.identifier.catalogName, cii.collator);
 }
 
 void PlanCacheIndexabilityState::processIndexCollation(const std::string& indexName,
@@ -159,10 +159,6 @@ IndexToDiscriminatorMap PlanCacheIndexabilityState::buildWildcardDiscriminators(
             cid.addDiscriminator(QueryPlannerIXSelect::nodeIsSupportedByWildcardIndex);
             cid.addDiscriminator(nodeIsConservativelySupportedBySparseIndex);
             cid.addDiscriminator(getCollatedIndexDiscriminator(wildcardDiscriminator.collator));
-            if (wildcardDiscriminator.filterExpr) {
-                cid.addDiscriminator(
-                    getPartialIndexDiscriminator(wildcardDiscriminator.filterExpr));
-            }
         }
     }
     return ret;
@@ -174,19 +170,28 @@ void PlanCacheIndexabilityState::updateDiscriminators(
     _wildcardIndexDiscriminators.clear();
 
     for (const auto& idx : indexCores) {
-        if (idx.type == IndexType::INDEX_WILDCARD) {
-            processWildcardIndex(idx);
-            continue;
-        }
-
-        if (idx.sparse) {
-            processSparseIndex(idx.identifier.catalogName, idx.keyPattern);
-        }
+        // If necessary, add discriminators for the paths mentioned in the partial filter
+        // expression. Unlike most of the discriminator logic, this is shared for wildcard and
+        // non-wildcard indexes.
         if (idx.filterExpr) {
             processPartialIndex(idx.identifier.catalogName, idx.filterExpr);
         }
 
-        processIndexCollation(idx.identifier.catalogName, idx.keyPattern, idx.collator);
+        if (idx.type == IndexType::INDEX_WILDCARD) {
+            // The set of paths for which we should add disciminators for wildcard indexes (outside
+            // of those paths mentioned in the partial filter expression) is not known a priori.
+            // Instead, we just record some information about the wildcard index so that the
+            // discriminators can be constructed on demand at query runtime.
+            processWildcardIndex(idx);
+        } else {
+            // If the index is not wildcard, add discriminators for fields mentioned in the key
+            // pattern.
+            if (idx.sparse) {
+                processSparseIndex(idx.identifier.catalogName, idx.keyPattern);
+            }
+
+            processIndexCollation(idx.identifier.catalogName, idx.keyPattern, idx.collator);
+        }
     }
 }
 
diff --git a/src/mongo/db/query/plan_cache_indexability.h b/src/mongo/db/query/plan_cache_indexability.h
index e98e7101fc8..9bc03494865 100644
--- a/src/mongo/db/query/plan_cache_indexability.h
+++ b/src/mongo/db/query/plan_cache_indexability.h
@@ -119,16 +119,11 @@ private:
     struct WildcardIndexDiscriminatorContext {
         WildcardIndexDiscriminatorContext(projection_executor::ProjectionExecutor* proj,
                                           std::string name,
-                                          const MatchExpression* filter,
                                           const CollatorInterface* coll)
-            : projectionExec(proj),
-              filterExpr(filter),
-              collator(coll),
-              catalogName(std::move(name)) {}
+            : projectionExec(proj), collator(coll), catalogName(std::move(name)) {}
 
         // These are owned by the catalog.
         projection_executor::ProjectionExecutor* projectionExec;
-        const MatchExpression* filterExpr;  // For partial indexes.
         const CollatorInterface* collator;
 
         std::string catalogName;
diff --git a/src/mongo/db/query/plan_cache_indexability_test.cpp b/src/mongo/db/query/plan_cache_indexability_test.cpp
index 2748bfd85ac..f76a7797e1c 100644
--- a/src/mongo/db/query/plan_cache_indexability_test.cpp
+++ b/src/mongo/db/query/plan_cache_indexability_test.cpp
@@ -595,15 +595,44 @@ TEST(PlanCacheIndexabilityTest, WildcardPartialIndexDiscriminator) {
     ASSERT_EQ(1U, discriminatorsA.size());
     ASSERT(discriminatorsA.find("indexName") != discriminatorsA.end());
 
-    const auto disc = discriminatorsA["indexName"];
-
-    // Match expression which queries for a value not included by the filter expression cannot use
-    // the index.
-    ASSERT_FALSE(disc.isMatchCompatibleWithIndex(parseMatchExpression(fromjson("{a: 0}")).get()));
+    const auto wildcardDiscriminators = discriminatorsA["indexName"];
+
+    // Since the fields in the partialFilterExpression are known a priori, they are _not_ part of
+    // the wildcard-discriminators, but rather the regular discriminators. Here we show that the
+    // wildcard discriminators consider all expressions on fields 'a' or 'b' to be compatible.
+    ASSERT_TRUE(wildcardDiscriminators.isMatchCompatibleWithIndex(
+        parseMatchExpression(fromjson("{a: 0}")).get()));
+    ASSERT_TRUE(wildcardDiscriminators.isMatchCompatibleWithIndex(
+        parseMatchExpression(fromjson("{a: 6}")).get()));
+    ASSERT_TRUE(wildcardDiscriminators.isMatchCompatibleWithIndex(
+        parseMatchExpression(fromjson("{b: 0}")).get()));
+    ASSERT_TRUE(wildcardDiscriminators.isMatchCompatibleWithIndex(
+        parseMatchExpression(fromjson("{b: 6}")).get()));
+
+    // The regular (non-wildcard) set of discriminators for the path "a" should reflect whether a
+    // predicate on "a" is compatible with the partial filter expression.
+    {
+        discriminatorsA = state.getDiscriminators("a");
+        auto discriminatorsIt = discriminatorsA.find("indexName");
+        ASSERT(discriminatorsIt != discriminatorsA.end());
+        auto disc = discriminatorsIt->second;
+
+        ASSERT_FALSE(
+            disc.isMatchCompatibleWithIndex(parseMatchExpression(fromjson("{a: 0}")).get()));
+        ASSERT_TRUE(
+            disc.isMatchCompatibleWithIndex(parseMatchExpression(fromjson("{a: 6}")).get()));
+
+        ASSERT_FALSE(
+            disc.isMatchCompatibleWithIndex(parseMatchExpression(fromjson("{b: 0}")).get()));
+        ASSERT_FALSE(
+            disc.isMatchCompatibleWithIndex(parseMatchExpression(fromjson("{b: 6}")).get()));
+    }
 
-    // Match expression which queries for a value included by the filter expression does not get
-    // discriminated out.
-    ASSERT_TRUE(disc.isMatchCompatibleWithIndex(parseMatchExpression(fromjson("{a: 6}")).get()));
+    // There shouldn't be any regular discriminators associated with path "b".
+    {
+        auto&& discriminatorsB = state.getDiscriminators("b");
+        ASSERT_FALSE(discriminatorsB.count("indexName"));
+    }
 }
 
 TEST(PlanCacheIndexabilityTest,
diff --git a/src/mongo/db/query/plan_cache_test.cpp b/src/mongo/db/query/plan_cache_test.cpp
index 41337bd9cc5..fc535b0db54 100644
--- a/src/mongo/db/query/plan_cache_test.cpp
+++ b/src/mongo/db/query/plan_cache_test.cpp
@@ -2013,6 +2013,164 @@ TEST(PlanCacheTest, ComputeKeyWildcardIndexDiscriminatesEqualityToEmptyObj) {
     ASSERT_EQ(planCache.computeKey(*inWithEmptyObj).getUnstablePart(), "<1>");
 }
 
+TEST(PlanCacheTest, ComputeKeyWildcardDiscriminatesCorrectlyBasedOnPartialFilterExpression) {
+    BSONObj filterObj = BSON("x" << BSON("$gt" << 0));
+    std::unique_ptr<MatchExpression> filterExpr(parseMatchExpression(filterObj));
+
+    auto entryProjUpdatePair = makeWildcardUpdate(BSON("$**" << 1));
+    auto indexInfo = std::move(entryProjUpdatePair.first);
+    indexInfo.filterExpr = filterExpr.get();
+
+    PlanCache planCache;
+    planCache.notifyOfIndexUpdates({indexInfo});
+
+    // Test that queries on field 'x' are discriminated based on their relationship with the partial
+    // filter expression.
+    {
+        auto compatibleWithFilter = canonicalize("{x: {$eq: 5}}");
+        auto incompatibleWithFilter = canonicalize("{x: {$eq: -5}}");
+        auto compatibleKey = planCache.computeKey(*compatibleWithFilter);
+        auto incompatibleKey = planCache.computeKey(*incompatibleWithFilter);
+
+        assertPlanCacheKeysUnequalDueToDiscriminators(compatibleKey, incompatibleKey);
+        // The discriminator strings have the format "<xx>". That is, there are two discriminator
+        // bits for the "x" predicate, the first pertaining to the partialFilterExpression and the
+        // second around applicability to the wildcard index.
+        ASSERT_EQ(compatibleKey.getUnstablePart(), "<11>");
+        ASSERT_EQ(incompatibleKey.getUnstablePart(), "<01>");
+    }
+
+    // The partialFilterExpression should lead to a discriminator over field 'x', but not over 'y'.
+    // (Separately, there are wildcard-related discriminator bits for both 'x' and 'y'.)
+    {
+        auto compatibleWithFilter = canonicalize("{x: {$eq: 5}, y: 1}");
+        auto incompatibleWithFilter = canonicalize("{x: {$eq: -5}, y: 1}");
+        auto compatibleKey = planCache.computeKey(*compatibleWithFilter);
+        auto incompatibleKey = planCache.computeKey(*incompatibleWithFilter);
+
+        assertPlanCacheKeysUnequalDueToDiscriminators(compatibleKey, incompatibleKey);
+        // The discriminator strings have the format "<xx><y>". That is, there are two discriminator
+        // bits for the "x" predicate (the first pertaining to the partialFilterExpression, the
+        // second around applicability to the wildcard index) and one discriminator bit for "y".
+        ASSERT_EQ(compatibleKey.getUnstablePart(), "<11><1>");
+        ASSERT_EQ(incompatibleKey.getUnstablePart(), "<01><1>");
+    }
+
+    // $eq:null predicates cannot be assigned to a wildcard index. Make sure that this is
+    // discrimated correctly. This test is designed to reproduce SERVER-48614.
+    {
+        auto compatibleQuery = canonicalize("{x: {$eq: 5}, y: 1}");
+        auto incompatibleQuery = canonicalize("{x: {$eq: 5}, y: null}");
+        auto compatibleKey = planCache.computeKey(*compatibleQuery);
+        auto incompatibleKey = planCache.computeKey(*incompatibleQuery);
+
+        assertPlanCacheKeysUnequalDueToDiscriminators(compatibleKey, incompatibleKey);
+        // The discriminator strings have the format "<xx><y>". That is, there are two discriminator
+        // bits for the "x" predicate (the first pertaining to the partialFilterExpression, the
+        // second around applicability to the wildcard index) and one discriminator bit for "y".
+        ASSERT_EQ(compatibleKey.getUnstablePart(), "<11><1>");
+        ASSERT_EQ(incompatibleKey.getUnstablePart(), "<11><0>");
+    }
+
+    // Test that the discriminators are correct for an $eq:null predicate on 'x'. This predicate is
+    // imcompatible for two reasons: null equality predicates cannot be answered by wildcard
+    // indexes, and the predicate is not compatible with the partial filter expression. This should
+    // result in two "0" bits inside the discriminator string.
+    {
+        auto key = planCache.computeKey(*canonicalize("{x: {$eq: null}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<00>");
+    }
+}
+
+TEST(PlanCacheTest, ComputeKeyWildcardDiscriminatesCorrectlyWithPartialFilterAndExpression) {
+    // Partial filter is an AND of multiple conditions.
+    BSONObj filterObj = BSON("x" << BSON("$gt" << 0) << "y" << BSON("$gt" << 0));
+    std::unique_ptr<MatchExpression> filterExpr(parseMatchExpression(filterObj));
+
+    auto entryProjUpdatePair = makeWildcardUpdate(BSON("$**" << 1));
+    auto indexInfo = std::move(entryProjUpdatePair.first);
+    indexInfo.filterExpr = filterExpr.get();
+
+    PlanCache planCache;
+    planCache.notifyOfIndexUpdates({indexInfo});
+
+    {
+        // The discriminators should have the format <xx><yy><z>. The 'z' predicate has just one
+        // discriminator because it is not referenced in the partial filter expression.  All
+        // predicates are compatible.
+        auto key = planCache.computeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, z: {$eq: 3}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<11><11><1>");
+    }
+
+    {
+        // The discriminators should have the format <xx><yy><z>. The 'y' predicate is not
+        // compatible with the partial filter expression, leading to one of the 'y' bits being set
+        // to zero.
+        auto key = planCache.computeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: -2}, z: {$eq: 3}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<11><01><1>");
+    }
+}
+
+TEST(PlanCacheTest, ComputeKeyWildcardDiscriminatesCorrectlyWithPartialFilterOnNestedField) {
+    BSONObj filterObj = BSON("x.y" << BSON("$gt" << 0));
+    std::unique_ptr<MatchExpression> filterExpr(parseMatchExpression(filterObj));
+
+    auto entryProjUpdatePair = makeWildcardUpdate(BSON("$**" << 1));
+    auto indexInfo = std::move(entryProjUpdatePair.first);
+    indexInfo.filterExpr = filterExpr.get();
+
+    PlanCache planCache;
+    planCache.notifyOfIndexUpdates({indexInfo});
+
+    {
+        // The discriminators have the format <x><(x.y)(x.y)<y>. All predicates are compatible
+        auto key =
+            planCache.computeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, 'x.y': {$eq: 3}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<1><11><1>");
+    }
+
+    {
+        // Here, the predicate on "x.y" is not compatible with the partial filter expression.
+        auto key =
+            planCache.computeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, 'x.y': {$eq: -3}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<1><01><1>");
+    }
+}
+
+TEST(PlanCacheTest, ComputeKeyDiscriminatesCorrectlyWithPartialFilterAndWildcardProjection) {
+    BSONObj filterObj = BSON("x" << BSON("$gt" << 0));
+    std::unique_ptr<MatchExpression> filterExpr(parseMatchExpression(filterObj));
+
+    auto entryProjUpdatePair = makeWildcardUpdate(BSON("y.$**" << 1));
+    auto indexInfo = std::move(entryProjUpdatePair.first);
+    indexInfo.filterExpr = filterExpr.get();
+
+    PlanCache planCache;
+    planCache.notifyOfIndexUpdates({indexInfo});
+
+    {
+        // The discriminators have the format <x><y>. The discriminator for 'x' indicates whether
+        // the predicate is compatible with the partial filter expression, whereas the disciminator
+        // for 'y' is about compatibility with the wildcard index.
+        auto key = planCache.computeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, z: {$eq: 3}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<1><1>");
+    }
+
+    {
+        // Similar to the previous case, except with an 'x' predicate that is incompatible with the
+        // partial filter expression.
+        auto key = planCache.computeKey(*canonicalize("{x: {$eq: -1}, y: {$eq: 2}, z: {$eq: 3}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<0><1>");
+    }
+
+    {
+        // Case where the 'y' predicate is not compatible with the wildcard index.
+        auto key =
+            planCache.computeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: null}, z: {$eq: 3}}"));
+        ASSERT_EQ(key.getUnstablePart(), "<1><0>");
+    }
+}
+
 TEST(PlanCacheTest, StableKeyDoesNotChangeAcrossIndexCreation) {
     PlanCache planCache;
     unique_ptr<CanonicalQuery> cq(canonicalize("{a: 0}}"));