SERVER-2001 filter queries on mongod using new key extraction path

Queries on mongod must filter out documents that do not currently belong
to the shard.  The belongsToMe function now uses the new key extraction
path to make this determination. An optimization is also added to the
clientcursor class so that it can make this determination using a
covered index if available.

6 files changed, 59 insertions, 12 deletions

diff --git a/src/mongo/db/clientcursor.cpp b/src/mongo/db/clientcursor.cpp
index 8d002038fd1..6424cd16e0a 100644
--- a/src/mongo/db/clientcursor.cpp
+++ b/src/mongo/db/clientcursor.cpp
@@ -427,6 +427,16 @@ namespace mongo {
         return b.obj();
     }
     
+    BSONObj ClientCursor::extractKey( const KeyPattern& usingKeyPattern ) const {
+        KeyPattern currentIndex( _c->indexKeyPattern() );
+        if ( usingKeyPattern.isCoveredBy( currentIndex ) && ! currentIndex.isSpecial() ){
+            BSONObj currKey = _c->currKey();
+            BSONObj prettyKey = currKey.replaceFieldNames( currentIndex.toBSON() );
+            return usingKeyPattern.extractSingleKey( prettyKey );
+        }
+        return usingKeyPattern.extractSingleKey( _c->current() );
+    }
+
     void ClientCursor::fillQueryResultFromObj( BufBuilder &b, const MatchDetails* details ) const {
         const Projection::KeyOnly *keyFieldsOnly = c()->keyFieldsOnly();
         if ( keyFieldsOnly ) {
diff --git a/src/mongo/db/clientcursor.h b/src/mongo/db/clientcursor.h
index 5498bd84acd..f77012c5fa2 100644
--- a/src/mongo/db/clientcursor.h
+++ b/src/mongo/db/clientcursor.h
@@ -38,6 +38,7 @@
 #include "matcher.h"
 #include "projection.h"
 #include "s/d_chunk_manager.h"
+#include "mongo/db/keypattern.h"
 
 namespace mongo {
 
@@ -290,6 +291,13 @@ namespace mongo {
         */
         BSONObj extractFields(const BSONObj &pattern , bool fillWithNull = false) ;
 
+        /** Extract elements from the object this cursor currently points to, using the expression
+         *  specified in KeyPattern. Will use a covered index if the one in this cursor is usable.
+         *  TODO: there are some cases where a covered index could be used but is not, for instance
+         *  if both this index and the keyPattern are {a : "hashed"}
+         */
+        BSONObj extractKey( const KeyPattern& usingKeyPattern ) const;
+
         void fillQueryResultFromObj( BufBuilder &b, const MatchDetails* details = NULL ) const;
 
         bool currentIsDup() { return _c->getsetdup( _c->currLoc() ); }
diff --git a/src/mongo/db/keypattern.cpp b/src/mongo/db/keypattern.cpp
index 1f3fd6233ec..96229b70fd8 100644
--- a/src/mongo/db/keypattern.cpp
+++ b/src/mongo/db/keypattern.cpp
@@ -47,6 +47,20 @@ namespace mongo {
         return false;
     }
 
+    bool KeyPattern::isCoveredBy( const KeyPattern& other ) const {
+        BSONForEach( e, _pattern ) {
+            BSONElement otherfield = other.getField( e.fieldName() );
+            if ( otherfield.eoo() ){
+                return false;
+            }
+
+            if ( otherfield.numberInt() != 1 && otherfield.numberInt() != -1 && otherfield != e ){
+                return false;
+            }
+        }
+        return true;
+    }
+
     typedef vector<pair<BSONObj,BSONObj> >::const_iterator BoundListIter;
 
     BoundList KeyPattern::keyBounds( const FieldRangeSet& queryConstraints ) const {
diff --git a/src/mongo/db/keypattern.h b/src/mongo/db/keypattern.h
index fa957d0c4a7..b8c96ba88e8 100644
--- a/src/mongo/db/keypattern.h
+++ b/src/mongo/db/keypattern.h
@@ -60,6 +60,12 @@ namespace mongo {
         bool hasField( const char* fieldname ) const { return _pattern.hasField( fieldname ); }
 
         /*
+         * Gets the element of this pattern corresponding to the given fieldname.
+         * Returns eoo if none exists.
+         */
+        BSONElement getField( const char* fieldname ) const { return _pattern[ fieldname ]; }
+
+        /*
          * Returns true if the key described by this KeyPattern is a prefix of
          * the (potentially) compound key described by 'other'
          */
@@ -75,6 +81,13 @@ namespace mongo {
          */
         bool isSpecial() const;
 
+        /**
+         * Returns true if the quantities stored in this KeyPattern can be used to compute all the
+         * quantities in "other". Useful for determining whether an index based on one KeyPattern
+         * can be used as a covered index for a query based on another.
+         */
+        bool isCoveredBy( const KeyPattern& other ) const;
+
         string toString() const{ return toBSON().toString(); }
 
         /* Given a document, extracts the index key corresponding to this KeyPattern
@@ -92,7 +105,6 @@ namespace mongo {
          */
         BSONObj extractSingleKey( const BSONObj& doc ) const;
 
-
         /**@param queryConstraints a FieldRangeSet, usually formed from parsing a query
          * @return an ordered list of bounds generated using this KeyPattern and the
          * constraints from the FieldRangeSet.  This function is used in sharding to
diff --git a/src/mongo/s/d_chunk_manager.cpp b/src/mongo/s/d_chunk_manager.cpp
index b5f890f1559..d1d9286e720 100644
--- a/src/mongo/s/d_chunk_manager.cpp
+++ b/src/mongo/s/d_chunk_manager.cpp
@@ -244,22 +244,24 @@ namespace mongo {
         if ( _rangesMap.size() == 0 )
             return false;
         
-        return _belongsToMe( cc->extractFields( _key , true ) );
+        KeyPattern pat( _key );
+        return _belongsToMe( cc->extractKey( pat ) );
     }
 
-    bool ShardChunkManager::belongsToMe( const BSONObj& obj ) const {
+    bool ShardChunkManager::belongsToMe( const BSONObj& doc ) const {
         if ( _rangesMap.size() == 0 )
             return false;
 
-        return _belongsToMe( obj.extractFields( _key , true ) );
+        KeyPattern pat( _key );
+        return _belongsToMe( pat.extractSingleKey( doc ) );
     }
 
-    bool ShardChunkManager::_belongsToMe( const BSONObj& x ) const {
-        RangeMap::const_iterator it = _rangesMap.upper_bound( x );
+    bool ShardChunkManager::_belongsToMe( const BSONObj& point ) const {
+        RangeMap::const_iterator it = _rangesMap.upper_bound( point );
         if ( it != _rangesMap.begin() )
             it--;
 
-        bool good = contains( it->first , it->second , x );
+        bool good = contains( it->first , it->second , point );
 
 #if 0
         if ( ! good ) {
diff --git a/src/mongo/s/d_chunk_manager.h b/src/mongo/s/d_chunk_manager.h
index 781cfb88ff8..b93772ab43c 100644
--- a/src/mongo/s/d_chunk_manager.h
+++ b/src/mongo/s/d_chunk_manager.h
@@ -105,12 +105,13 @@ namespace mongo {
          * @param obj document containing sharding keys (and, optionally, other attributes)
          * @return true if shards hold the object
          */
-        bool belongsToMe( const BSONObj& obj ) const;
+        bool belongsToMe( const BSONObj& doc ) const;
 
         /**
-         * Checks whether a document belongs to this shard.
+         * Checks whether the document currently pointed to by this cursor belongs to this shard.
+         * This version of the function will use a covered index if there is one in the cursor.
          *
-         * @param obj document containing sharding keys (and, optionally, other attributes)
+         * @param cc cursor pointing to an object
          * @return true if shards hold the object
          */
         bool belongsToMe( ClientCursor* cc ) const;
@@ -137,9 +138,9 @@ namespace mongo {
         void _init( const string& configServer , const string& ns , const string& shardName, ShardChunkManagerPtr oldManager = ShardChunkManagerPtr() );
 
         /**
-         * @same as belongsToMe to but key has to be the shard key
+         * @same as belongsToMe but point is the extracted shard key
          */
-        bool _belongsToMe( const BSONObj& key ) const;
+        bool _belongsToMe( const BSONObj& point ) const;
 
         ShardChunkVersion _collVersion;
         // highest ShardChunkVersion for which this ShardChunkManager's information is accurate