SERVER-68434 Fix plan cache key encoding to account for $or in partial index expression

(cherry picked from commit f15f2bf8958557b4e8fccc6e8e1c7c8c5834d209)

2 files changed, 192 insertions, 3 deletions

diff --git a/jstests/core/query/partial_index_logical.js b/jstests/core/query/partial_index_logical.js
new file mode 100644
index 00000000000..145eb206dd7
--- /dev/null
+++ b/jstests/core/query/partial_index_logical.js
@@ -0,0 +1,180 @@
+/**
+ * Test the planners ability to distinguish parameterized queries in the presence of a partial index
+ * containing logical expressions ($and, $or).
+ *
+ * @tags: [
+ *  # TODO SERVER-67607: Test plan cache with CQF enabled.
+ *  cqf_incompatible,
+ *  # Since the plan cache is per-node state, this test assumes that all operations are happening
+ *  # against the same mongod.
+ *  assumes_read_preference_unchanged,
+ *  assumes_read_concern_unchanged,
+ *  does_not_support_stepdowns,
+ *  # If all chunks are moved off of a shard, it can cause the plan cache to miss commands.
+ *  assumes_balancer_off,
+ *  assumes_unsharded_collection,
+ *  requires_fcv_63,
+ *  # Plan cache state is node-local and will not get migrated alongside tenant data.
+ *  tenant_migration_incompatible,
+ * ]
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/analyze_plan.js");  // For getPlanCacheKeyFromShape.
+
+(function partialIndexMixedFields() {
+    db.test.drop();
+
+    // Create enough competing indexes such that a query is eligible for caching (single plan
+    // queries are not cached).
+    assert.commandWorked(
+        db.test.createIndex({num: 1}, {partialFilterExpression: {num: 5, foo: 6}}));
+    assert.commandWorked(db.test.createIndex({num: -1}));
+    assert.commandWorked(db.test.createIndex({num: -1, not_num: 1}));
+
+    assert.commandWorked(db.test.insert([
+        {_id: 0, num: 5, foo: 6},
+        {_id: 1, num: 5, foo: 7},
+    ]));
+
+    // Run a query which is eligible to use the {num: 1} index as it is covered by the partial
+    // filter expression.
+    assert.eq(db.test.find({num: 5, foo: 6}).itcount(), 1);
+    assert.eq(db.test.find({num: 5, foo: 6}).itcount(), 1);
+    const matchingKey =
+        getPlanCacheKeyFromShape({query: {num: 5, foo: 6}, collection: db.test, db: db});
+    assert.eq(1,
+              db.test.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey: matchingKey}}])
+                  .itcount());
+
+    // This query should not be eligible for the {num: 1} index despite the path 'num' being
+    // compatible (per the plan cache key encoding).
+    assert.eq(1, db.test.find({num: 5, foo: 7}).itcount());
+    const nonCoveredKey =
+        getPlanCacheKeyFromShape({query: {num: 5, foo: 7}, collection: db.test, db: db});
+    assert.eq(1,
+              db.test.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey: nonCoveredKey}}])
+                  .itcount());
+
+    // Sanity check that the generated keys are different due to the index compatibility.
+    assert.neq(nonCoveredKey, matchingKey);
+})();
+
+(function partialIndexDisjunction() {
+    db.test.drop();
+
+    // Create enough competing indexes such that a query is eligible for caching (single plan
+    // queries are not cached).
+    assert.commandWorked(db.test.createIndex(
+        {num: 1},
+        {partialFilterExpression: {$or: [{num: {$exists: true}}, {num: {$type: 'number'}}]}}));
+    assert.commandWorked(db.test.createIndex({num: -1}));
+    assert.commandWorked(db.test.createIndex({num: -1, not_num: 1}));
+
+    assert.commandWorked(db.test.insert([
+        {_id: 0},
+        {_id: 1, num: null},
+        {_id: 2, num: 5},
+    ]));
+
+    // Run a query which is eligible to use the {num: 1} index as it is covered by the partial
+    // filter expression.
+    assert.eq(db.test.find({num: 5}).itcount(), 1);
+    assert.eq(db.test.find({num: 5}).itcount(), 1);
+    const numericKey = getPlanCacheKeyFromShape({query: {num: 5}, collection: db.test, db: db});
+    assert.eq(
+        1,
+        db.test.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey: numericKey}}]).itcount());
+
+    // The plan for the query above should now be in the cache and active. Now execute a query with
+    // a very similar shape, however the predicate parameters are not satisfied by the partial
+    // filter expression. This is because {num: null} should match both explicit null as well as
+    // missing values (the latter are not indexed).
+    assert.eq(2, db.test.find({num: null}).itcount());
+    const nullKey = getPlanCacheKeyFromShape({query: {num: null}, collection: db.test, db: db});
+    assert.eq(
+        1, db.test.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey: nullKey}}]).itcount());
+
+    // Sanity check that the generated keys are different due to the index compatibility.
+    assert.neq(nullKey, numericKey);
+})();
+
+(function partialIndexDisjunctionWithCollation() {
+    db.test.drop();
+
+    const caseInsensitive = {locale: "en_US", strength: 2};
+
+    // Create enough competing indexes such that a query is eligible for caching (single plan
+    // queries are not cached).
+    assert.commandWorked(db.test.createIndex({a: 1}, {
+        partialFilterExpression: {$or: [{a: {$gt: 0}}, {a: {$gt: ""}}]},
+        collation: caseInsensitive,
+    }));
+    assert.commandWorked(db.test.createIndex({a: -1}));
+    assert.commandWorked(db.test.createIndex({a: -1, b: 1}));
+
+    assert.commandWorked(db.test.insert([
+        {_id: 0, a: "some"},
+        {_id: 1, a: "string"},
+    ]));
+
+    // Populate the plan cache for a query which is eligible for the partial index. This is true
+    // without an explicit collation because the query text does not contain any string comparisons.
+    assert.eq(db.test.aggregate({$match: {a: {$in: [1, 3]}}}).itcount(), 0);
+    assert.eq(db.test.aggregate({$match: {a: {$in: [1, 3]}}}).itcount(), 0);
+    const simpleCollationKey =
+        getPlanCacheKeyFromShape({query: {a: {$in: [1, 3]}}, collection: db.test, db: db});
+    assert.eq(
+        1,
+        db.test.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey: simpleCollationKey}}])
+            .itcount());
+
+    // A collation-sensitive query should _not_ use the cached plan since the default simple
+    // collation does not match the collation on the index.
+    assert.eq(db.test.aggregate({$match: {a: {$in: ["a", "Some"]}}}).itcount(), 0);
+    assert.eq(db.test.aggregate({$match: {a: {$in: ["a", "Some"]}}}).itcount(), 0);
+    const collationSensitiveKey =
+        getPlanCacheKeyFromShape({query: {a: {$in: ["a", "Some"]}}, collection: db.test, db: db});
+    assert.eq(
+        1,
+        db.test.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey: collationSensitiveKey}}])
+            .itcount());
+
+    // Sanity check that the generated keys are different due to the collation and index
+    // compatibility.
+    assert.neq(collationSensitiveKey, simpleCollationKey);
+})();
+
+(function partialIndexConjunction() {
+    db.test.drop();
+
+    // Create enough competing indexes such that a query is eligible for caching (single plan
+    // queries are not cached).
+    assert.commandWorked(
+        db.test.createIndex({num: 1}, {partialFilterExpression: {num: {$gt: 0, $lt: 10}}}));
+    assert.commandWorked(db.test.createIndex({num: -1}));
+    assert.commandWorked(db.test.createIndex({num: -1, not_num: 1}));
+
+    assert.commandWorked(db.test.insert([
+        {_id: 0},
+        {_id: 1, num: 1},
+        {_id: 2, num: 11},
+    ]));
+
+    // Run a query which is eligible to use the {num: 1} index as it is covered by the partial
+    // filter expression.
+    assert.eq(db.test.find({num: {$gt: 0, $lt: 10}}).itcount(), 1);
+    assert.eq(db.test.find({num: {$gt: 0, $lt: 10}}).itcount(), 1);
+    const validKey =
+        getPlanCacheKeyFromShape({query: {num: {$gt: 0, $lt: 10}}, collection: db.test, db: db});
+    assert.eq(
+        1,
+        db.test.aggregate([{$planCacheStats: {}}, {$match: {planCacheKey: validKey}}]).itcount());
+
+    // The plan for the query above should now be in the cache and active. Now execute a query with
+    // a very similar shape, however the predicate parameters are not satisfied by the partial
+    // filter expression.
+    assert.eq(2, db.test.find({num: {$gt: 0, $lt: 12}}).itcount());
+})();
+})();
diff --git a/jstests/libs/analyze_plan.js b/jstests/libs/analyze_plan.js
index dcfb3f11221..9ae5e62ba65 100644
--- a/jstests/libs/analyze_plan.js
+++ b/jstests/libs/analyze_plan.js
@@ -526,9 +526,18 @@ function getPlanCacheKeyFromExplain(explainRes, db) {
  * Helper to run a explain on the given query shape and get the "planCacheKey" from the explain
  * result.
  */
-function getPlanCacheKeyFromShape({query = {}, projection = {}, sort = {}, collection, db}) {
-    const explainRes =
-        assert.commandWorked(collection.explain().find(query, projection).sort(sort).finish());
+function getPlanCacheKeyFromShape({
+    query = {},
+    projection = {},
+    sort = {},
+    collation = {
+        locale: "simple"
+    },
+    collection,
+    db
+}) {
+    const explainRes = assert.commandWorked(
+        collection.explain().find(query, projection).collation(collation).sort(sort).finish());
 
     return getPlanCacheKeyFromExplain(explainRes, db);
 }