SERVER-60121 Enable hint to use the clusterKey of a clustered collection

2 files changed, 330 insertions, 12 deletions

diff --git a/jstests/core/clustered_collection_hint.js b/jstests/core/clustered_collection_hint.js
new file mode 100644
index 00000000000..f0ac75f1e31
--- /dev/null
+++ b/jstests/core/clustered_collection_hint.js
@@ -0,0 +1,276 @@
+/**
+ * Tests that a collection with a clustered index can use and interpret a query hint.
+ * @tags: [
+ *   requires_fcv_52,
+ *   // Does not support sharding
+ *   assumes_against_mongod_not_mongos,
+ *   assumes_unsharded_collection,
+ * ]
+ */
+(function() {
+"use strict";
+load("jstests/libs/analyze_plan.js");
+load("jstests/libs/collection_drop_recreate.js");
+
+const clusteredIndexesEnabled = assert
+                                    .commandWorked(db.getMongo().adminCommand(
+                                        {getParameter: 1, featureFlagClusteredIndexes: 1}))
+                                    .featureFlagClusteredIndexes.value;
+
+if (!clusteredIndexesEnabled) {
+    jsTestLog('Skipping test because the clustered indexes feature flag is disabled');
+    return;
+}
+
+const testDB = db.getSiblingDB(jsTestName());
+const collName = "coll";
+const coll = testDB[collName];
+assertDropCollection(testDB, collName);
+
+const validateHint = ({expectedNReturned, cmd, expectedWinningPlanStats = {}}) => {
+    const explain = assert.commandWorked(coll.runCommand({explain: cmd}));
+    assert.eq(explain.executionStats.nReturned, expectedNReturned, tojson(explain));
+
+    const actualWinningPlan = getWinningPlan(explain.queryPlanner);
+    const stageOfInterest = getPlanStage(actualWinningPlan, expectedWinningPlanStats.stage);
+    assert.neq(null, stageOfInterest);
+
+    for (const [key, value] of Object.entries(expectedWinningPlanStats)) {
+        assert(stageOfInterest[key], tojson(explain));
+        assert.eq(stageOfInterest[key], value, tojson(explain));
+    }
+
+    // Explicitly check that the plan is not bounded by default.
+    if (!expectedWinningPlanStats.hasOwnProperty("minRecord")) {
+        assert(!actualWinningPlan["minRecord"], tojson(explain));
+    }
+    if (!expectedWinningPlanStats.hasOwnProperty("maxRecord")) {
+        assert(!actualWinningPlan["maxRecord"], tojson(explain));
+    }
+};
+
+assert.commandWorked(
+    testDB.createCollection(collName, {clusteredIndex: {key: {_id: 1}, unique: true}}));
+
+// Create an index that the query planner would consider preferable to using the cluster key for
+// point predicates on 'a'.
+const idxA = {
+    a: -1
+};
+assert.commandWorked(coll.createIndex(idxA));
+
+const batchSize = 100;
+const bulk = coll.initializeUnorderedBulkOp();
+for (let i = 0; i < batchSize; i++) {
+    bulk.insert({_id: i, a: -i});
+}
+assert.commandWorked(bulk.execute());
+assert.eq(coll.find().itcount(), batchSize);
+
+// Basic find with hints on cluster key.
+validateHint({
+    expectedNReturned: batchSize,
+    cmd: {
+        find: collName,
+        hint: {_id: 1},
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+    }
+});
+validateHint({
+    expectedNReturned: batchSize,
+    cmd: {
+        find: collName,
+        hint: "_id_",
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+    }
+});
+validateHint({
+    expectedNReturned: 1,
+    cmd: {
+        find: collName,
+        filter: {a: -2},
+        hint: {_id: 1},
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+    }
+});
+validateHint({
+    expectedNReturned: 1,
+    cmd: {
+        find: collName,
+        filter: {a: -2},
+        hint: "_id_",
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+    }
+});
+
+// Find with hints on cluster key that generate bounded collection scans.
+const arbitraryDocId = 12;
+validateHint({
+    expectedNReturned: 1,
+    cmd: {
+        find: collName,
+        filter: {_id: arbitraryDocId},
+        hint: {_id: 1},
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+        minRecord: arbitraryDocId,
+        maxRecord: arbitraryDocId
+    }
+});
+validateHint({
+    expectedNReturned: 1,
+    cmd: {
+        find: collName,
+        filter: {_id: arbitraryDocId},
+        hint: "_id_",
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+        minRecord: arbitraryDocId,
+        maxRecord: arbitraryDocId
+    }
+});
+validateHint({
+    expectedNReturned: arbitraryDocId,
+    cmd: {
+        find: collName,
+        filter: {_id: {$lt: arbitraryDocId}},
+        hint: {_id: 1},
+    },
+    expectedWinningPlanStats: {stage: "COLLSCAN", direction: "forward", maxRecord: arbitraryDocId}
+});
+validateHint({
+    expectedNReturned: batchSize - arbitraryDocId,
+    cmd: {
+        find: collName,
+        filter: {_id: {$gte: arbitraryDocId}},
+        hint: {_id: 1},
+    },
+    expectedWinningPlanStats: {stage: "COLLSCAN", direction: "forward", minRecord: arbitraryDocId}
+});
+
+// Find with $natural hints.
+validateHint({
+    expectedNReturned: batchSize,
+    cmd: {
+        find: collName,
+        hint: {$natural: -1},
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "backward",
+    }
+});
+validateHint({
+    expectedNReturned: batchSize,
+    cmd: {
+        find: collName,
+        hint: {$natural: 1},
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+    }
+});
+validateHint({
+    expectedNReturned: 1,
+    cmd: {
+        find: collName,
+        filter: {a: -2},
+        hint: {$natural: -1},
+    },
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "backward",
+    }
+});
+
+// Find on a standard index.
+validateHint({
+    expectedNReturned: batchSize,
+    cmd: {find: collName, hint: idxA},
+    expectedWinningPlanStats: {
+        stage: "IXSCAN",
+        keyPattern: idxA,
+    }
+});
+
+// Update with hint on cluster key.
+validateHint({
+    expectedNReturned: 0,
+    cmd: {update: collName, updates: [{q: {_id: 3}, u: {$inc: {a: -2}}, hint: {_id: 1}}]},
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+    }
+});
+
+// Update with reverse $natural hint.
+validateHint({
+    expectedNReturned: 0,
+    cmd: {update: collName, updates: [{q: {_id: 80}, u: {$inc: {a: 80}}, hint: {$natural: -1}}]},
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "backward",
+    }
+});
+
+// Update with hint on secondary index.
+validateHint({
+    expectedNReturned: 0,
+    cmd: {update: collName, updates: [{q: {a: -2}, u: {$set: {a: 2}}, hint: idxA}]},
+    expectedWinningPlanStats: {
+        stage: "IXSCAN",
+        keyPattern: idxA,
+    }
+});
+
+// Delete with hint on cluster key.
+validateHint({
+    expectedNReturned: 0,
+    cmd: {delete: collName, deletes: [{q: {_id: 2}, limit: 0, hint: {_id: 1}}]},
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "forward",
+    }
+});
+
+// Delete reverse $natural hint.
+validateHint({
+    expectedNReturned: 0,
+    cmd: {delete: collName, deletes: [{q: {_id: 30}, limit: 0, hint: {$natural: -1}}]},
+    expectedWinningPlanStats: {
+        stage: "COLLSCAN",
+        direction: "backward",
+    }
+});
+
+// Delete with hint on standard index.
+validateHint({
+    expectedNReturned: 0,
+    cmd: {delete: collName, deletes: [{q: {a: -5}, limit: 0, hint: idxA}]},
+    expectedWinningPlanStats: {
+        stage: "IXSCAN",
+        keyPattern: idxA,
+    }
+});
+
+// Reverse 'hint' on the cluster key is illegal.
+assert.commandFailedWithCode(testDB.runCommand({find: collName, hint: {_id: -1}}),
+                             ErrorCodes.BadValue);
+})();
diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp
index 600f106f25f..37d89750d2a 100644
--- a/src/mongo/db/query/query_planner.cpp
+++ b/src/mongo/db/query/query_planner.cpp
@@ -129,6 +129,40 @@ Status tagOrChildAccordingToCache(PlanCacheIndexTree* compositeCacheData,
 
     return Status::OK();
 }
+
+/**
+ * Returns whether the hintedIndex matches the cluster key. When hinting by index name,
+ * 'hintObj' takes the shape of {$hint: <indexName>}. When hinting by key pattern,
+ * 'hintObj' represents the actual key pattern (eg: {_id: 1}).
+ */
+bool hintMatchesClusterKey(const boost::optional<ClusteredCollectionInfo>& clusteredInfo,
+                           const BSONObj& hintObj) {
+    if (!clusteredInfo) {
+        // The collection isn't clustered.
+        return false;
+    }
+
+    auto clusteredIndexSpec = clusteredInfo->getIndexSpec();
+
+    BSONElement firstHintElt = hintObj.firstElement();
+    if (firstHintElt.fieldNameStringData() == "$hint"_sd &&
+        firstHintElt.type() == BSONType::String) {
+        // An index name is provided by the hint.
+
+        // The clusteredIndex's name should always be filled in with a default value when not
+        // specified upon creation.
+        tassert(6012100,
+                "clusteredIndex's 'name' field should be filled in by default after creation",
+                clusteredIndexSpec.getName());
+
+        auto hintName = firstHintElt.valueStringData();
+        return hintName == clusteredIndexSpec.getName().get();
+    }
+
+    // An index spec is provided by the hint.
+    return hintObj.woCompare(clusteredIndexSpec.getKey()) == 0;
+}
+
 }  // namespace
 
 using std::numeric_limits;
@@ -671,21 +705,29 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::planForMul
         return {std::move(out)};
     }
 
-    // The hint can be {$natural: +/-1}. If this happens, output a collscan. We expect any $natural
-    // sort to have been normalized to a $natural hint upstream.
     if (!query.getFindCommandRequest().getHint().isEmpty()) {
         const BSONObj& hintObj = query.getFindCommandRequest().getHint();
-        if (hintObj[query_request_helper::kNaturalSortField]) {
-            LOGV2_DEBUG(20969, 5, "Forcing a table scan due to hinted $natural");
-            if (!canTableScan) {
-                return Status(ErrorCodes::NoQueryExecutionPlans,
-                              "hint $natural is not allowed, because 'notablescan' is enabled");
-            }
-            if (!query.getFindCommandRequest().getMin().isEmpty() ||
-                !query.getFindCommandRequest().getMax().isEmpty()) {
-                return Status(ErrorCodes::NoQueryExecutionPlans,
-                              "min and max are incompatible with $natural");
+        const auto naturalHint = hintObj[query_request_helper::kNaturalSortField];
+        if (naturalHint || hintMatchesClusterKey(params.clusteredInfo, hintObj)) {
+            // The hint can be {$natural: +/-1}. If this happens, output a collscan. We expect any
+            // $natural sort to have been normalized to a $natural hint upstream. Additionally, if
+            // the hint matches the collection's cluster key, we also output a collscan utilizing
+            // the cluster key.
+
+            if (naturalHint) {
+                // Perform validation specific to $natural.
+                LOGV2_DEBUG(20969, 5, "Forcing a table scan due to hinted $natural");
+                if (!canTableScan) {
+                    return Status(ErrorCodes::NoQueryExecutionPlans,
+                                  "hint $natural is not allowed, because 'notablescan' is enabled");
+                }
+                if (!query.getFindCommandRequest().getMin().isEmpty() ||
+                    !query.getFindCommandRequest().getMax().isEmpty()) {
+                    return Status(ErrorCodes::NoQueryExecutionPlans,
+                                  "min and max are incompatible with $natural");
+                }
             }
+
             auto soln = buildCollscanSoln(query, isTailable, params);
             if (!soln) {
                 return Status(ErrorCodes::NoQueryExecutionPlans,