From a4fc389f6be6cd97382774ff4873deac79c84601 Mon Sep 17 00:00:00 2001 From: Aaron Date: Fri, 20 Feb 2009 15:19:41 -0500 Subject: [PATCH] Count function using new query optimizer --- db/queryoptimizer.cpp | 68 ++++++++++++++++++++++++++++++--- db/queryoptimizer.h | 4 +- dbtests/queryoptimizertests.cpp | 12 +++++- 3 files changed, 76 insertions(+), 8 deletions(-) diff --git a/db/queryoptimizer.cpp b/db/queryoptimizer.cpp index 65296b54183..de38d9904f2 100644 --- a/db/queryoptimizer.cpp +++ b/db/queryoptimizer.cpp @@ -224,10 +224,14 @@ namespace mongo { if ( !index_ ) return theDataFileMgr.findAll( fbs_.ns() ); else - return auto_ptr< Cursor >( new BtreeCursor( *const_cast< IndexDetails* >( index_ ), startKey_, endKey_, direction_ ) ); + return auto_ptr< Cursor >( new BtreeCursor( *const_cast< IndexDetails* >( index_ ), startKey_, endKey_, direction_ >= 0 ? 1 : -1 ) ); //TODO This constructor should really take a const ref to the index details. } + BSONObj QueryPlan::indexKey() const { + return index_->keyPattern(); + } + QueryPlanSet::QueryPlanSet( const char *ns, const BSONObj &query, const BSONObj &order, const BSONElement *hint ) : fbs_( ns, query ) { NamespaceDetails *d = nsdetails( ns ); @@ -298,7 +302,6 @@ namespace mongo { threads.create_thread( r ); } threads.join_all(); - cout << "really done" << endl; for( int i = 0; i < plans_.nPlans(); ++i ) if ( ops[ i ]->done() ) return ops[ i ]; @@ -308,21 +311,74 @@ namespace mongo { class CountOp : public QueryOp { public: + CountOp( const BSONObj &spec ) : spec_( spec ), count_() {} virtual void run( const QueryPlan &qp, QueryAborter &qa ) { - for( int i = 0; i < 100000; ++i ) + BSONObj query = spec_.getObjectField("query"); + set< string > fields; + spec_.getObjectField("fields").getFieldNames( fields ); + + auto_ptr c = qp.newCursor(); + + // TODO We could check if all fields in the key are in 'fields' + if ( qp.exactKeyMatch() && fields.empty() ) { + /* Here we only look at the btree keys to determine if a match, instead of looking + into the records, which would be much slower. + */ + BtreeCursor *bc = dynamic_cast(c.get()); + if ( c->ok() && !query.woCompare( bc->currKeyNode().key, BSONObj(), false ) ) { + BSONObj firstMatch = bc->currKeyNode().key; + count_++; + while ( c->advance() ) { + qa.mayAbort(); + if ( !firstMatch.woEqual( bc->currKeyNode().key ) ) + break; + count_++; + } + } + return; + } + + auto_ptr matcher(new JSMatcher(query, c->indexKeyPattern())); + while ( c->ok() ) { qa.mayAbort(); - cout << "done" << endl; + BSONObj js = c->current(); + bool deep; + if ( !matcher->matches(js, &deep) ) { + } + else if ( !deep || !c->getsetdup(c->currLoc()) ) { // i.e., check for dups on deep items only + bool match = true; + for( set< string >::iterator i = fields.begin(); i != fields.end(); ++i ) { + if ( js.getFieldDotted( i->c_str() ).eoo() ) { + match = false; + break; + } + } + if ( match ) + ++count_; + } + c->advance(); + } } virtual QueryOp *clone() const { return new CountOp( *this ); } - int count() const { return 1; } + int count() const { return count_; } + private: + BSONObj spec_; + int count_; }; int doCount( const char *ns, const BSONObj &cmd, string &err ) { BSONObj query = cmd.getObjectField("query"); + BSONObj fields = cmd.getObjectField("fields"); + // count of all objects + if ( query.isEmpty() && fields.isEmpty() ) { + NamespaceDetails *d = nsdetails( ns ); + massert( "ns missing", d ); + return d->nrecords; + } QueryPlanSet qps( ns, query, emptyObj ); - auto_ptr< QueryOp > original( new CountOp () ); + auto_ptr< QueryOp > original( new CountOp( cmd ) ); auto_ptr< QueryOp > o = qps.runOp( *original ); return dynamic_cast< CountOp* >( o.get() )->count(); } diff --git a/db/queryoptimizer.h b/db/queryoptimizer.h index cd69c5ab7fe..56348293c7f 100644 --- a/db/queryoptimizer.h +++ b/db/queryoptimizer.h @@ -93,6 +93,8 @@ namespace mongo { BSONObj startKey() const { return startKey_; } BSONObj endKey() const { return endKey_; } auto_ptr< Cursor > newCursor() const; + BSONObj indexKey() const; + const char *ns() const { return fbs_.ns(); } private: const FieldBoundSet &fbs_; const BSONObj &order_; @@ -122,7 +124,7 @@ namespace mongo { class QueryOp { public: - QueryOp() : done_( false ) {} + QueryOp() : done_() {} virtual ~QueryOp() {} virtual void run( const QueryPlan &qp, QueryAborter &qa ) = 0; virtual QueryOp *clone() const = 0; diff --git a/dbtests/queryoptimizertests.cpp b/dbtests/queryoptimizertests.cpp index f0f5622be9e..862d4011018 100644 --- a/dbtests/queryoptimizertests.cpp +++ b/dbtests/queryoptimizertests.cpp @@ -506,7 +506,17 @@ namespace QueryOptimizerTests { Helpers::ensureIndex( ns(), BSON( "a" << 1 ), "a_1" ); Helpers::ensureIndex( ns(), BSON( "b" << 1 ), "b_1" ); string err; - doCount( ns(), BSON( "query" << BSON( "a" << 4 ) ), err ); + ASSERT_EQUALS( 0, doCount( ns(), BSON( "query" << BSON( "a" << 4 ) ), err ) ); + BSONObj one = BSON( "a" << 1 ); + BSONObj four = BSON( "a" << 4 ); + theDataFileMgr.insert( ns(), one ); + ASSERT_EQUALS( 0, doCount( ns(), BSON( "query" << BSON( "a" << 4 ) ), err ) ); + theDataFileMgr.insert( ns(), four ); + ASSERT_EQUALS( 1, doCount( ns(), BSON( "query" << BSON( "a" << 4 ) ), err ) ); + theDataFileMgr.insert( ns(), four ); + ASSERT_EQUALS( 2, doCount( ns(), BSON( "query" << BSON( "a" << 4 ) ), err ) ); + ASSERT_EQUALS( 3, doCount( ns(), BSON( "query" << emptyObj ), err ) ); + ASSERT_EQUALS( 3, doCount( ns(), BSON( "query" << BSON( "a" << GT << 0 ) ), err ) ); } };