SERVER-19507 allow DISTINCT_SCAN over trailing field of an index

Closes #1101

Signed-off-by: David Storch <david.storch@10gen.com>

2 files changed, 77 insertions, 6 deletions

diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 0963df541a9..4742f915d0d 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -1137,6 +1137,10 @@ bool getDistinctNodeIndex(const std::vector<IndexEntry>& indices,
         if (indices[i].multikey && isDottedField) {
             continue;
         }
+        // Skip indices where the first key is not field.
+        if (indices[i].keyPattern.firstElement().fieldNameStringData() != StringData(field)) {
+            continue;
+        }
         int nFields = indices[i].keyPattern.nFields();
         // Pick the index with the lowest number of fields.
         if (nFields < minFields) {
@@ -1352,9 +1356,7 @@ bool turnIxscanIntoDistinctIxscan(QuerySolution* soln, const string& field) {
     distinctNode->direction = indexScanNode->direction;
     distinctNode->bounds = indexScanNode->bounds;
 
-    // Figure out which field we're skipping to the next value of.  TODO: We currently only
-    // try to distinct-hack when there is an index prefixed by the field we're distinct-ing
-    // over.  Consider removing this code if we stick with that policy.
+    // Figure out which field we're skipping to the next value of.
     distinctNode->fieldNo = 0;
     BSONObjIterator it(indexScanNode->index.keyPattern);
     while (it.more()) {
@@ -1431,9 +1433,7 @@ StatusWith<unique_ptr<PlanExecutor>> getExecutorDistinct(OperationContext* txn,
     while (ii.more()) {
         const IndexDescriptor* desc = ii.next();
         IndexCatalogEntry* ice = ii.catalogEntry(desc);
-        // The distinct hack can work if any field is in the index but it's not always clear
-        // if it's a win unless it's the first field.
-        if (desc->keyPattern().firstElement().fieldName() == parsedDistinct->getKey()) {
+        if (desc->keyPattern().hasField(parsedDistinct->getKey())) {
             plannerParams.indices.push_back(IndexEntry(desc->keyPattern(),
                                                        desc->getAccessMethodName(),
                                                        desc->isMultikey(txn),
diff --git a/src/mongo/dbtests/query_stage_distinct.cpp b/src/mongo/dbtests/query_stage_distinct.cpp
index 1cbc7b440c4..1558f62574a 100644
--- a/src/mongo/dbtests/query_stage_distinct.cpp
+++ b/src/mongo/dbtests/query_stage_distinct.cpp
@@ -236,6 +236,76 @@ public:
     }
 };
 
+class QueryStageDistinctCompoundIndex : public DistinctBase {
+public:
+    void run() {
+        // insert documents with a: 1 and b: 1
+        for (size_t i = 0; i < 1000; ++i) {
+            insert(BSON("a" << 1 << "b" << 1));
+        }
+        // insert documents with a: 1 and b: 2
+        for (size_t i = 0; i < 1000; ++i) {
+            insert(BSON("a" << 1 << "b" << 2));
+        }
+        // insert documents with a: 2 and b: 1
+        for (size_t i = 0; i < 1000; ++i) {
+            insert(BSON("a" << 2 << "b" << 1));
+        }
+        // insert documents with a: 2 and b: 3
+        for (size_t i = 0; i < 1000; ++i) {
+            insert(BSON("a" << 2 << "b" << 3));
+        }
+
+        addIndex(BSON("a" << 1 << "b" << 1));
+
+        AutoGetCollectionForRead ctx(&_txn, ns());
+        Collection* coll = ctx.getCollection();
+
+        std::vector<IndexDescriptor*> indices;
+        coll->getIndexCatalog()->findIndexesByKeyPattern(
+            &_txn, BSON("a" << 1 << "b" << 1), false, &indices);
+        ASSERT_EQ(1U, indices.size());
+
+        DistinctParams params;
+        params.descriptor = indices[0];
+        ASSERT_TRUE(params.descriptor);
+
+        params.direction = 1;
+        params.fieldNo = 1;
+        params.bounds.isSimpleRange = false;
+
+        OrderedIntervalList aOil{"a"};
+        aOil.intervals.push_back(IndexBoundsBuilder::allValues());
+        params.bounds.fields.push_back(aOil);
+
+        OrderedIntervalList bOil{"b"};
+        bOil.intervals.push_back(IndexBoundsBuilder::allValues());
+        params.bounds.fields.push_back(bOil);
+
+        WorkingSet ws;
+        DistinctScan distinct(&_txn, params, &ws);
+
+        WorkingSetID wsid;
+        PlanStage::StageState state;
+
+        std::vector<int> seen;
+
+        while (PlanStage::IS_EOF != (state = distinct.work(&wsid))) {
+            ASSERT_NE(PlanStage::FAILURE, state);
+            ASSERT_NE(PlanStage::DEAD, state);
+            if (PlanStage::ADVANCED == state) {
+                seen.push_back(getIntFieldDotted(ws, wsid, "b"));
+            }
+        }
+
+        ASSERT_EQUALS(4U, seen.size());
+        ASSERT_EQUALS(1, seen[0]);
+        ASSERT_EQUALS(2, seen[1]);
+        ASSERT_EQUALS(1, seen[2]);
+        ASSERT_EQUALS(3, seen[3]);
+    }
+};
+
 // XXX: add a test case with bounds where skipping to the next key gets us a result that's not
 // valid w.r.t. our query.
 
@@ -246,6 +316,7 @@ public:
     void setupTests() {
         add<QueryStageDistinctBasic>();
         add<QueryStageDistinctMultiKey>();
+        add<QueryStageDistinctCompoundIndex>();
     }
 };