diff options
Diffstat (limited to 'src/mongo/dbtests/query_stage_near.cpp')
| -rw-r--r-- | src/mongo/dbtests/query_stage_near.cpp | 378 |
1 files changed, 180 insertions, 198 deletions
diff --git a/src/mongo/dbtests/query_stage_near.cpp b/src/mongo/dbtests/query_stage_near.cpp index 2ce80d94a7b..935aa2af34d 100644 --- a/src/mongo/dbtests/query_stage_near.cpp +++ b/src/mongo/dbtests/query_stage_near.cpp @@ -38,247 +38,229 @@ namespace { - using namespace mongo; - using std::shared_ptr; - using std::vector; - - /** - * Stage which takes in an array of BSONObjs and returns them. - * If the BSONObj is in the form of a Status, returns the Status as a FAILURE. - */ - class MockStage : public PlanStage { - public: - - MockStage(const vector<BSONObj>& data, WorkingSet* workingSet) : - _data(data), _pos(0), _workingSet(workingSet), _stats("MOCK_STAGE") { - } - - virtual ~MockStage() { - } - - virtual StageState work(WorkingSetID* out) { - ++_stats.works; - - if (isEOF()) - return PlanStage::IS_EOF; - - BSONObj next = _data[_pos++]; - - if (WorkingSetCommon::isValidStatusMemberObject(next)) { - Status status = WorkingSetCommon::getMemberObjectStatus(next); - *out = WorkingSetCommon::allocateStatusMember(_workingSet, status); - return PlanStage::FAILURE; - } +using namespace mongo; +using std::shared_ptr; +using std::vector; - *out = _workingSet->allocate(); - WorkingSetMember* member = _workingSet->get(*out); - member->state = WorkingSetMember::OWNED_OBJ; - member->obj = Snapshotted<BSONObj>(SnapshotId(), next); - - return PlanStage::ADVANCED; - } - - virtual bool isEOF() { - return _pos == static_cast<int>(_data.size()); - } - - virtual void saveState() { - } +/** + * Stage which takes in an array of BSONObjs and returns them. + * If the BSONObj is in the form of a Status, returns the Status as a FAILURE. + */ +class MockStage : public PlanStage { +public: + MockStage(const vector<BSONObj>& data, WorkingSet* workingSet) + : _data(data), _pos(0), _workingSet(workingSet), _stats("MOCK_STAGE") {} - virtual void restoreState(OperationContext* opCtx) { - } + virtual ~MockStage() {} - virtual void invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type) { - } - virtual vector<PlanStage*> getChildren() const { - return vector<PlanStage*>(); - } + virtual StageState work(WorkingSetID* out) { + ++_stats.works; - virtual StageType stageType() const { - return STAGE_UNKNOWN; - } + if (isEOF()) + return PlanStage::IS_EOF; - virtual PlanStageStats* getStats() { - return new PlanStageStats(_stats, STAGE_UNKNOWN); - } + BSONObj next = _data[_pos++]; - virtual const CommonStats* getCommonStats() const { - return &_stats; + if (WorkingSetCommon::isValidStatusMemberObject(next)) { + Status status = WorkingSetCommon::getMemberObjectStatus(next); + *out = WorkingSetCommon::allocateStatusMember(_workingSet, status); + return PlanStage::FAILURE; } - virtual const SpecificStats* getSpecificStats() const { - return NULL; - } + *out = _workingSet->allocate(); + WorkingSetMember* member = _workingSet->get(*out); + member->state = WorkingSetMember::OWNED_OBJ; + member->obj = Snapshotted<BSONObj>(SnapshotId(), next); - private: + return PlanStage::ADVANCED; + } - vector<BSONObj> _data; - int _pos; + virtual bool isEOF() { + return _pos == static_cast<int>(_data.size()); + } - // Not owned here - WorkingSet* const _workingSet; + virtual void saveState() {} - CommonStats _stats; - }; + virtual void restoreState(OperationContext* opCtx) {} - /** - * Stage which implements a basic distance search, and interprets the "distance" field of - * fetched documents as the distance. - */ - class MockNearStage : public NearStage { - public: + virtual void invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type) {} + virtual vector<PlanStage*> getChildren() const { + return vector<PlanStage*>(); + } - struct MockInterval { + virtual StageType stageType() const { + return STAGE_UNKNOWN; + } - MockInterval(const vector<BSONObj>& data, double min, double max) : - data(data), min(min), max(max) { - } + virtual PlanStageStats* getStats() { + return new PlanStageStats(_stats, STAGE_UNKNOWN); + } - vector<BSONObj> data; - double min; - double max; - }; + virtual const CommonStats* getCommonStats() const { + return &_stats; + } - MockNearStage(WorkingSet* workingSet) : - NearStage(NULL, workingSet, NULL, - new PlanStageStats(CommonStats("MOCK_DISTANCE_SEARCH_STAGE"), STAGE_UNKNOWN)), - _pos(0) { - } + virtual const SpecificStats* getSpecificStats() const { + return NULL; + } - virtual ~MockNearStage() { - } +private: + vector<BSONObj> _data; + int _pos; - void addInterval(vector<BSONObj> data, double min, double max) { - _intervals.mutableVector().push_back(new MockInterval(data, min, max)); - } + // Not owned here + WorkingSet* const _workingSet; - virtual StatusWith<CoveredInterval*> nextInterval(OperationContext* txn, - WorkingSet* workingSet, - Collection* collection) { + CommonStats _stats; +}; - if (_pos == static_cast<int>(_intervals.size())) - return StatusWith<CoveredInterval*>(NULL); +/** + * Stage which implements a basic distance search, and interprets the "distance" field of + * fetched documents as the distance. + */ +class MockNearStage : public NearStage { +public: + struct MockInterval { + MockInterval(const vector<BSONObj>& data, double min, double max) + : data(data), min(min), max(max) {} + + vector<BSONObj> data; + double min; + double max; + }; - const MockInterval& interval = *_intervals.vector()[_pos++]; + MockNearStage(WorkingSet* workingSet) + : NearStage(NULL, + workingSet, + NULL, + new PlanStageStats(CommonStats("MOCK_DISTANCE_SEARCH_STAGE"), STAGE_UNKNOWN)), + _pos(0) {} - bool lastInterval = _pos == static_cast<int>(_intervals.vector().size()); - return StatusWith<CoveredInterval*>(new CoveredInterval(new MockStage(interval.data, - workingSet), - true, - interval.min, - interval.max, - lastInterval)); - } + virtual ~MockNearStage() {} - virtual StatusWith<double> computeDistance(WorkingSetMember* member) { - ASSERT(member->hasObj()); - return StatusWith<double>(member->obj.value()["distance"].numberDouble()); - } + void addInterval(vector<BSONObj> data, double min, double max) { + _intervals.mutableVector().push_back(new MockInterval(data, min, max)); + } - virtual StageState initialize(OperationContext* txn, - WorkingSet* workingSet, - Collection* collection, - WorkingSetID* out) { - return IS_EOF; - } + virtual StatusWith<CoveredInterval*> nextInterval(OperationContext* txn, + WorkingSet* workingSet, + Collection* collection) { + if (_pos == static_cast<int>(_intervals.size())) + return StatusWith<CoveredInterval*>(NULL); + + const MockInterval& interval = *_intervals.vector()[_pos++]; + + bool lastInterval = _pos == static_cast<int>(_intervals.vector().size()); + return StatusWith<CoveredInterval*>( + new CoveredInterval(new MockStage(interval.data, workingSet), + true, + interval.min, + interval.max, + lastInterval)); + } - virtual void finishSaveState() { } + virtual StatusWith<double> computeDistance(WorkingSetMember* member) { + ASSERT(member->hasObj()); + return StatusWith<double>(member->obj.value()["distance"].numberDouble()); + } - virtual void finishRestoreState(OperationContext* txn) { } + virtual StageState initialize(OperationContext* txn, + WorkingSet* workingSet, + Collection* collection, + WorkingSetID* out) { + return IS_EOF; + } - virtual void finishInvalidate(OperationContext* txn, - const RecordId& dl, - InvalidationType type) { } + virtual void finishSaveState() {} - private: + virtual void finishRestoreState(OperationContext* txn) {} - OwnedPointerVector<MockInterval> _intervals; - int _pos; - }; + virtual void finishInvalidate(OperationContext* txn, + const RecordId& dl, + InvalidationType type) {} - static vector<BSONObj> advanceStage(PlanStage* stage, WorkingSet* workingSet) { +private: + OwnedPointerVector<MockInterval> _intervals; + int _pos; +}; - vector<BSONObj> results; +static vector<BSONObj> advanceStage(PlanStage* stage, WorkingSet* workingSet) { + vector<BSONObj> results; - WorkingSetID nextMemberID; - PlanStage::StageState state = PlanStage::NEED_TIME; + WorkingSetID nextMemberID; + PlanStage::StageState state = PlanStage::NEED_TIME; - while (PlanStage::NEED_TIME == state) { - while (PlanStage::ADVANCED == (state = stage->work(&nextMemberID))) { - results.push_back(workingSet->get(nextMemberID)->obj.value()); - } + while (PlanStage::NEED_TIME == state) { + while (PlanStage::ADVANCED == (state = stage->work(&nextMemberID))) { + results.push_back(workingSet->get(nextMemberID)->obj.value()); } - - return results; } - static void assertAscendingAndValid(const vector<BSONObj>& results) { - double lastDistance = -1.0; - for (vector<BSONObj>::const_iterator it = results.begin(); it != results.end(); ++it) { - double distance = (*it)["distance"].numberDouble(); - bool shouldInclude = (*it)["$included"].eoo() || (*it)["$included"].trueValue(); - ASSERT(shouldInclude); - ASSERT_GREATER_THAN_OR_EQUALS(distance, lastDistance); - lastDistance = distance; - } - } + return results; +} - TEST(query_stage_near, Basic) { - - vector<BSONObj> mockData; - WorkingSet workingSet; - - MockNearStage nearStage(&workingSet); - - // First set of results - mockData.clear(); - mockData.push_back(BSON("distance" << 0.5)); - mockData.push_back(BSON("distance" << 2.0 << "$included" << false)); // Not included - mockData.push_back(BSON("distance" << 0.0)); - nearStage.addInterval(mockData, 0.0, 1.0); - - // Second set of results - mockData.clear(); - mockData.push_back(BSON("distance" << 1.5)); - mockData.push_back(BSON("distance" << 2.0 << "$included" << false)); // Not included - mockData.push_back(BSON("distance" << 1.0)); - nearStage.addInterval(mockData, 1.0, 2.0); - - // Last set of results - mockData.clear(); - mockData.push_back(BSON("distance" << 2.5)); - mockData.push_back(BSON("distance" << 3.0)); // Included - mockData.push_back(BSON("distance" << 2.0)); - nearStage.addInterval(mockData, 2.0, 3.0); - - vector<BSONObj> results = advanceStage(&nearStage, &workingSet); - ASSERT_EQUALS(results.size(), 7u); - assertAscendingAndValid(results); +static void assertAscendingAndValid(const vector<BSONObj>& results) { + double lastDistance = -1.0; + for (vector<BSONObj>::const_iterator it = results.begin(); it != results.end(); ++it) { + double distance = (*it)["distance"].numberDouble(); + bool shouldInclude = (*it)["$included"].eoo() || (*it)["$included"].trueValue(); + ASSERT(shouldInclude); + ASSERT_GREATER_THAN_OR_EQUALS(distance, lastDistance); + lastDistance = distance; } +} - TEST(query_stage_near, EmptyResults) { - - vector<BSONObj> mockData; - WorkingSet workingSet; - - MockNearStage nearStage(&workingSet); - - // Empty set of results - mockData.clear(); - nearStage.addInterval(mockData, 0.0, 1.0); +TEST(query_stage_near, Basic) { + vector<BSONObj> mockData; + WorkingSet workingSet; + + MockNearStage nearStage(&workingSet); + + // First set of results + mockData.clear(); + mockData.push_back(BSON("distance" << 0.5)); + mockData.push_back(BSON("distance" << 2.0 << "$included" << false)); // Not included + mockData.push_back(BSON("distance" << 0.0)); + nearStage.addInterval(mockData, 0.0, 1.0); + + // Second set of results + mockData.clear(); + mockData.push_back(BSON("distance" << 1.5)); + mockData.push_back(BSON("distance" << 2.0 << "$included" << false)); // Not included + mockData.push_back(BSON("distance" << 1.0)); + nearStage.addInterval(mockData, 1.0, 2.0); + + // Last set of results + mockData.clear(); + mockData.push_back(BSON("distance" << 2.5)); + mockData.push_back(BSON("distance" << 3.0)); // Included + mockData.push_back(BSON("distance" << 2.0)); + nearStage.addInterval(mockData, 2.0, 3.0); + + vector<BSONObj> results = advanceStage(&nearStage, &workingSet); + ASSERT_EQUALS(results.size(), 7u); + assertAscendingAndValid(results); +} - // Non-empty sest of results - mockData.clear(); - mockData.push_back(BSON("distance" << 1.5)); - mockData.push_back(BSON("distance" << 2.0)); - mockData.push_back(BSON("distance" << 1.0)); - nearStage.addInterval(mockData, 1.0, 2.0); +TEST(query_stage_near, EmptyResults) { + vector<BSONObj> mockData; + WorkingSet workingSet; - vector<BSONObj> results = advanceStage(&nearStage, &workingSet); - ASSERT_EQUALS(results.size(), 3u); - assertAscendingAndValid(results); - } + MockNearStage nearStage(&workingSet); + // Empty set of results + mockData.clear(); + nearStage.addInterval(mockData, 0.0, 1.0); + // Non-empty sest of results + mockData.clear(); + mockData.push_back(BSON("distance" << 1.5)); + mockData.push_back(BSON("distance" << 2.0)); + mockData.push_back(BSON("distance" << 1.0)); + nearStage.addInterval(mockData, 1.0, 2.0); + vector<BSONObj> results = advanceStage(&nearStage, &workingSet); + ASSERT_EQUALS(results.size(), 3u); + assertAscendingAndValid(results); +} } |