SERVER-2001 KeyPattern class; utilities for more general index & shard key specs

The KeyPattern class is an abstraction for defining more general
expression-based keys (both index keys and shard keys). This class
provide some utility functions for extracting keys based on an
expression, and computing range bounds based on an expression.
This patch lays the groundwork and begins to make use of KeyPatterns.
The idea is that to implement more general key expressions, we will only
need to enhance the functions in this class.

6 files changed, 240 insertions, 103 deletions

diff --git a/src/mongo/db/hashindex.cpp b/src/mongo/db/hashindex.cpp
index 034306bf7fd..c33c9fd4515 100644
--- a/src/mongo/db/hashindex.cpp
+++ b/src/mongo/db/hashindex.cpp
@@ -26,13 +26,11 @@ namespace mongo {
     const string HashedIndexType::HASHED_INDEX_TYPE_IDENTIFIER = "hashed";
 
     HashedIndexType::HashedIndexType( const IndexPlugin* plugin , const IndexSpec* spec ) :
-            IndexType( plugin , spec ) {
-
-        _keyPattern = spec->keyPattern;
+            IndexType( plugin , spec ) , _keyPattern( spec->keyPattern ) {
 
         //change these if single-field limitation lifted later
         uassert( 16241 , "Currently only single field hashed index supported." ,
-                _keyPattern.nFields() == 1 );
+                _keyPattern.toBSON().nFields() == 1 );
         uassert( 16242 , "Currently hashed indexes cannot guarantee uniqueness. Use a regular index." ,
                 ! (spec->info).getField("unique").booleanSafe() );
 
@@ -52,7 +50,7 @@ namespace mongo {
         _hashVersion = (spec->info).getField("hashVersion").numberInt();
 
         //Get the hashfield name
-        BSONElement firstElt = _keyPattern.firstElement();
+        BSONElement firstElt = spec->keyPattern.firstElement();
         massert( 16243 , "error: no hashed index field" ,
                 firstElt.str().compare( HASHED_INDEX_TYPE_IDENTIFIER ) == 0 );
         _hashedField = firstElt.fieldName();
@@ -75,13 +73,12 @@ namespace mongo {
         uassert( 16244 , "Error: hashed indexes do not currently support array values" , fieldVal.type() != Array );
 
         if ( ! fieldVal.eoo() ) {
-            BSONObj key = BSON( "" << makeSingleKey( fieldVal , _seed , _hashVersion  ) );
+            BSONObj key = _keyPattern.extractSingleKey( obj ) ;
             keys.insert( key );
         }
         else if (! _isSparse ) {
             BSONObj nullobj = BSON( _hashedField << BSONNULL );
-            BSONElement nullElt = nullobj.firstElement();
-            BSONObj key = BSON( "" << makeSingleKey( nullElt , _seed , _hashVersion  ) );
+            BSONObj key = _keyPattern.extractSingleKey( nullobj );
             keys.insert( key );
         }
     }
diff --git a/src/mongo/db/hashindex.h b/src/mongo/db/hashindex.h
index a360ea8669a..3a96f82178c 100644
--- a/src/mongo/db/hashindex.h
+++ b/src/mongo/db/hashindex.h
@@ -19,6 +19,7 @@
 #include "mongo/db/btree.h"
 #include "mongo/db/hasher.h"
 #include "mongo/db/index.h"
+#include "mongo/db/keypattern.h"
 #include "mongo/db/matcher.h"
 #include "mongo/db/namespace-inl.h"
 #include "mongo/db/pdfile.h"
@@ -110,7 +111,7 @@ namespace mongo {
 
     private:
         string _hashedField;
-        BSONObj _keyPattern;
+        KeyPattern _keyPattern;
         HashSeed _seed; //defaults to zero if not in the IndexSpec
         HashVersion _hashVersion; //defaults to zero if not in the IndexSpec
         bool _isSparse;
diff --git a/src/mongo/db/keypattern.cpp b/src/mongo/db/keypattern.cpp
new file mode 100644
index 00000000000..93b4ab050fc
--- /dev/null
+++ b/src/mongo/db/keypattern.cpp
@@ -0,0 +1,117 @@
+// @file keypattern.cpp
+
+/**
+*    Copyright (C) 2012 10gen Inc.
+*
+*    This program is free software: you can redistribute it and/or  modify
+*    it under the terms of the GNU Affero General Public License, version 3,
+*    as published by the Free Software Foundation.
+*
+*    This program is distributed in the hope that it will be useful,
+*    but WITHOUT ANY WARRANTY; without even the implied warranty of
+*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+*    GNU Affero General Public License for more details.
+*
+*    You should have received a copy of the GNU Affero General Public License
+*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+#include "mongo/db/keypattern.h"
+#include "mongo/util/mongoutils/str.h"
+#include "mongo/db/hasher.h"
+#include "mongo/db/queryutil.h"
+
+namespace mongo {
+
+    BSONObj KeyPattern::extractSingleKey(const BSONObj& doc ) const {
+        if ( _pattern.isEmpty() )
+            return BSONObj();
+
+        if ( mongoutils::str::equals( _pattern.firstElement().valuestrsafe() , "hashed" ) ){
+            BSONElement fieldVal = doc.getFieldDotted( _pattern.firstElementFieldName() );
+            return BSON( _pattern.firstElementFieldName() <<
+                         BSONElementHasher::hash64( fieldVal ,
+                                                    BSONElementHasher::DEFAULT_HASH_SEED ) );
+        }
+
+        return doc.extractFields( _pattern );
+    }
+
+    bool KeyPattern::isSpecial() const {
+        BSONForEach(e, _pattern) {
+            int fieldVal = e.numberInt();
+            if ( fieldVal != 1 && fieldVal != -1 ){
+                return true;
+            }
+        }
+        return false;
+    }
+
+    BoundList KeyPattern::keyBounds( const FieldRangeSet& queryConstraints ) const {
+        typedef vector< pair< shared_ptr<BSONObjBuilder> ,
+                              shared_ptr<BSONObjBuilder> > > BoundBuilders;
+        BoundBuilders builders;
+        builders.push_back( make_pair( shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ),
+                                       shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ) ) );
+        BSONObjIterator i( _pattern );
+        // until equalityOnly is false, we are just dealing with equality (no range or $in querys).
+        bool equalityOnly = true;
+        while( i.more() ) {
+            BSONElement e = i.next();
+            const FieldRange &fr = queryConstraints.range( e.fieldName() );
+            int number = (int) e.number(); // returns 0.0 if not numeric
+            bool forward = ( number >= 0 );
+            if ( equalityOnly ) {
+                if ( fr.equality() ) {
+                    for( BoundBuilders::const_iterator j = builders.begin(); j != builders.end(); ++j ) {
+                        j->first->appendAs( fr.min(), "" );
+                        j->second->appendAs( fr.min(), "" );
+                    }
+                }
+                else {
+                    equalityOnly = false;
+
+                    BoundBuilders newBuilders;
+                    const vector<FieldInterval> &intervals = fr.intervals();
+                    for( BoundBuilders::const_iterator i = builders.begin(); i != builders.end(); ++i ) {
+                        BSONObj first = i->first->obj();
+                        BSONObj second = i->second->obj();
+
+                        if ( forward ) {
+                            for( vector<FieldInterval>::const_iterator j = intervals.begin(); j != intervals.end(); ++j ) {
+                                uassert( 16449, "combinatorial limit of $in partitioning of result set exceeded", newBuilders.size() < MAX_IN_COMBINATIONS );
+                                newBuilders.push_back( make_pair( shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ), shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ) ) );
+                                newBuilders.back().first->appendElements( first );
+                                newBuilders.back().second->appendElements( second );
+                                newBuilders.back().first->appendAs( j->_lower._bound, "" );
+                                newBuilders.back().second->appendAs( j->_upper._bound, "" );
+                            }
+                        }
+                        else {
+                            for( vector<FieldInterval>::const_reverse_iterator j = intervals.rbegin(); j != intervals.rend(); ++j ) {
+                                uassert( 16450, "combinatorial limit of $in partitioning of result set exceeded", newBuilders.size() < MAX_IN_COMBINATIONS );
+                                newBuilders.push_back( make_pair( shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ), shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ) ) );
+                                newBuilders.back().first->appendElements( first );
+                                newBuilders.back().second->appendElements( second );
+                                newBuilders.back().first->appendAs( j->_upper._bound, "" );
+                                newBuilders.back().second->appendAs( j->_lower._bound, "" );
+                            }
+                        }
+                    }
+                    builders = newBuilders;
+                }
+            }
+            else {
+                for( BoundBuilders::const_iterator j = builders.begin(); j != builders.end(); ++j ) {
+                    j->first->appendAs( forward ? fr.min() : fr.max(), "" );
+                    j->second->appendAs( forward ? fr.max() : fr.min(), "" );
+                }
+            }
+        }
+        BoundList ret;
+        for( BoundBuilders::const_iterator i = builders.begin(); i != builders.end(); ++i )
+            ret.push_back( make_pair( i->first->obj(), i->second->obj() ) );
+        return ret;
+    }
+
+} // namespace mongo
diff --git a/src/mongo/db/keypattern.h b/src/mongo/db/keypattern.h
new file mode 100644
index 00000000000..afb7a04af76
--- /dev/null
+++ b/src/mongo/db/keypattern.h
@@ -0,0 +1,108 @@
+// @file keypattern.h - Utilities for manipulating index/shard key patterns.
+
+/**
+*    Copyright (C) 2012 10gen Inc.
+*
+*    This program is free software: you can redistribute it and/or  modify
+*    it under the terms of the GNU Affero General Public License, version 3,
+*    as published by the Free Software Foundation.
+*
+*    This program is distributed in the hope that it will be useful,
+*    but WITHOUT ANY WARRANTY; without even the implied warranty of
+*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+*    GNU Affero General Public License for more details.
+*
+*    You should have received a copy of the GNU Affero General Public License
+*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+#pragma once
+
+#include "mongo/db/jsobj.h"
+
+namespace mongo {
+
+    class FieldRangeSet;
+
+    /**
+     * A BoundList contains intervals specified by inclusive start
+     * and end bounds.  The intervals should be nonoverlapping and occur in
+     * the specified direction of traversal.  For example, given a simple index {i:1}
+     * and direction +1, one valid BoundList is: (1, 2); (4, 6).  The same BoundList
+     * would be valid for index {i:-1} with direction -1.
+     */
+    typedef vector<pair<BSONObj,BSONObj> > BoundList;
+
+    /** A KeyPattern is an expression describing a transformation of a document into a
+     *  document key.  Document keys are used to store documents in indices and to target
+     *  sharded queries.
+     *
+     *  Examples:
+     *    { a : 1 }
+     *    { a : 1 , b  : -1 }
+     *    { a : "hashed" }
+     */
+    class KeyPattern {
+    public:
+        KeyPattern( const BSONObj& pattern ): _pattern( pattern ) {}
+
+        /*
+         *  Returns a BSON representation of this KeyPattern.
+         */
+        BSONObj toBSON() const { return _pattern; }
+
+        /*
+         * Returns true if the given fieldname is the name of one element of the (potentially)
+         * compound key described by this KeyPattern.
+         */
+        bool hasField( const char* fieldname ) const { return _pattern.hasField( fieldname ); }
+
+        /*
+         * Returns true if the key described by this KeyPattern is a prefix of
+         * the (potentially) compound key described by 'other'
+         */
+        bool isPrefixOf( const KeyPattern& other ) const {
+            return _pattern.isPrefixOf( other.toBSON() );
+        }
+
+        /**
+         * Returns true if this KeyPattern contains any computed values, (e.g. {a : "hashed"}),
+         * and false if this KeyPattern consists of only ascending/descending fields
+         * (e.g. {a : 1, b : -1}). With our current index expression language, "special" patterns
+         * are any patterns that are not a simple list of field names and 1/-1 values.
+         */
+        bool isSpecial() const;
+
+        string toString() const{ return toBSON().toString(); }
+
+        /* Given a document, extracts the index key corresponding to this KeyPattern
+         * Warning: assumes that there is a *single* key to be extracted!
+         *
+         * Examples:
+         *  If 'this' KeyPattern is { a  : 1 }
+         *   { a: "hi" , b : 4} --> returns { a : "hi" }
+         *   { c : 4 , a : 2 } -->  returns { a : 2 }
+         *   { b : 2 }  (bad input, don't call with this)
+         *   { a : [1,2] }  (bad input, don't call with this)
+         *
+         *  If 'this' KeyPattern is { a  : "hashed" }
+         *   { a: 1 } --> returns { a : NumberLong("5902408780260971510")  }
+         */
+        BSONObj extractSingleKey( const BSONObj& doc ) const;
+
+        /**@param fromQuery a FieldRangeSet formed from parsing a query
+         * @return an ordered list of bounds generated using this KeyPattern
+         * and the constraints from the FieldRangeSet
+         *
+         * The value of frsp->matchPossibleForSingleKeyFRS(fromQuery) should be true,
+         * otherwise this function could throw.
+         *
+         */
+        BoundList keyBounds( const FieldRangeSet& queryConstraints ) const;
+
+    private:
+        BSONObj _pattern;
+    };
+
+
+} // namespace mongo
diff --git a/src/mongo/db/queryutil.cpp b/src/mongo/db/queryutil.cpp
index b926160664e..bdfc561ac90 100644
--- a/src/mongo/db/queryutil.cpp
+++ b/src/mongo/db/queryutil.cpp
@@ -25,8 +25,6 @@
 
 namespace mongo {
 
-    static const unsigned maxCombinations = 4000000;
-
     ParsedQuery::ParsedQuery( QueryMessage& qm )
     : _ns( qm.ns ) , _ntoskip( qm.ntoskip ) , _ntoreturn( qm.ntoreturn ) , _options( qm.queryOptions ) {
         init( qm.query );
@@ -1163,7 +1161,7 @@ namespace mongo {
             verify( !_ranges.back().empty() );
         }
         uassert( 13385, "combinatorial limit of $in partitioning of result set exceeded",
-                size() < maxCombinations );
+                size() < MAX_IN_COMBINATIONS );
     }    
 
     BSONObj FieldRangeVector::startKey() const {
@@ -1260,71 +1258,6 @@ namespace mongo {
         return QueryPattern( *this, sort );
     }
 
-    // TODO get rid of this
-    BoundList FieldRangeSet::indexBounds( const BSONObj &keyPattern, int direction ) const {
-        typedef vector<pair<shared_ptr<BSONObjBuilder>, shared_ptr<BSONObjBuilder> > > BoundBuilders;
-        BoundBuilders builders;
-        builders.push_back( make_pair( shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ), shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ) ) );
-        BSONObjIterator i( keyPattern );
-        bool equalityOnly = true; // until equalityOnly is false, we are just dealing with equality (no range or $in querys).
-        while( i.more() ) {
-            BSONElement e = i.next();
-            const FieldRange &fr = range( e.fieldName() );
-            int number = (int) e.number(); // returns 0.0 if not numeric
-            bool forward = ( ( number >= 0 ? 1 : -1 ) * ( direction >= 0 ? 1 : -1 ) > 0 );
-            if ( equalityOnly ) {
-                if ( fr.equality() ) {
-                    for( BoundBuilders::const_iterator j = builders.begin(); j != builders.end(); ++j ) {
-                        j->first->appendAs( fr.min(), "" );
-                        j->second->appendAs( fr.min(), "" );
-                    }
-                }
-                else {
-                    equalityOnly = false;
-
-                    BoundBuilders newBuilders;
-                    const vector<FieldInterval> &intervals = fr.intervals();
-                    for( BoundBuilders::const_iterator i = builders.begin(); i != builders.end(); ++i ) {
-                        BSONObj first = i->first->obj();
-                        BSONObj second = i->second->obj();
-
-                        if ( forward ) {
-                            for( vector<FieldInterval>::const_iterator j = intervals.begin(); j != intervals.end(); ++j ) {
-                                uassert( 13303, "combinatorial limit of $in partitioning of result set exceeded", newBuilders.size() < maxCombinations );
-                                newBuilders.push_back( make_pair( shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ), shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ) ) );
-                                newBuilders.back().first->appendElements( first );
-                                newBuilders.back().second->appendElements( second );
-                                newBuilders.back().first->appendAs( j->_lower._bound, "" );
-                                newBuilders.back().second->appendAs( j->_upper._bound, "" );
-                            }
-                        }
-                        else {
-                            for( vector<FieldInterval>::const_reverse_iterator j = intervals.rbegin(); j != intervals.rend(); ++j ) {
-                                uassert( 13304, "combinatorial limit of $in partitioning of result set exceeded", newBuilders.size() < maxCombinations );
-                                newBuilders.push_back( make_pair( shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ), shared_ptr<BSONObjBuilder>( new BSONObjBuilder() ) ) );
-                                newBuilders.back().first->appendElements( first );
-                                newBuilders.back().second->appendElements( second );
-                                newBuilders.back().first->appendAs( j->_upper._bound, "" );
-                                newBuilders.back().second->appendAs( j->_lower._bound, "" );
-                            }
-                        }
-                    }
-                    builders = newBuilders;
-                }
-            }
-            else {
-                for( BoundBuilders::const_iterator j = builders.begin(); j != builders.end(); ++j ) {
-                    j->first->appendAs( forward ? fr.min() : fr.max(), "" );
-                    j->second->appendAs( forward ? fr.max() : fr.min(), "" );
-                }
-            }
-        }
-        BoundList ret;
-        for( BoundBuilders::const_iterator i = builders.begin(); i != builders.end(); ++i )
-            ret.push_back( make_pair( i->first->obj(), i->second->obj() ) );
-        return ret;
-    }
-
     int FieldRangeSet::numNonUniversalRanges() const {
         int count = 0;
         for( map<string,FieldRange>::const_iterator i = _ranges.begin(); i != _ranges.end(); ++i ) {
diff --git a/src/mongo/db/queryutil.h b/src/mongo/db/queryutil.h
index bdea78b7361..32bf3541e88 100644
--- a/src/mongo/db/queryutil.h
+++ b/src/mongo/db/queryutil.h
@@ -25,6 +25,9 @@ namespace mongo {
     
     extern const int MaxBytesToReturnToClientAtOnce;
     
+    //maximum number of intervals produced by $in queries.
+    static const unsigned MAX_IN_COMBINATIONS = 4000000;
+
     /* This is for languages whose "objects" are not well ordered (JSON is well ordered).
      [ { a : ... } , { b : ... } ] -> { a : ..., b : ... }
      */
@@ -381,15 +384,6 @@ namespace mongo {
                                        // _elemMatchContext for the FieldRange on 'a.b' is the query
                                        // element having field name '$elemMatch'.
     };
-    
-    /**
-     * A BoundList contains intervals specified by inclusive start
-     * and end bounds.  The intervals should be nonoverlapping and occur in
-     * the specified direction of traversal.  For example, given a simple index {i:1}
-     * and direction +1, one valid BoundList is: (1, 2); (4, 6).  The same BoundList
-     * would be valid for index {i:-1} with direction -1.
-     */
-    typedef vector<pair<BSONObj,BSONObj> > BoundList;
 
     class QueryPattern;
     
@@ -477,18 +471,6 @@ namespace mongo {
         const FieldRangeSet &operator-=( const FieldRangeSet &other );
         /** @return intersection of 'this' with 'other'. */
         const FieldRangeSet &operator&=( const FieldRangeSet &other );
-        
-        /**
-         * @return an ordered list of bounds generated using an index key pattern
-         * and traversal direction.
-         *
-         * The value of matchPossible() should be true, otherwise this function
-         * may @throw.
-         *
-         * NOTE This function is deprecated in the query optimizer and only
-         * currently used by sharding code.
-         */
-        BoundList indexBounds( const BSONObj &keyPattern, int direction ) const;
 
         /**
          * @return - A new FieldRangeSet based on this FieldRangeSet, but with only
@@ -597,17 +579,16 @@ namespace mongo {
          * already been scanned.
          */
         FieldRangeSetPair &operator-=( const FieldRangeSet &scanned );
-
-        BoundList shardKeyIndexBounds( const BSONObj &keyPattern ) const {
-            return _singleKey.indexBounds( keyPattern, 1 );
-        }
         
-        bool matchPossibleForShardKey( const BSONObj &keyPattern ) const {
+        bool matchPossibleForSingleKeyFRS( const BSONObj &keyPattern ) const {
             return _singleKey.matchPossibleForIndex( keyPattern );
         }
         
         BSONObj originalQuery() const { return _singleKey.originalQuery(); }
 
+        const FieldRangeSet getSingleKeyFRS() const { return _singleKey; }
+        const FieldRangeSet getMultiKeyFRS() const { return _singleKey; }
+
         string toString() const;
     private:
         FieldRangeSetPair( const FieldRangeSet &singleKey, const FieldRangeSet &multiKey )