SERVER-35336 Add tests for allPaths indexes being a partial index

6 files changed, 200 insertions, 7 deletions

diff --git a/jstests/noPassthroughWithMongod/all_paths_cached_plans.js b/jstests/noPassthroughWithMongod/all_paths_cached_plans.js
index 6cf576882a8..bd54cba889f 100644
--- a/jstests/noPassthroughWithMongod/all_paths_cached_plans.js
+++ b/jstests/noPassthroughWithMongod/all_paths_cached_plans.js
@@ -110,6 +110,27 @@
                    queryWithStringExplain.queryPlanner.queryHash);
     })();
 
-    // TODO SERVER-35336: Update this test to use a partial $** index, and be sure indexability
-    // discriminators also work for partial indices.
+    // Check that indexability discriminators work with partial allPaths indexes.
+    (function() {
+        assert.eq(coll.drop(), true);
+        assert.commandWorked(db.createCollection(coll.getName()));
+        assert.commandWorked(
+            coll.createIndex({"$**": 1}, {partialFilterExpression: {a: {$lte: 5}}}));
+
+        // Run a query for a value included by the partial filter expression.
+        const queryIndexedExplain = coll.find({a: 4}).explain();
+        let ixScans = getPlanStages(queryIndexedExplain.queryPlanner.winningPlan, "IXSCAN");
+        assert.eq(ixScans.length, 1);
+        assert.eq(ixScans[0].keyPattern, {$_path: 1, a: 1});
+
+        // Run a query which tries to get a value not included by the partial filter expression.
+        const queryUnindexedExplain = coll.find({a: 100}).explain();
+        ixScans = getPlanStages(queryUnindexedExplain.queryPlanner.winningPlan, "IXSCAN");
+        assert.eq(ixScans.length, 0);
+
+        // Check that the shapes are different since the query which searches for a value not
+        // included by the partial filter expression won't be eligible to use the $** index.
+        assert.neq(queryIndexedExplain.queryPlanner.queryHash,
+                   queryUnindexedExplain.queryPlanner.queryHash);
+    })();
 })();
diff --git a/jstests/noPassthroughWithMongod/all_paths_partial_index.js b/jstests/noPassthroughWithMongod/all_paths_partial_index.js
new file mode 100644
index 00000000000..8cd6f80c9db
--- /dev/null
+++ b/jstests/noPassthroughWithMongod/all_paths_partial_index.js
@@ -0,0 +1,58 @@
+/**
+ * Test that $** indexes work when provided with a partial filter expression.
+ *
+ * TODO: SERVER-36198: Move this test back to jstests/core/
+ */
+(function() {
+    "use strict";
+
+    load("jstests/libs/analyze_plan.js");  // For isIxScan, isCollscan.
+
+    const coll = db.all_paths_partial_index;
+
+    function testPartialAllPathsIndex(indexKeyPattern, indexOptions) {
+        coll.drop();
+
+        assert.commandWorked(coll.createIndex(indexKeyPattern, indexOptions));
+        assert.commandWorked(coll.insert({x: 5, a: 2}));  // Not in index.
+        assert.commandWorked(coll.insert({x: 6, a: 1}));  // In index.
+
+        // find() operations that should use the index.
+        let explain = coll.explain("executionStats").find({x: 6, a: 1}).finish();
+        assert.eq(1, explain.executionStats.nReturned);
+        assert(isIxscan(db, explain.queryPlanner.winningPlan));
+        explain = coll.explain("executionStats").find({x: {$gt: 1}, a: 1}).finish();
+        assert.eq(1, explain.executionStats.nReturned);
+        assert(isIxscan(db, explain.queryPlanner.winningPlan));
+        explain = coll.explain("executionStats").find({x: 6, a: {$lte: 1}}).finish();
+        assert.eq(1, explain.executionStats.nReturned);
+        assert(isIxscan(db, explain.queryPlanner.winningPlan));
+
+        // find() operations that should not use the index.
+        explain = coll.explain("executionStats").find({x: 6, a: {$lt: 1.6}}).finish();
+        assert.eq(1, explain.executionStats.nReturned);
+        assert(isCollscan(db, explain.queryPlanner.winningPlan));
+
+        explain = coll.explain("executionStats").find({x: 6}).finish();
+        assert.eq(1, explain.executionStats.nReturned);
+        assert(isCollscan(db, explain.queryPlanner.winningPlan));
+
+        explain = coll.explain("executionStats").find({a: {$gte: 0}}).finish();
+        assert.eq(2, explain.executionStats.nReturned);
+        assert(isCollscan(db, explain.queryPlanner.winningPlan));
+    }
+
+    try {
+        // Required in order to build $** indexes.
+        assert.commandWorked(
+            db.adminCommand({setParameter: 1, internalQueryAllowAllPathsIndexes: true}));
+
+        testPartialAllPathsIndex({"$**": 1}, {partialFilterExpression: {a: {$lte: 1.5}}});
+
+        // Case where the partial filter expression is on a field not included in the index.
+        testPartialAllPathsIndex({"x.$**": 1}, {partialFilterExpression: {a: {$lte: 1.5}}});
+    } finally {
+        // Disable $** indexes once the tests have either completed or failed.
+        db.adminCommand({setParameter: 1, internalQueryAllowAllPathsIndexes: false});
+    }
+})();
diff --git a/src/mongo/db/query/plan_cache_indexability_test.cpp b/src/mongo/db/query/plan_cache_indexability_test.cpp
index bf409593859..6531b5cae30 100644
--- a/src/mongo/db/query/plan_cache_indexability_test.cpp
+++ b/src/mongo/db/query/plan_cache_indexability_test.cpp
@@ -500,7 +500,35 @@ TEST(PlanCacheIndexabilityTest, AllPathsWithCollationDiscriminator) {
         parseMatchExpression(fromjson("{a: \"hello world\"}"), &collator).get()));
 }
 
-// TODO SERVER-35336: Update this test to use a partial $** index, and be sure indexability
-// discriminators also work for partial indices.
+TEST(PlanCacheIndexabilityTest, AllPathsPartialIndexDiscriminator) {
+    PlanCacheIndexabilityState state;
+
+    // Need to keep the filter BSON object around for the duration of the test since the match
+    // expression will store (unowned) pointers into it.
+    BSONObj filterObj = fromjson("{a: {$gt: 5}}");
+    auto filterExpr = parseMatchExpression(filterObj);
+    IndexEntry entry(BSON("$**" << 1),
+                     false,  // multikey
+                     false,  // sparse
+                     false,  // unique
+                     IndexEntry::Identifier{"indexName"},
+                     filterExpr.get(),
+                     BSONObj());
+    state.updateDiscriminators({entry});
+
+    auto discriminatorsA = state.buildAllPathsDiscriminators("a");
+    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()));
+
+    // 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()));
+}
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_planner_all_paths_index_test.cpp b/src/mongo/db/query/query_planner_all_paths_index_test.cpp
index c758a0c48a8..7ea870083f8 100644
--- a/src/mongo/db/query/query_planner_all_paths_index_test.cpp
+++ b/src/mongo/db/query/query_planner_all_paths_index_test.cpp
@@ -50,7 +50,8 @@ protected:
 
     void addAllPathsIndex(BSONObj keyPattern,
                           const std::set<std::string>& multikeyPathSet = {},
-                          BSONObj starPathsTempName = BSONObj{}) {
+                          BSONObj starPathsTempName = BSONObj{},
+                          MatchExpression* partialFilterExpr = nullptr) {
         // Convert the set of std::string to a set of FieldRef.
         std::set<FieldRef> multikeyFieldRefs;
         for (auto&& path : multikeyPathSet) {
@@ -69,7 +70,7 @@ protected:
                                             false,  // sparse
                                             false,  // unique
                                             IndexEntry::Identifier{kIndexName},
-                                            nullptr,  // partialFilterExpression
+                                            partialFilterExpr,
                                             std::move(infoObj),
                                             nullptr});  // collator
     }
@@ -585,6 +586,82 @@ TEST_F(QueryPlannerAllPathsTest, InBasicOrEquivalent) {
         "bounds: {'$_path': [['a','a',true,true]], a: [[1,1,true,true],[2,2,true,true]]}}}}}");
 }
 
+TEST_F(QueryPlannerAllPathsTest, PartialIndexCanAnswerPredicateOnFilteredField) {
+    auto filterObj = fromjson("{a: {$gt: 0}}");
+    auto filterExpr = QueryPlannerTest::parseMatchExpression(filterObj);
+    addAllPathsIndex(BSON("$**" << 1), {}, BSONObj{}, filterExpr.get());
+
+    runQuery(fromjson("{a: {$gte: 5}}"));
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: null, node: "
+        "{ixscan: {filter: null, pattern: {'$_path': 1, a: 1},"
+        "bounds: {'$_path': [['a','a',true,true]], a: [[5,Infinity,true,true]]}}}}}");
+
+    runQuery(fromjson("{a: 5}"));
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: null, node: "
+        "{ixscan: {filter: null, pattern: {'$_path': 1, a: 1},"
+        "bounds: {'$_path': [['a','a',true,true]], a: [[5,5,true,true]]}}}}}");
+
+    runQuery(fromjson("{a: {$gte: 1, $lte: 10}}"));
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: null, node: "
+        "{ixscan: {filter: null, pattern: {'$_path': 1, a: 1},"
+        "bounds: {'$_path': [['a','a',true,true]], a: [[1,10,true,true]]}}}}}");
+}
+
+TEST_F(QueryPlannerAllPathsTest, PartialIndexDoesNotAnswerPredicatesExcludedByFilter) {
+    // Must keep 'filterObj' around since match expressions will store pointers into the BSON they
+    // were parsed from.
+    auto filterObj = fromjson("{a: {$gt: 0}}");
+    auto filterExpr = QueryPlannerTest::parseMatchExpression(filterObj);
+    addAllPathsIndex(BSON("$**" << 1), {}, BSONObj{}, filterExpr.get());
+
+    runQuery(fromjson("{a: {$gte: -1}}"));
+    assertHasOnlyCollscan();
+
+    runQuery(fromjson("{a: {$lte: 10}}"));
+    assertHasOnlyCollscan();
+
+    runQuery(fromjson("{a: {$eq: 0}}"));
+    assertHasOnlyCollscan();
+}
+
+TEST_F(QueryPlannerAllPathsTest, PartialIndexCanAnswerPredicateOnUnrelatedField) {
+    auto filterObj = fromjson("{a: {$gt: 0}}");
+    auto filterExpr = QueryPlannerTest::parseMatchExpression(filterObj);
+    addAllPathsIndex(BSON("$**" << 1), {}, BSONObj{}, filterExpr.get());
+
+    // Test when the field query is not included by the partial filter expression.
+    runQuery(fromjson("{b: {$gte: -1}, a: {$gte: 5}}"));
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{fetch: {filter: {a: {$gte: 5}}, node: "
+        "{ixscan: {filter: null, pattern: {'$_path': 1, b: 1},"
+        "bounds: {'$_path': [['b','b',true,true]], b: [[-1,Infinity,true,true]]}}}}}");
+    assertSolutionExists(
+        "{fetch: {filter: {b: {$gte: -1}}, node: "
+        "{ixscan: {filter: null, pattern: {'$_path': 1, a: 1},"
+        "bounds: {'$_path': [['a','a',true,true]], a: [[5,Infinity,true,true]]}}}}}");
+}
+
+TEST_F(QueryPlannerAllPathsTest, PartialIndexWithExistsTrueFilterCanAnswerExistenceQuery) {
+    auto filterObj = fromjson("{x: {$exists: true}}");
+    auto filterExpr = QueryPlannerTest::parseMatchExpression(filterObj);
+    addAllPathsIndex(BSON("$**" << 1), {}, BSONObj{}, filterExpr.get());
+    runQuery(fromjson("{x: {$exists: true}}"));
+
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: null, node: "
+        "{ixscan: {filter: null, pattern: {'$_path': 1, x: 1},"
+        "bounds: {'$_path': [['x','x',true,true],['x.','x/',true,false]], x: "
+        "[['MinKey','MaxKey',true,true]]}}}}}");
+}
+
 //
 // Index intersection tests.
 //
@@ -692,7 +769,6 @@ TEST_F(QueryPlannerAllPathsTest, AllPathsIndexDoesNotSupplyCandidatePlanForTextS
 
 // TODO SERVER-35335: Add testing for Min/Max.
 // TODO SERVER-36517: Add testing for DISTINCT_SCAN.
-// TODO SERVER-35336: Add testing for partialFilterExpression.
 // TODO SERVER-35331: Add testing for hints.
 // TODO SERVER-36145: Add testing for non-blocking sort.
 
diff --git a/src/mongo/db/query/query_planner_test_fixture.cpp b/src/mongo/db/query/query_planner_test_fixture.cpp
index 4d11ba10903..fc6e782895e 100644
--- a/src/mongo/db/query/query_planner_test_fixture.cpp
+++ b/src/mongo/db/query/query_planner_test_fixture.cpp
@@ -462,6 +462,11 @@ void QueryPlannerTest::assertHasOneSolutionOf(const std::vector<std::string>& so
     FAIL(ss);
 }
 
+void QueryPlannerTest::assertHasOnlyCollscan() const {
+    assertNumSolutions(1U);
+    assertSolutionExists("{cscan: {dir: 1}}");
+}
+
 std::unique_ptr<MatchExpression> QueryPlannerTest::parseMatchExpression(
     const BSONObj& obj, const CollatorInterface* collator) {
     boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
diff --git a/src/mongo/db/query/query_planner_test_fixture.h b/src/mongo/db/query/query_planner_test_fixture.h
index 2bdd7d5b168..628d9ee3fca 100644
--- a/src/mongo/db/query/query_planner_test_fixture.h
+++ b/src/mongo/db/query/query_planner_test_fixture.h
@@ -192,6 +192,11 @@ protected:
     void assertHasOneSolutionOf(const std::vector<std::string>& solnStrs) const;
 
     /**
+     * Check that the only solution available is an ascending collection scan.
+     */
+    void assertHasOnlyCollscan() const;
+
+    /**
      * Helper function to parse a MatchExpression.
      */
     static std::unique_ptr<MatchExpression> parseMatchExpression(