PlanStage::StageState EnsureSortedStage::doWork(WorkingSetID* out) { StageState stageState = child()->work(out); if (PlanStage::ADVANCED == stageState) { // We extract the sort key from the WSM's computed data. This must have been generated // by a SortKeyGeneratorStage descendent in the execution tree. WorkingSetMember* member = _ws->get(*out); auto sortKeyComputedData = static_cast<const SortKeyComputedData*>(member->getComputed(WSM_SORT_KEY)); BSONObj curSortKey = sortKeyComputedData->getSortKey(); invariant(!curSortKey.isEmpty()); if (!_prevSortKey.isEmpty() && !isInOrder(_prevSortKey, curSortKey)) { // 'member' is out of order. Drop it from the result set. _ws->free(*out); ++_specificStats.nDropped; return PlanStage::NEED_TIME; } invariant(curSortKey.isOwned()); _prevSortKey = curSortKey; return PlanStage::ADVANCED; } return stageState; }
// static void WorkingSetCommon::initFrom(WorkingSetMember* dest, const WorkingSetMember& src) { dest->loc = src.loc; dest->obj = src.obj; dest->keyData = src.keyData; dest->state = src.state; // Merge computed data. typedef WorkingSetComputedDataType WSCD; for (WSCD i = WSCD(0); i < WSM_COMPUTED_NUM_TYPES; i = WSCD(i + 1)) { if (src.hasComputed(i)) { dest->addComputed(src.getComputed(i)->clone()); } } }
Status SortKeyGenerator::getSortKey(const WorkingSetMember& member, BSONObj* objOut) const { StatusWith<BSONObj> sortKey = BSONObj(); if (member.hasObj()) { sortKey = getSortKeyFromObject(member); } else { sortKey = getSortKeyFromIndexKey(member); } if (!sortKey.isOK()) { return sortKey.getStatus(); } if (!_sortHasMeta) { *objOut = sortKey.getValue(); return Status::OK(); } BSONObjBuilder mergedKeyBob; // Merge metadata into the key. BSONObjIterator it(_rawSortSpec); BSONObjIterator sortKeyIt(sortKey.getValue()); while (it.more()) { BSONElement elt = it.next(); if (elt.isNumber()) { // Merge btree key elt. mergedKeyBob.append(sortKeyIt.next()); } else if (LiteParsedQuery::isTextScoreMeta(elt)) { // Add text score metadata double score = 0.0; if (member.hasComputed(WSM_COMPUTED_TEXT_SCORE)) { const TextScoreComputedData* scoreData = static_cast<const TextScoreComputedData*>( member.getComputed(WSM_COMPUTED_TEXT_SCORE)); score = scoreData->getScore(); } mergedKeyBob.append("$metaTextScore", score); } } *objOut = mergedKeyBob.obj(); return Status::OK(); }
Status SortStageKeyGenerator::getSortKey(const WorkingSetMember& member, BSONObj* objOut) const { BSONObj btreeKeyToUse; Status btreeStatus = getBtreeKey(member.obj, &btreeKeyToUse); if (!btreeStatus.isOK()) { return btreeStatus; } if (!_sortHasMeta) { *objOut = btreeKeyToUse; return Status::OK(); } BSONObjBuilder mergedKeyBob; // Merge metadata into the key. BSONObjIterator it(_rawSortSpec); BSONObjIterator btreeIt(btreeKeyToUse); while (it.more()) { BSONElement elt = it.next(); if (elt.isNumber()) { // Merge btree key elt. mergedKeyBob.append(btreeIt.next()); } else if (LiteParsedQuery::isTextScoreMeta(elt)) { // Add text score metadata double score = 0.0; if (member.hasComputed(WSM_COMPUTED_TEXT_SCORE)) { const TextScoreComputedData* scoreData = static_cast<const TextScoreComputedData*>( member.getComputed(WSM_COMPUTED_TEXT_SCORE)); score = scoreData->getScore(); } mergedKeyBob.append("$metaTextScore", score); } } *objOut = mergedKeyBob.obj(); return Status::OK(); }
PlanStage::StageState SortStage::doWork(WorkingSetID* out) { const size_t maxBytes = static_cast<size_t>(internalQueryExecMaxBlockingSortBytes); if (_memUsage > maxBytes) { mongoutils::str::stream ss; ss << "Sort operation used more than the maximum " << maxBytes << " bytes of RAM. Add an index, or specify a smaller limit."; Status status(ErrorCodes::OperationFailed, ss); *out = WorkingSetCommon::allocateStatusMember(_ws, status); return PlanStage::FAILURE; } if (isEOF()) { return PlanStage::IS_EOF; } // Still reading in results to sort. if (!_sorted) { WorkingSetID id = WorkingSet::INVALID_ID; StageState code = child()->work(&id); if (PlanStage::ADVANCED == code) { // Add it into the map for quick invalidation if it has a valid RecordId. // A RecordId may be invalidated at any time (during a yield). We need to get into // the WorkingSet as quickly as possible to handle it. WorkingSetMember* member = _ws->get(id); // Planner must put a fetch before we get here. verify(member->hasObj()); // We might be sorting something that was invalidated at some point. if (member->hasLoc()) { _wsidByDiskLoc[member->loc] = id; } SortableDataItem item; item.wsid = id; // We extract the sort key from the WSM's computed data. This must have been generated // by a SortKeyGeneratorStage descendent in the execution tree. auto sortKeyComputedData = static_cast<const SortKeyComputedData*>(member->getComputed(WSM_SORT_KEY)); item.sortKey = sortKeyComputedData->getSortKey(); if (member->hasLoc()) { // The RecordId breaks ties when sorting two WSMs with the same sort key. item.loc = member->loc; } addToBuffer(item); return PlanStage::NEED_TIME; } else if (PlanStage::IS_EOF == code) { // TODO: We don't need the lock for this. We could ask for a yield and do this work // unlocked. Also, this is performing a lot of work for one call to work(...) sortBuffer(); _resultIterator = _data.begin(); _sorted = true; return PlanStage::NEED_TIME; } else if (PlanStage::FAILURE == code || PlanStage::DEAD == code) { *out = id; // If a stage fails, it may create a status WSM to indicate why it // failed, in which case 'id' is valid. If ID is invalid, we // create our own error message. if (WorkingSet::INVALID_ID == id) { mongoutils::str::stream ss; ss << "sort stage failed to read in results to sort from child"; Status status(ErrorCodes::InternalError, ss); *out = WorkingSetCommon::allocateStatusMember(_ws, status); } return code; } else if (PlanStage::NEED_YIELD == code) { *out = id; } return code; } // Returning results. verify(_resultIterator != _data.end()); verify(_sorted); *out = _resultIterator->wsid; _resultIterator++; // If we're returning something, take it out of our DL -> WSID map so that future // calls to invalidate don't cause us to take action for a DL we're done with. WorkingSetMember* member = _ws->get(*out); if (member->hasLoc()) { _wsidByDiskLoc.erase(member->loc); } return PlanStage::ADVANCED; }
PlanStage::StageState SortStage::doWork(WorkingSetID* out) { const size_t maxBytes = static_cast<size_t>(internalQueryExecMaxBlockingSortBytes.load()); if (_memUsage > maxBytes) { str::stream ss; ss << "Sort operation used more than the maximum " << maxBytes << " bytes of RAM. Add an index, or specify a smaller limit."; Status status(ErrorCodes::OperationFailed, ss); *out = WorkingSetCommon::allocateStatusMember(_ws, status); return PlanStage::FAILURE; } if (isEOF()) { return PlanStage::IS_EOF; } // Still reading in results to sort. if (!_sorted) { WorkingSetID id = WorkingSet::INVALID_ID; StageState code = child()->work(&id); if (PlanStage::ADVANCED == code) { WorkingSetMember* member = _ws->get(id); SortableDataItem item; item.wsid = id; // We extract the sort key from the WSM's computed data. This must have been generated // by a SortKeyGeneratorStage descendent in the execution tree. auto sortKeyComputedData = static_cast<const SortKeyComputedData*>(member->getComputed(WSM_SORT_KEY)); item.sortKey = sortKeyComputedData->getSortKey(); if (member->hasRecordId()) { // The RecordId breaks ties when sorting two WSMs with the same sort key. item.recordId = member->recordId; } addToBuffer(item); return PlanStage::NEED_TIME; } else if (PlanStage::IS_EOF == code) { // TODO: We don't need the lock for this. We could ask for a yield and do this work // unlocked. Also, this is performing a lot of work for one call to work(...) sortBuffer(); _resultIterator = _data.begin(); _sorted = true; return PlanStage::NEED_TIME; } else if (PlanStage::FAILURE == code) { // The stage which produces a failure is responsible for allocating a working set member // with error details. invariant(WorkingSet::INVALID_ID != id); *out = id; return code; } else if (PlanStage::NEED_YIELD == code) { *out = id; } return code; } // Returning results. verify(_resultIterator != _data.end()); verify(_sorted); *out = _resultIterator->wsid; _resultIterator++; return PlanStage::ADVANCED; }