diff options
author | Kyle Suarez <kyle.suarez@mongodb.com> | 2018-06-18 23:34:49 -0400 |
---|---|---|
committer | Kyle Suarez <kyle.suarez@mongodb.com> | 2018-06-18 23:34:49 -0400 |
commit | 7bc7864fc042b69d36a88c6839c5dd5b4eb20693 (patch) | |
tree | e103e752e5d708aa22dca3ae99ef269d79a5d917 /src/mongo | |
parent | 7c89f48c4f1f1e3ede2931ab602fa118281530a2 (diff) | |
download | mongo-7bc7864fc042b69d36a88c6839c5dd5b4eb20693.tar.gz |
SERVER-35043, SERVER-22949: move geoNear implementation into aggregation
This commit removes the geoNear command and moves its implementation
into the aggregation framework. Users should use the aggregate command
with a $geoNear stage.
The implementation rewrite additionally removes the limit in the
$geoNear aggregation stage. To limit the number of results, use a $limit
stage.
Diffstat (limited to 'src/mongo')
42 files changed, 1000 insertions, 1005 deletions
diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript index d329ef5b1ea..736619ab543 100644 --- a/src/mongo/db/SConscript +++ b/src/mongo/db/SConscript @@ -976,6 +976,7 @@ env.Library( 'query/explain.cpp', 'query/find.cpp', 'pipeline/document_source_cursor.cpp', + 'pipeline/document_source_geo_near_cursor.cpp', 'pipeline/pipeline_d.cpp', 'query/get_executor.cpp', 'query/internal_plans.cpp', diff --git a/src/mongo/db/commands/SConscript b/src/mongo/db/commands/SConscript index 6efa6d3e148..8e718100a33 100644 --- a/src/mongo/db/commands/SConscript +++ b/src/mongo/db/commands/SConscript @@ -212,7 +212,6 @@ env.Library( "explain_cmd.cpp", "find_and_modify.cpp", "find_cmd.cpp", - "geo_near_cmd.cpp", "get_last_error.cpp", "getmore_cmd.cpp", "index_filter_commands.cpp", diff --git a/src/mongo/db/commands/geo_near_cmd.cpp b/src/mongo/db/commands/geo_near_cmd.cpp deleted file mode 100644 index 7385f6645f5..00000000000 --- a/src/mongo/db/commands/geo_near_cmd.cpp +++ /dev/null @@ -1,404 +0,0 @@ -/** -* Copyright (C) 2012-2014 MongoDB 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/>. -* -* As a special exception, the copyright holders give permission to link the -* code of portions of this program with the OpenSSL library under certain -* conditions as described in each individual source file and distribute -* linked combinations including the program with the OpenSSL library. You -* must comply with the GNU Affero General Public License in all respects for -* all of the code used other than as permitted herein. If you modify file(s) -* with this exception, you may extend this exception to your version of the -* file(s), but you are not obligated to do so. If you do not wish to do so, -* delete this exception statement from your version. If you delete this -* exception statement from all source files in the program, then also delete -* it in the license file. -*/ - -#define MONGO_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kCommand - -#include <vector> - -#include "mongo/bson/util/bson_extract.h" -#include "mongo/db/auth/action_set.h" -#include "mongo/db/auth/action_type.h" -#include "mongo/db/auth/privilege.h" -#include "mongo/db/catalog/collection.h" -#include "mongo/db/catalog/database.h" -#include "mongo/db/catalog/index_catalog.h" -#include "mongo/db/client.h" -#include "mongo/db/commands.h" -#include "mongo/db/curop.h" -#include "mongo/db/db_raii.h" -#include "mongo/db/exec/working_set_common.h" -#include "mongo/db/geo/geoconstants.h" -#include "mongo/db/geo/geoparser.h" -#include "mongo/db/index/index_descriptor.h" -#include "mongo/db/index_names.h" -#include "mongo/db/jsobj.h" -#include "mongo/db/matcher/expression_geo.h" -#include "mongo/db/matcher/extensions_callback_real.h" -#include "mongo/db/pipeline/document_source_geo_near.h" -#include "mongo/db/query/explain.h" -#include "mongo/db/query/find_common.h" -#include "mongo/db/query/get_executor.h" -#include "mongo/db/query/plan_summary_stats.h" -#include "mongo/util/log.h" - -namespace mongo { - -using std::unique_ptr; -using std::stringstream; - -/** - * The geoNear command is deprecated. Users should prefer the $near query operator, the $nearSphere - * query operator, or the $geoNear aggregation stage. See - * http://dochub.mongodb.org/core/geoNear-deprecation for more detail. - */ -class Geo2dFindNearCmd : public ErrmsgCommandDeprecated { -public: - Geo2dFindNearCmd() : ErrmsgCommandDeprecated("geoNear") {} - - virtual bool supportsWriteConcern(const BSONObj& cmd) const override { - return false; - } - AllowedOnSecondary secondaryAllowed(ServiceContext*) const override { - return AllowedOnSecondary::kAlways; - } - bool supportsReadConcern(const std::string& dbName, - const BSONObj& cmdObj, - repl::ReadConcernLevel level) const final { - return true; - } - - ReadWriteType getReadWriteType() const { - return ReadWriteType::kRead; - } - - std::size_t reserveBytesForReply() const override { - return FindCommon::kInitReplyBufferSize; - } - - std::string help() const override { - return "http://dochub.mongodb.org/core/geo#GeospatialIndexing-geoNearCommand"; - } - - virtual void addRequiredPrivileges(const std::string& dbname, - const BSONObj& cmdObj, - std::vector<Privilege>* out) const { - ActionSet actions; - actions.addAction(ActionType::find); - out->push_back(Privilege(parseResourcePattern(dbname, cmdObj), actions)); - } - - bool errmsgRun(OperationContext* opCtx, - const string& dbname, - const BSONObj& cmdObj, - string& errmsg, - BSONObjBuilder& result) { - // Do not log the deprecation warning when in a direct client, since the $geoNear - // aggregation stage runs the geoNear command in a direct client. - RARELY if (!opCtx->getClient()->isInDirectClient()) { - warning() << "Support for the geoNear command has been deprecated. Please plan to " - "rewrite geoNear commands using the $near query operator, the $nearSphere " - "query operator, or the $geoNear aggregation stage. See " - "http://dochub.mongodb.org/core/geoNear-deprecation."; - } - - if (!cmdObj["start"].eoo()) { - errmsg = "using deprecated 'start' argument to geoNear"; - return false; - } - - const NamespaceString nss(CommandHelpers::parseNsCollectionRequired(dbname, cmdObj)); - AutoGetCollectionForReadCommand ctx(opCtx, nss); - - Collection* collection = ctx.getCollection(); - if (!collection) { - errmsg = "can't find ns"; - return false; - } - - auto nearFieldName = getFieldName(opCtx, collection, cmdObj); - - PointWithCRS point; - uassert(17304, - "'near' field must be point", - GeoParser::parseQueryPoint(cmdObj["near"], &point).isOK()); - - bool isSpherical = cmdObj["spherical"].trueValue(); - - // Build the $near expression for the query. - BSONObjBuilder nearBob; - if (isSpherical) { - nearBob.append("$nearSphere", cmdObj["near"].Obj()); - } else { - nearBob.append("$near", cmdObj["near"].Obj()); - } - - if (!cmdObj["maxDistance"].eoo()) { - uassert(17299, "maxDistance must be a number", cmdObj["maxDistance"].isNumber()); - nearBob.append("$maxDistance", cmdObj["maxDistance"].number()); - } - - if (!cmdObj["minDistance"].eoo()) { - uassert(17300, "minDistance must be a number", cmdObj["minDistance"].isNumber()); - nearBob.append("$minDistance", cmdObj["minDistance"].number()); - } - - if (!cmdObj["uniqueDocs"].eoo()) { - warning() << nss << ": ignoring deprecated uniqueDocs option in geoNear command"; - } - - // And, build the full query expression. - BSONObjBuilder queryBob; - queryBob.append(nearFieldName, nearBob.obj()); - if (!cmdObj["query"].eoo() && cmdObj["query"].isABSONObj()) { - queryBob.appendElements(cmdObj["query"].Obj()); - } - BSONObj rewritten = queryBob.obj(); - - // Extract the collation, if it exists. - BSONObj collation; - { - BSONElement collationElt; - Status collationEltStatus = - bsonExtractTypedField(cmdObj, "collation", BSONType::Object, &collationElt); - if (!collationEltStatus.isOK() && (collationEltStatus != ErrorCodes::NoSuchKey)) { - uassertStatusOK(collationEltStatus); - } - if (collationEltStatus.isOK()) { - collation = collationElt.Obj(); - } - } - - long long numWanted = 100; - const char* limitName = !cmdObj["num"].eoo() ? "num" : "limit"; - BSONElement eNumWanted = cmdObj[limitName]; - if (!eNumWanted.eoo()) { - uassert(17303, "limit must be number", eNumWanted.isNumber()); - numWanted = eNumWanted.safeNumberLong(); - uassert(17302, "limit must be >=0", numWanted >= 0); - } - - bool includeLocs = false; - if (!cmdObj["includeLocs"].eoo()) { - includeLocs = cmdObj["includeLocs"].trueValue(); - } - - double distanceMultiplier = 1.0; - BSONElement eDistanceMultiplier = cmdObj["distanceMultiplier"]; - if (!eDistanceMultiplier.eoo()) { - uassert(17296, "distanceMultiplier must be a number", eDistanceMultiplier.isNumber()); - distanceMultiplier = eDistanceMultiplier.number(); - uassert(17297, "distanceMultiplier must be non-negative", distanceMultiplier >= 0); - } - - BSONObj projObj = BSON("$pt" << BSON("$meta" << QueryRequest::metaGeoNearPoint) << "$dis" - << BSON("$meta" << QueryRequest::metaGeoNearDistance)); - - auto qr = stdx::make_unique<QueryRequest>(nss); - qr->setFilter(rewritten); - qr->setProj(projObj); - qr->setLimit(numWanted); - qr->setCollation(collation); - const ExtensionsCallbackReal extensionsCallback(opCtx, &nss); - const boost::intrusive_ptr<ExpressionContext> expCtx; - auto statusWithCQ = - CanonicalQuery::canonicalize(opCtx, - std::move(qr), - expCtx, - extensionsCallback, - MatchExpressionParser::kAllowAllSpecialFeatures); - if (!statusWithCQ.isOK()) { - errmsg = "Can't parse filter / create query"; - return false; - } - unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); - - // Prevent chunks from being cleaned up during yields - this allows us to only check the - // version on initial entry into geoNear. - auto rangePreserver = CollectionShardingState::get(opCtx, nss)->getMetadata(opCtx); - - const auto& readConcernArgs = repl::ReadConcernArgs::get(opCtx); - const PlanExecutor::YieldPolicy yieldPolicy = - readConcernArgs.getLevel() == repl::ReadConcernLevel::kSnapshotReadConcern - ? PlanExecutor::INTERRUPT_ONLY - : PlanExecutor::YIELD_AUTO; - auto exec = uassertStatusOK(getExecutor(opCtx, collection, std::move(cq), yieldPolicy, 0)); - - auto curOp = CurOp::get(opCtx); - { - stdx::lock_guard<Client> lk(*opCtx->getClient()); - curOp->setPlanSummary_inlock(Explain::getPlanSummary(exec.get())); - } - - double totalDistance = 0; - BSONObjBuilder resultBuilder(result.subarrayStart("results")); - double farthestDist = 0; - - BSONObj currObj; - long long results = 0; - PlanExecutor::ExecState state; - while (PlanExecutor::ADVANCED == (state = exec->getNext(&currObj, NULL))) { - // Come up with the correct distance. - double dist = currObj["$dis"].number() * distanceMultiplier; - totalDistance += dist; - if (dist > farthestDist) { - farthestDist = dist; - } - - // Strip out '$dis' and '$pt' from the result obj. The rest gets added as 'obj' - // in the command result. - BSONObjIterator resIt(currObj); - BSONObjBuilder resBob; - while (resIt.more()) { - BSONElement elt = resIt.next(); - if (!mongoutils::str::equals("$pt", elt.fieldName()) && - !mongoutils::str::equals("$dis", elt.fieldName())) { - resBob.append(elt); - } - } - BSONObj resObj = resBob.obj(); - - // Don't make a too-big result object. - if (resultBuilder.len() + resObj.objsize() > BSONObjMaxUserSize) { - warning() << "Too many geoNear results for query " << redact(rewritten) - << ", truncating output."; - break; - } - - // Add the next result to the result builder. - BSONObjBuilder oneResultBuilder( - resultBuilder.subobjStart(BSONObjBuilder::numStr(results))); - oneResultBuilder.append("dis", dist); - if (includeLocs) { - oneResultBuilder.appendAs(currObj["$pt"], "loc"); - } - oneResultBuilder.append("obj", resObj); - oneResultBuilder.done(); - - ++results; - - // Break if we have the number of requested result documents. - if (results >= numWanted) { - break; - } - } - - resultBuilder.done(); - - // Return an error if execution fails for any reason. - if (PlanExecutor::FAILURE == state || PlanExecutor::DEAD == state) { - log() << "Plan executor error during geoNear command: " << PlanExecutor::statestr(state) - << ", stats: " << redact(Explain::getWinningPlanStats(exec.get())); - - uassertStatusOK(WorkingSetCommon::getMemberObjectStatus(currObj).withContext( - "Executor error during geoNear command")); - } - - PlanSummaryStats summary; - Explain::getSummaryStats(*exec, &summary); - - // Fill out the stats subobj. - BSONObjBuilder stats(result.subobjStart("stats")); - - stats.appendNumber("nscanned", summary.totalKeysExamined); - stats.appendNumber("objectsLoaded", summary.totalDocsExamined); - - if (results > 0) { - stats.append("avgDistance", totalDistance / results); - } - stats.append("maxDistance", farthestDist); - stats.appendIntOrLL("time", - durationCount<Microseconds>(curOp->elapsedTimeExcludingPauses())); - stats.done(); - - collection->infoCache()->notifyOfQuery(opCtx, summary.indexesUsed); - - curOp->debug().setPlanSummaryMetrics(summary); - - if (curOp->shouldDBProfile()) { - BSONObjBuilder execStatsBob; - Explain::getWinningPlanStats(exec.get(), &execStatsBob); - curOp->debug().execStats = execStatsBob.obj(); - } - - return true; - } - -private: - /** - * Given a collection and the geoNear command parameters, returns the field path over which - * the geoNear should operate. - * - * Throws an assertion with ErrorCodes::IndexNotFound if there is no single geo index - * which this geoNear command should use. - */ - StringData getFieldName(OperationContext* opCtx, Collection* collection, BSONObj cmdObj) { - if (auto keyElt = cmdObj[DocumentSourceGeoNear::kKeyFieldName]) { - uassert(ErrorCodes::TypeMismatch, - str::stream() << "geoNear parameter '" << DocumentSourceGeoNear::kKeyFieldName - << "' must be of type string but found type: " - << typeName(keyElt.type()), - keyElt.type() == BSONType::String); - auto fieldName = keyElt.valueStringData(); - uassert(ErrorCodes::BadValue, - str::stream() << "$geoNear parameter '" << DocumentSourceGeoNear::kKeyFieldName - << "' cannot be the empty string", - !fieldName.empty()); - return fieldName; - } - - vector<IndexDescriptor*> idxs; - - // First, try 2d. - collection->getIndexCatalog()->findIndexByType(opCtx, IndexNames::GEO_2D, idxs); - uassert(ErrorCodes::IndexNotFound, - "more than one 2d index, not sure which to run geoNear on", - idxs.size() <= 1u); - - if (1 == idxs.size()) { - BSONObj indexKp = idxs[0]->keyPattern(); - BSONObjIterator kpIt(indexKp); - while (kpIt.more()) { - BSONElement elt = kpIt.next(); - if (BSONType::String == elt.type() && IndexNames::GEO_2D == elt.valuestr()) { - return elt.fieldNameStringData(); - } - } - } - - // Next, 2dsphere. - idxs.clear(); - collection->getIndexCatalog()->findIndexByType(opCtx, IndexNames::GEO_2DSPHERE, idxs); - uassert(ErrorCodes::IndexNotFound, "no geo indices for geoNear", !idxs.empty()); - uassert(ErrorCodes::IndexNotFound, - "more than one 2dsphere index, not sure which to run geoNear on", - idxs.size() == 1u); - - // 1 == idx.size(). - BSONObj indexKp = idxs[0]->keyPattern(); - BSONObjIterator kpIt(indexKp); - while (kpIt.more()) { - BSONElement elt = kpIt.next(); - if (BSONType::String == elt.type() && IndexNames::GEO_2DSPHERE == elt.valuestr()) { - return elt.fieldNameStringData(); - } - } - - MONGO_UNREACHABLE; - } -} geo2dFindNearCmd; -} // namespace mongo diff --git a/src/mongo/db/pipeline/cluster_aggregation_planner.cpp b/src/mongo/db/pipeline/cluster_aggregation_planner.cpp index 53ea48dd68f..5148d56f888 100644 --- a/src/mongo/db/pipeline/cluster_aggregation_planner.cpp +++ b/src/mongo/db/pipeline/cluster_aggregation_planner.cpp @@ -100,10 +100,7 @@ void moveFinalUnwindFromShardsToMerger(Pipeline* shardPipe, Pipeline* mergePipe) * Documents. */ void limitFieldsSentFromShardsToMerger(Pipeline* shardPipe, Pipeline* mergePipe) { - auto depsMetadata = DocumentSourceMatch::isTextQuery(shardPipe->getInitialQuery()) - ? DepsTracker::MetadataAvailable::kTextScore - : DepsTracker::MetadataAvailable::kNoMetadata; - DepsTracker mergeDeps(mergePipe->getDependencies(depsMetadata)); + DepsTracker mergeDeps(mergePipe->getDependencies(DepsTracker::kAllMetadataAvailable)); if (mergeDeps.needWholeDocument) return; // the merge needs all fields, so nothing we can do. @@ -113,7 +110,7 @@ void limitFieldsSentFromShardsToMerger(Pipeline* shardPipe, Pipeline* mergePipe) // Remove metadata from dependencies since it automatically flows through projection and we // don't want to project it in to the document. - mergeDeps.setNeedTextScore(false); + mergeDeps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, false); // HEURISTIC: only apply optimization if none of the shard stages have an exhaustive list of // field dependencies. While this may not be 100% ideal in all cases, it is simple and @@ -125,7 +122,7 @@ void limitFieldsSentFromShardsToMerger(Pipeline* shardPipe, Pipeline* mergePipe) // 2) Optimization IS NOT applied immediately following a $project or $group since it would // add an unnecessary project (and therefore a deep-copy). for (auto&& source : shardPipe->getSources()) { - DepsTracker dt(depsMetadata); + DepsTracker dt(DepsTracker::kAllMetadataAvailable); if (source->getDependencies(&dt) & DocumentSource::EXHAUSTIVE_FIELDS) return; } diff --git a/src/mongo/db/pipeline/dependencies.cpp b/src/mongo/db/pipeline/dependencies.cpp index a1fc3b88a19..df24a8bb9ed 100644 --- a/src/mongo/db/pipeline/dependencies.cpp +++ b/src/mongo/db/pipeline/dependencies.cpp @@ -41,6 +41,8 @@ using std::vector; namespace str = mongoutils::str; +constexpr DepsTracker::MetadataAvailable DepsTracker::kAllGeoNearDataAvailable; + bool DepsTracker::_appendMetaProjections(BSONObjBuilder* projectionBuilder) const { if (_needTextScore) { projectionBuilder->append(Document::metaFieldTextScore, @@ -52,7 +54,17 @@ bool DepsTracker::_appendMetaProjections(BSONObjBuilder* projectionBuilder) cons BSON("$meta" << "sortKey")); } - return (_needTextScore || _needSortKey); + if (_needGeoNearDistance) { + projectionBuilder->append(Document::metaFieldGeoNearDistance, + BSON("$meta" + << "geoNearDistance")); + } + if (_needGeoNearPoint) { + projectionBuilder->append(Document::metaFieldGeoNearPoint, + BSON("$meta" + << "geoNearPoint")); + } + return (_needTextScore || _needSortKey || _needGeoNearDistance || _needGeoNearPoint); } BSONObj DepsTracker::toProjection() const { @@ -134,6 +146,81 @@ boost::optional<ParsedDeps> DepsTracker::toParsedDeps() const { return ParsedDeps(md.freeze()); } +bool DepsTracker::getNeedsMetadata(MetadataType type) const { + switch (type) { + case MetadataType::TEXT_SCORE: + return _needTextScore; + case MetadataType::SORT_KEY: + return _needSortKey; + case MetadataType::GEO_NEAR_DISTANCE: + return _needGeoNearDistance; + case MetadataType::GEO_NEAR_POINT: + return _needGeoNearPoint; + } + MONGO_UNREACHABLE; +} + +bool DepsTracker::isMetadataAvailable(MetadataType type) const { + switch (type) { + case MetadataType::TEXT_SCORE: + return _metadataAvailable & MetadataAvailable::kTextScore; + case MetadataType::SORT_KEY: + MONGO_UNREACHABLE; + case MetadataType::GEO_NEAR_DISTANCE: + return _metadataAvailable & MetadataAvailable::kGeoNearDistance; + case MetadataType::GEO_NEAR_POINT: + return _metadataAvailable & MetadataAvailable::kGeoNearPoint; + } + MONGO_UNREACHABLE; +} + +void DepsTracker::setNeedsMetadata(MetadataType type, bool required) { + switch (type) { + case MetadataType::TEXT_SCORE: + uassert(40218, + "pipeline requires text score metadata, but there is no text score available", + !required || isMetadataAvailable(type)); + _needTextScore = required; + return; + case MetadataType::SORT_KEY: + invariant(required || !_needSortKey); + _needSortKey = required; + return; + case MetadataType::GEO_NEAR_DISTANCE: + uassert(50860, + "pipeline requires $geoNear distance metadata, but it is not available", + !required || isMetadataAvailable(type)); + invariant(required || !_needGeoNearDistance); + _needGeoNearDistance = required; + return; + case MetadataType::GEO_NEAR_POINT: + uassert(50859, + "pipeline requires $geoNear point metadata, but it is not available", + !required || isMetadataAvailable(type)); + invariant(required || !_needGeoNearPoint); + _needGeoNearPoint = required; + return; + } + MONGO_UNREACHABLE; +} + +std::vector<DepsTracker::MetadataType> DepsTracker::getAllRequiredMetadataTypes() const { + std::vector<MetadataType> reqs; + if (_needTextScore) { + reqs.push_back(MetadataType::TEXT_SCORE); + } + if (_needSortKey) { + reqs.push_back(MetadataType::SORT_KEY); + } + if (_needGeoNearDistance) { + reqs.push_back(MetadataType::GEO_NEAR_DISTANCE); + } + if (_needGeoNearPoint) { + reqs.push_back(MetadataType::GEO_NEAR_POINT); + } + return reqs; +} + namespace { // Mutually recursive with arrayHelper Document documentHelper(const BSONObj& bson, const Document& neededFields, int nFieldsNeeded = -1); diff --git a/src/mongo/db/pipeline/dependencies.h b/src/mongo/db/pipeline/dependencies.h index 133fc1d9893..dec5296060f 100644 --- a/src/mongo/db/pipeline/dependencies.h +++ b/src/mongo/db/pipeline/dependencies.h @@ -43,9 +43,43 @@ class ParsedDeps; */ struct DepsTracker { /** + * Represents the type of metadata a pipeline might request. + */ + enum class MetadataType { + // The score associated with a text match. + TEXT_SCORE, + + // The key to use for sorting. + SORT_KEY, + + // The computed distance for a near query. + GEO_NEAR_DISTANCE, + + // The point used in the computation of the GEO_NEAR_DISTANCE. + GEO_NEAR_POINT, + }; + + /** * Represents what metadata is available on documents that are input to the pipeline. */ - enum MetadataAvailable { kNoMetadata = 0, kTextScore = 1 }; + enum MetadataAvailable { + kNoMetadata = 0, + kTextScore = 1 << 1, + kGeoNearDistance = 1 << 2, + kGeoNearPoint = 1 << 3, + }; + + /** + * Represents a state where all geo metadata is available. + */ + static constexpr auto kAllGeoNearDataAvailable = + MetadataAvailable(MetadataAvailable::kGeoNearDistance | MetadataAvailable::kGeoNearPoint); + + /** + * Represents a state where all metadata is available. + */ + static constexpr auto kAllMetadataAvailable = + MetadataAvailable(kTextScore | kGeoNearDistance | kGeoNearPoint); DepsTracker(MetadataAvailable metadataAvailable = kNoMetadata) : _metadataAvailable(metadataAvailable) {} @@ -71,36 +105,44 @@ struct DepsTracker { return !match.empty(); } + /** + * Returns a value with bits set indicating the types of metadata available. + */ MetadataAvailable getMetadataAvailable() const { return _metadataAvailable; } - bool isTextScoreAvailable() const { - return _metadataAvailable & MetadataAvailable::kTextScore; - } + /** + * Returns true if the DepsTracker the metadata 'type' is available to the pipeline. It is + * illegal to call this with MetadataType::SORT_KEY, since the sort key will always be available + * if needed. + */ + bool isMetadataAvailable(MetadataType type) const; - bool getNeedTextScore() const { - return _needTextScore; - } + /** + * Sets whether or not metadata 'type' is required. Throws if 'required' is true but that + * metadata is not available to the pipeline. + * + * Except for MetadataType::SORT_KEY, once 'type' is required, it cannot be unset. + */ + void setNeedsMetadata(MetadataType type, bool required); - void setNeedTextScore(bool needTextScore) { - if (needTextScore && !isTextScoreAvailable()) { - uasserted( - 40218, - "pipeline requires text score metadata, but there is no text score available"); - } - _needTextScore = needTextScore; - } + /** + * Returns true if the DepsTracker requires that metadata of type 'type' is present. + */ + bool getNeedsMetadata(MetadataType type) const; - bool getNeedSortKey() const { - return _needSortKey; + /** + * Returns true if there exists a type of metadata required by the DepsTracker. + */ + bool getNeedsAnyMetadata() const { + return _needTextScore || _needSortKey || _needGeoNearDistance || _needGeoNearPoint; } - void setNeedSortKey(bool needSortKey) { - // We don't expect to ever unset '_needSortKey'. - invariant(!_needSortKey || needSortKey); - _needSortKey = needSortKey; - } + /** + * Returns a vector containing all the types of metadata required by this DepsTracker. + */ + std::vector<MetadataType> getAllRequiredMetadataTypes() const; std::set<std::string> fields; // Names of needed fields in dotted notation. std::set<Variables::Id> vars; // IDs of referenced variables. @@ -114,8 +156,12 @@ private: bool _appendMetaProjections(BSONObjBuilder* bb) const; MetadataAvailable _metadataAvailable; - bool _needTextScore = false; // if true, add a {$meta: "textScore"} to the projection. - bool _needSortKey = false; // if true, add a {$meta: "sortKey"} to the projection. + + // Each member variable influences a different $meta projection. + bool _needTextScore = false; // {$meta: "textScore"} + bool _needSortKey = false; // {$meta: "sortKey"} + bool _needGeoNearDistance = false; // {$meta: "geoNearDistance"} + bool _needGeoNearPoint = false; // {$meta: "geoNearPoint"} }; /** diff --git a/src/mongo/db/pipeline/dependencies_test.cpp b/src/mongo/db/pipeline/dependencies_test.cpp index cb33e1fcc9f..733cc94dd60 100644 --- a/src/mongo/db/pipeline/dependencies_test.cpp +++ b/src/mongo/db/pipeline/dependencies_test.cpp @@ -116,7 +116,7 @@ TEST(DependenciesToProjectionTest, ShouldOnlyRequestTextScoreIfEntireDocumentAnd DepsTracker deps(DepsTracker::MetadataAvailable::kTextScore); deps.fields = arrayToSet(array); deps.needWholeDocument = true; - deps.setNeedTextScore(true); + deps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); ASSERT_BSONOBJ_EQ(deps.toProjection(), BSON(Document::metaFieldTextScore << metaTextScore)); } @@ -125,7 +125,7 @@ TEST(DependenciesToProjectionTest, const char* array[] = {"a"}; // needTextScore without needWholeDocument DepsTracker deps(DepsTracker::MetadataAvailable::kTextScore); deps.fields = arrayToSet(array); - deps.setNeedTextScore(true); + deps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); ASSERT_BSONOBJ_EQ( deps.toProjection(), BSON(Document::metaFieldTextScore << metaTextScore << "a" << 1 << "_id" << 0)); @@ -135,7 +135,7 @@ TEST(DependenciesToProjectionTest, ShouldProduceEmptyObjectIfThereAreNoDependenc DepsTracker deps(DepsTracker::MetadataAvailable::kTextScore); deps.fields = {}; deps.needWholeDocument = false; - deps.setNeedTextScore(false); + deps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, false); ASSERT_BSONOBJ_EQ(deps.toProjection(), BSONObj()); } @@ -143,7 +143,7 @@ TEST(DependenciesToProjectionTest, ShouldAttemptToExcludeOtherFieldsIfOnlyTextSc DepsTracker deps(DepsTracker::MetadataAvailable::kTextScore); deps.fields = {}; deps.needWholeDocument = false; - deps.setNeedTextScore(true); + deps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); ASSERT_BSONOBJ_EQ(deps.toProjection(), BSON(Document::metaFieldTextScore << metaTextScore << "_id" << 0 << "$noFieldsNeeded" @@ -155,7 +155,7 @@ TEST(DependenciesToProjectionTest, DepsTracker deps(DepsTracker::MetadataAvailable::kTextScore); deps.fields = {}; deps.needWholeDocument = true; - deps.setNeedTextScore(true); + deps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); ASSERT_BSONOBJ_EQ(deps.toProjection(), BSON(Document::metaFieldTextScore << metaTextScore)); } diff --git a/src/mongo/db/pipeline/document.cpp b/src/mongo/db/pipeline/document.cpp index 4f531fe72f3..f6f88b504fc 100644 --- a/src/mongo/db/pipeline/document.cpp +++ b/src/mongo/db/pipeline/document.cpp @@ -45,8 +45,11 @@ using std::vector; const DocumentStorage DocumentStorage::kEmptyDoc; -const std::vector<StringData> Document::allMetadataFieldNames = { - Document::metaFieldTextScore, Document::metaFieldRandVal, Document::metaFieldSortKey}; +const std::vector<StringData> Document::allMetadataFieldNames = {Document::metaFieldTextScore, + Document::metaFieldRandVal, + Document::metaFieldSortKey, + Document::metaFieldGeoNearDistance, + Document::metaFieldGeoNearPoint}; Position DocumentStorage::findField(StringData requested) const { int reqSize = requested.size(); // get size calculation out of the way if needed @@ -205,6 +208,8 @@ intrusive_ptr<DocumentStorage> DocumentStorage::clone() const { out->_textScore = _textScore; out->_randVal = _randVal; out->_sortKey = _sortKey.getOwned(); + out->_geoNearDistance = _geoNearDistance; + out->_geoNearPoint = _geoNearPoint.getOwned(); // Tell values that they have been memcpyed (updates ref counts) for (DocumentStorageIterator it = out->iteratorAll(); !it.atEnd(); it.advance()) { @@ -272,6 +277,8 @@ BSONObj Document::toBson() const { constexpr StringData Document::metaFieldTextScore; constexpr StringData Document::metaFieldRandVal; constexpr StringData Document::metaFieldSortKey; +constexpr StringData Document::metaFieldGeoNearDistance; +constexpr StringData Document::metaFieldGeoNearPoint; BSONObj Document::toBsonWithMetaData() const { BSONObjBuilder bb; @@ -282,6 +289,10 @@ BSONObj Document::toBsonWithMetaData() const { bb.append(metaFieldRandVal, getRandMetaField()); if (hasSortKeyMetaField()) bb.append(metaFieldSortKey, getSortKeyMetaField()); + if (hasGeoNearDistance()) + bb.append(metaFieldGeoNearDistance, getGeoNearDistance()); + if (hasGeoNearPoint()) + getGeoNearPoint().addToBsonObj(&bb, metaFieldGeoNearPoint); return bb.obj(); } @@ -302,6 +313,20 @@ Document Document::fromBsonWithMetaData(const BSONObj& bson) { } else if (fieldName == metaFieldSortKey) { md.setSortKeyMetaField(elem.Obj()); continue; + } else if (fieldName == metaFieldGeoNearDistance) { + md.setGeoNearDistance(elem.Double()); + continue; + } else if (fieldName == metaFieldGeoNearPoint) { + Value val; + if (elem.type() == BSONType::Array) { + val = Value(BSONArray(elem.embeddedObject())); + } else { + invariant(elem.type() == BSONType::Object); + val = Value(elem.embeddedObject()); + } + + md.setGeoNearPoint(val); + continue; } } diff --git a/src/mongo/db/pipeline/document.h b/src/mongo/db/pipeline/document.h index eb7438b776e..76d207276ce 100644 --- a/src/mongo/db/pipeline/document.h +++ b/src/mongo/db/pipeline/document.h @@ -93,6 +93,8 @@ public: static constexpr StringData metaFieldTextScore = "$textScore"_sd; static constexpr StringData metaFieldRandVal = "$randVal"_sd; static constexpr StringData metaFieldSortKey = "$sortKey"_sd; + static constexpr StringData metaFieldGeoNearDistance = "$dis"_sd; + static constexpr StringData metaFieldGeoNearPoint = "$pt"_sd; static const std::vector<StringData> allMetadataFieldNames; @@ -266,6 +268,20 @@ public: return storage().getSortKeyMetaField(); } + bool hasGeoNearDistance() const { + return storage().hasGeoNearDistance(); + } + double getGeoNearDistance() const { + return storage().getGeoNearDistance(); + } + + bool hasGeoNearPoint() const { + return storage().hasGeoNearPoint(); + } + Value getGeoNearPoint() const { + return storage().getGeoNearPoint(); + } + /// members for Sorter struct SorterDeserializeSettings {}; // unused void serializeForSorter(BufBuilder& buf) const; @@ -518,6 +534,14 @@ public: storage().setSortKeyMetaField(sortKey); } + void setGeoNearDistance(double dist) { + storage().setGeoNearDistance(dist); + } + + void setGeoNearPoint(Value point) { + storage().setGeoNearPoint(std::move(point)); + } + /** Convert to a read-only document and release reference. * * Call this to indicate that you are done with this Document and will diff --git a/src/mongo/db/pipeline/document_internal.h b/src/mongo/db/pipeline/document_internal.h index baa68e60658..e1f7c299fc6 100644 --- a/src/mongo/db/pipeline/document_internal.h +++ b/src/mongo/db/pipeline/document_internal.h @@ -189,7 +189,8 @@ public: _hashTabMask(0), _metaFields(), _textScore(0), - _randVal(0) {} + _randVal(0), + _geoNearDistance(0) {} ~DocumentStorage(); @@ -197,7 +198,10 @@ public: TEXT_SCORE, RAND_VAL, SORT_KEY, + GEONEAR_DIST, + GEONEAR_POINT, + // New fields must be added before the NUM_FIELDS sentinel. NUM_FIELDS }; @@ -284,6 +288,12 @@ public: if (source.hasSortKeyMetaField()) { setSortKeyMetaField(source.getSortKeyMetaField()); } + if (source.hasGeoNearDistance()) { + setGeoNearDistance(source.getGeoNearDistance()); + } + if (source.hasGeoNearPoint()) { + setGeoNearPoint(source.getGeoNearPoint()); + } } bool hasTextScore() const { @@ -319,6 +329,28 @@ public: _sortKey = sortKey.getOwned(); } + bool hasGeoNearDistance() const { + return _metaFields.test(MetaType::GEONEAR_DIST); + } + double getGeoNearDistance() const { + return _geoNearDistance; + } + void setGeoNearDistance(double dist) { + _metaFields.set(MetaType::GEONEAR_DIST); + _geoNearDistance = dist; + } + + bool hasGeoNearPoint() const { + return _metaFields.test(MetaType::GEONEAR_POINT); + } + Value getGeoNearPoint() const { + return _geoNearPoint; + } + void setGeoNearPoint(Value point) { + _metaFields.set(MetaType::GEONEAR_POINT); + _geoNearPoint = std::move(point); + } + private: /// Same as lastElement->next() or firstElement() if empty. const ValueElement* end() const { @@ -401,6 +433,8 @@ private: double _textScore; double _randVal; BSONObj _sortKey; + double _geoNearDistance; + Value _geoNearPoint; // When adding a field, make sure to update clone() method // Defined in document.cpp diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h index de383b154e2..6e706876aff 100644 --- a/src/mongo/db/pipeline/document_source.h +++ b/src/mongo/db/pipeline/document_source.h @@ -648,7 +648,7 @@ public: /** * Get the dependencies this operation needs to do its job. If overridden, subclasses must add * all paths needed to apply their transformation to 'deps->fields', and call - * 'deps->setNeedTextScore()' if the text score is required. + * 'deps->setNeedsMetadata()' to indicate what metadata (e.g. text score), if any, is required. * * See GetDepsReturn above for the possible return values and what they mean. */ diff --git a/src/mongo/db/pipeline/document_source_add_fields_test.cpp b/src/mongo/db/pipeline/document_source_add_fields_test.cpp index fa9b27f5d61..8fa400e731a 100644 --- a/src/mongo/db/pipeline/document_source_add_fields_test.cpp +++ b/src/mongo/db/pipeline/document_source_add_fields_test.cpp @@ -133,7 +133,7 @@ TEST_F(AddFieldsTest, ShouldAddReferencedFieldsToDependencies) { ASSERT_EQUALS(1U, dependencies.fields.count("c")); ASSERT_EQUALS(1U, dependencies.fields.count("d")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(true, dependencies.getNeedTextScore()); + ASSERT_EQUALS(true, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(AddFieldsTest, ShouldPropagatePauses) { diff --git a/src/mongo/db/pipeline/document_source_bucket_auto_test.cpp b/src/mongo/db/pipeline/document_source_bucket_auto_test.cpp index df19b4eb497..157c2aad4f1 100644 --- a/src/mongo/db/pipeline/document_source_bucket_auto_test.cpp +++ b/src/mongo/db/pipeline/document_source_bucket_auto_test.cpp @@ -447,7 +447,7 @@ TEST_F(BucketAutoTests, ShouldAddDependenciesOfGroupByFieldAndComputedFields) { ASSERT_EQUALS(1U, dependencies.fields.count("b")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(BucketAutoTests, ShouldNeedTextScoreInDependenciesFromGroupByField) { @@ -459,7 +459,7 @@ TEST_F(BucketAutoTests, ShouldNeedTextScoreInDependenciesFromGroupByField) { ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(true, dependencies.getNeedTextScore()); + ASSERT_EQUALS(true, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(BucketAutoTests, ShouldNeedTextScoreInDependenciesFromOutputField) { @@ -475,7 +475,7 @@ TEST_F(BucketAutoTests, ShouldNeedTextScoreInDependenciesFromOutputField) { ASSERT_EQUALS(1U, dependencies.fields.count("x")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(true, dependencies.getNeedTextScore()); + ASSERT_EQUALS(true, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(BucketAutoTests, SerializesDefaultAccumulatorIfOutputFieldIsNotSpecified) { diff --git a/src/mongo/db/pipeline/document_source_cursor.cpp b/src/mongo/db/pipeline/document_source_cursor.cpp index 10fe39a1484..aeb16094abb 100644 --- a/src/mongo/db/pipeline/document_source_cursor.cpp +++ b/src/mongo/db/pipeline/document_source_cursor.cpp @@ -63,6 +63,10 @@ DocumentSource::GetNextResult DocumentSourceCursor::getNext() { return std::move(out); } +Document DocumentSourceCursor::transformBSONObjToDocument(const BSONObj& obj) const { + return _dependencies ? _dependencies->extractFields(obj) : Document::fromBsonWithMetaData(obj); +} + void DocumentSourceCursor::loadBatch() { if (!_exec || _exec->isDisposed()) { // No more documents. @@ -85,10 +89,8 @@ void DocumentSourceCursor::loadBatch() { while ((state = _exec->getNext(&resultObj, nullptr)) == PlanExecutor::ADVANCED) { if (_shouldProduceEmptyDocs) { _currentBatch.push_back(Document()); - } else if (_dependencies) { - _currentBatch.push_back(_dependencies->extractFields(resultObj)); } else { - _currentBatch.push_back(Document::fromBsonWithMetaData(resultObj)); + _currentBatch.push_back(transformBSONObjToDocument(resultObj)); } if (_limit) { @@ -302,6 +304,8 @@ DocumentSourceCursor::DocumentSourceCursor( _docsAddedToBatches(0), _exec(std::move(exec)), _outputSorts(_exec->getOutputSorts()) { + // Later code in the DocumentSourceCursor lifecycle expects that '_exec' is in a saved state. + _exec->saveState(); _planSummary = Explain::getPlanSummary(_exec.get()); recordPlanSummaryStats(); diff --git a/src/mongo/db/pipeline/document_source_cursor.h b/src/mongo/db/pipeline/document_source_cursor.h index fc9b37f9403..fa10f450d4e 100644 --- a/src/mongo/db/pipeline/document_source_cursor.h +++ b/src/mongo/db/pipeline/document_source_cursor.h @@ -43,14 +43,17 @@ namespace mongo { /** * Constructs and returns Documents from the BSONObj objects produced by a supplied PlanExecutor. */ -class DocumentSourceCursor final : public DocumentSource { +class DocumentSourceCursor : public DocumentSource { public: // virtuals from DocumentSource GetNextResult getNext() final; - const char* getSourceName() const final; - BSONObjSet getOutputSorts() final { + + const char* getSourceName() const override; + + BSONObjSet getOutputSorts() override { return _outputSorts; } + Value serialize(boost::optional<ExplainOptions::Verbosity> explain = boost::none) const final; StageConstraints constraints(Pipeline::SplitState pipeState) const final { @@ -150,6 +153,12 @@ public: } protected: + DocumentSourceCursor(Collection* collection, + std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec, + const boost::intrusive_ptr<ExpressionContext>& pExpCtx); + + ~DocumentSourceCursor(); + /** * Disposes of '_exec' if it hasn't been disposed already. This involves taking a collection * lock. @@ -162,12 +171,15 @@ protected: Pipeline::SourceContainer::iterator doOptimizeAt(Pipeline::SourceContainer::iterator itr, Pipeline::SourceContainer* container) final; -private: - DocumentSourceCursor(Collection* collection, - std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec, - const boost::intrusive_ptr<ExpressionContext>& pExpCtx); - ~DocumentSourceCursor(); + /** + * If '_shouldProduceEmptyDocs' is false, this function hook is called on each 'obj' returned by + * '_exec' when loading a batch and returns a Document to be added to '_currentBatch'. + * + * The default implementation is a dependency-aware BSONObj-to-Document transformation. + */ + virtual Document transformBSONObjToDocument(const BSONObj& obj) const; +private: /** * Acquires the appropriate locks, then destroys and de-registers '_exec'. '_exec' must be * non-null. @@ -180,12 +192,14 @@ private: void cleanupExecutor(const AutoGetCollectionForRead& readLock); /** - * Reads a batch of data from '_exec'. + * Reads a batch of data from '_exec'. Subclasses can specify custom behavior to be performed on + * each document by overloading transformBSONObjToDocument(). */ void loadBatch(); void recordPlanSummaryStats(); + // Batches results returned from the underlying PlanExecutor. std::deque<Document> _currentBatch; // BSONObj members must outlive _projection and cursor. diff --git a/src/mongo/db/pipeline/document_source_facet.cpp b/src/mongo/db/pipeline/document_source_facet.cpp index bc49ab6371f..8b1ce60f6e1 100644 --- a/src/mongo/db/pipeline/document_source_facet.cpp +++ b/src/mongo/db/pipeline/document_source_facet.cpp @@ -276,11 +276,17 @@ DocumentSource::GetDepsReturn DocumentSourceFacet::getDependencies(DepsTracker* deps->vars.insert(subDepsTracker.vars.begin(), subDepsTracker.vars.end()); deps->needWholeDocument = deps->needWholeDocument || subDepsTracker.needWholeDocument; - deps->setNeedTextScore(deps->getNeedTextScore() || subDepsTracker.getNeedTextScore()); + + // The text score is the only type of metadata that could be needed by $facet. + deps->setNeedsMetadata( + DepsTracker::MetadataType::TEXT_SCORE, + deps->getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE) || + subDepsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); // If there are variables defined at this stage's scope, there may be dependencies upon // them in subsequent pipelines. Keep enumerating. - if (deps->needWholeDocument && deps->getNeedTextScore() && !scopeHasVariables) { + if (deps->needWholeDocument && + deps->getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE) && !scopeHasVariables) { break; } } diff --git a/src/mongo/db/pipeline/document_source_facet_test.cpp b/src/mongo/db/pipeline/document_source_facet_test.cpp index 25f8cdc7a07..bc88b491229 100644 --- a/src/mongo/db/pipeline/document_source_facet_test.cpp +++ b/src/mongo/db/pipeline/document_source_facet_test.cpp @@ -564,7 +564,7 @@ TEST_F(DocumentSourceFacetTest, ShouldUnionDependenciesOfInnerPipelines) { DepsTracker deps(DepsTracker::MetadataAvailable::kNoMetadata); ASSERT_EQ(facetStage->getDependencies(&deps), DocumentSource::GetDepsReturn::EXHAUSTIVE_ALL); ASSERT_FALSE(deps.needWholeDocument); - ASSERT_FALSE(deps.getNeedTextScore()); + ASSERT_FALSE(deps.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); ASSERT_EQ(deps.fields.size(), 2UL); ASSERT_EQ(deps.fields.count("a"), 1UL); ASSERT_EQ(deps.fields.count("b"), 1UL); @@ -602,7 +602,7 @@ TEST_F(DocumentSourceFacetTest, ShouldRequireWholeDocumentIfAnyPipelineRequiresW DepsTracker deps(DepsTracker::MetadataAvailable::kNoMetadata); ASSERT_EQ(facetStage->getDependencies(&deps), DocumentSource::GetDepsReturn::EXHAUSTIVE_ALL); ASSERT_TRUE(deps.needWholeDocument); - ASSERT_FALSE(deps.getNeedTextScore()); + ASSERT_FALSE(deps.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } /** @@ -611,7 +611,7 @@ TEST_F(DocumentSourceFacetTest, ShouldRequireWholeDocumentIfAnyPipelineRequiresW class DocumentSourceNeedsOnlyTextScore : public DocumentSourcePassthrough { public: GetDepsReturn getDependencies(DepsTracker* deps) const override { - deps->setNeedTextScore(true); + deps->setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); return GetDepsReturn::EXHAUSTIVE_ALL; } static boost::intrusive_ptr<DocumentSourceNeedsOnlyTextScore> create() { @@ -641,7 +641,7 @@ TEST_F(DocumentSourceFacetTest, ShouldRequireTextScoreIfAnyPipelineRequiresTextS DepsTracker deps(DepsTracker::MetadataAvailable::kTextScore); ASSERT_EQ(facetStage->getDependencies(&deps), DocumentSource::GetDepsReturn::EXHAUSTIVE_ALL); ASSERT_TRUE(deps.needWholeDocument); - ASSERT_TRUE(deps.getNeedTextScore()); + ASSERT_TRUE(deps.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceFacetTest, ShouldThrowIfAnyPipelineRequiresTextScoreButItIsNotAvailable) { diff --git a/src/mongo/db/pipeline/document_source_geo_near.cpp b/src/mongo/db/pipeline/document_source_geo_near.cpp index 39a31253573..7ba5e2a4f8e 100644 --- a/src/mongo/db/pipeline/document_source_geo_near.cpp +++ b/src/mongo/db/pipeline/document_source_geo_near.cpp @@ -33,84 +33,26 @@ #include "mongo/db/pipeline/document_source_geo_near.h" #include "mongo/db/pipeline/document.h" -#include "mongo/db/pipeline/document_source_limit.h" #include "mongo/db/pipeline/document_source_sort.h" #include "mongo/db/pipeline/lite_parsed_document_source.h" -#include "mongo/rpc/get_status_from_command_result.h" #include "mongo/util/log.h" namespace mongo { using boost::intrusive_ptr; +constexpr StringData DocumentSourceGeoNear::kKeyFieldName; +constexpr const char* DocumentSourceGeoNear::kStageName; + REGISTER_DOCUMENT_SOURCE(geoNear, LiteParsedDocumentSourceDefault::parse, DocumentSourceGeoNear::createFromBson); -const long long DocumentSourceGeoNear::kDefaultLimit = 100; - -constexpr StringData DocumentSourceGeoNear::kKeyFieldName; - -const char* DocumentSourceGeoNear::getSourceName() const { - return "$geoNear"; -} - -DocumentSource::GetNextResult DocumentSourceGeoNear::getNext() { - pExpCtx->checkForInterrupt(); - - if (!resultsIterator) - runCommand(); - - if (!resultsIterator->more()) - return GetNextResult::makeEOF(); - - // Each result from the geoNear command is wrapped in a wrapper object with "obj", - // "dis" and maybe "loc" fields. We want to take the object from "obj" and inject the - // other fields into it. - Document result(resultsIterator->next().embeddedObject()); - MutableDocument output(result["obj"].getDocument()); - output.setNestedField(*distanceField, result["dis"]); - if (includeLocs) - output.setNestedField(*includeLocs, result["loc"]); - - // In a cluster, $geoNear output will be merged via $sort, so add the sort key. - if (pExpCtx->needsMerge) { - output.setSortKeyMetaField(BSON("" << result["dis"])); - } - - return output.freeze(); -} - -Pipeline::SourceContainer::iterator DocumentSourceGeoNear::doOptimizeAt( - Pipeline::SourceContainer::iterator itr, Pipeline::SourceContainer* container) { - invariant(*itr == this); - - auto nextLimit = dynamic_cast<DocumentSourceLimit*>((*std::next(itr)).get()); - - if (nextLimit) { - // If the next stage is a $limit, we can combine it with ourselves. - limit = std::min(limit, nextLimit->getLimit()); - container->erase(std::next(itr)); - return itr; - } - return std::next(itr); -} - -// This command is sent as-is to the shards. -intrusive_ptr<DocumentSource> DocumentSourceGeoNear::getShardSource() { - return this; -} -// On mongoS this becomes a merge sort by distance (nearest-first) with limit. -std::list<intrusive_ptr<DocumentSource>> DocumentSourceGeoNear::getMergeSources() { - return {DocumentSourceSort::create( - pExpCtx, BSON(distanceField->fullPath() << 1 << "$mergePresorted" << true), limit)}; -} - Value DocumentSourceGeoNear::serialize(boost::optional<ExplainOptions::Verbosity> explain) const { MutableDocument result; - if (!keyFieldPath.empty()) { - result.setField(kKeyFieldName, Value(keyFieldPath)); + if (keyFieldPath) { + result.setField(kKeyFieldName, Value(keyFieldPath->fullPath())); } if (coordsIsArray) { @@ -119,20 +61,21 @@ Value DocumentSourceGeoNear::serialize(boost::optional<ExplainOptions::Verbosity result.setField("near", Value(coords)); } - // not in buildGeoNearCmd result.setField("distanceField", Value(distanceField->fullPath())); - result.setField("limit", Value(limit)); - - if (maxDistance > 0) - result.setField("maxDistance", Value(maxDistance)); + if (maxDistance) { + result.setField("maxDistance", Value(*maxDistance)); + } - if (minDistance > 0) - result.setField("minDistance", Value(minDistance)); + if (minDistance) { + result.setField("minDistance", Value(*minDistance)); + } result.setField("query", Value(query)); result.setField("spherical", Value(spherical)); - result.setField("distanceMultiplier", Value(distanceMultiplier)); + if (distanceMultiplier) { + result.setField("distanceMultiplier", Value(*distanceMultiplier)); + } if (includeLocs) result.setField("includeLocs", Value(includeLocs->fullPath())); @@ -140,60 +83,6 @@ Value DocumentSourceGeoNear::serialize(boost::optional<ExplainOptions::Verbosity return Value(DOC(getSourceName() << result.freeze())); } -BSONObj DocumentSourceGeoNear::buildGeoNearCmd() const { - // this is very similar to sourceToBson, but slightly different. - // differences will be noted. - - BSONObjBuilder geoNear; // not building a subField - - geoNear.append("geoNear", pExpCtx->ns.coll()); // not in toBson - - if (coordsIsArray) { - geoNear.appendArray("near", coords); - } else { - geoNear.append("near", coords); - } - - geoNear.append("num", limit); // called limit in toBson - - if (maxDistance > 0) - geoNear.append("maxDistance", maxDistance); - - if (minDistance > 0) - geoNear.append("minDistance", minDistance); - - geoNear.append("query", query); - if (pExpCtx->getCollator()) { - geoNear.append("collation", pExpCtx->getCollator()->getSpec().toBSON()); - } else { - geoNear.append("collation", CollationSpec::kSimpleSpec); - } - - geoNear.append("spherical", spherical); - geoNear.append("distanceMultiplier", distanceMultiplier); - - if (includeLocs) - geoNear.append("includeLocs", true); // String in toBson - - if (!keyFieldPath.empty()) { - geoNear.append(kKeyFieldName, keyFieldPath); - } - - return geoNear.obj(); -} - -void DocumentSourceGeoNear::runCommand() { - massert(16603, "Already ran geoNearCommand", !resultsIterator); - - bool ok = pExpCtx->mongoProcessInterface->directClient()->runCommand( - pExpCtx->ns.db().toString(), buildGeoNearCmd(), cmdOutput); - if (!ok) { - uassertStatusOK(getStatusFromCommandResult(cmdOutput)); - } - - resultsIterator.reset(new BSONObjIterator(cmdOutput["results"].embeddedObject())); -} - intrusive_ptr<DocumentSourceGeoNear> DocumentSourceGeoNear::create( const intrusive_ptr<ExpressionContext>& pCtx) { intrusive_ptr<DocumentSourceGeoNear> source(new DocumentSourceGeoNear(pCtx)); @@ -208,8 +97,25 @@ intrusive_ptr<DocumentSource> DocumentSourceGeoNear::createFromBson( } void DocumentSourceGeoNear::parseOptions(BSONObj options) { - // near and distanceField are required + // First, check for explicitly-disallowed fields. + // The old geoNear command used to accept a collation. We explicitly ban it here, since the + // $geoNear stage should respect the collation associated with the entire pipeline. + uassert(40227, + "$geoNear does not accept the 'collation' parameter. Instead, specify a collation " + "for the entire aggregation command.", + !options["collation"]); + + // The following fields were present in older versions but are no longer supported. + uassert(50858, + "$geoNear no longer supports the 'limit' parameter. Use a $limit stage instead.", + !options["limit"]); + uassert(50857, + "$geoNear no longer supports the 'num' parameter. Use a $limit stage instead.", + !options["num"]); + uassert(50856, "$geoNear no longer supports the 'start' argument.", !options["start"]); + + // The "near" and "distanceField" parameters are required. uassert(16605, "$geoNear requires a 'near' option as an Array", options["near"].isABSONObj()); // Array or Object (Object is deprecated) @@ -221,33 +127,47 @@ void DocumentSourceGeoNear::parseOptions(BSONObj options) { options["distanceField"].type() == String); distanceField.reset(new FieldPath(options["distanceField"].str())); - // remaining fields are optional - - // num and limit are synonyms - if (options["limit"].isNumber()) - limit = options["limit"].numberLong(); - if (options["num"].isNumber()) - limit = options["num"].numberLong(); - - if (options["maxDistance"].isNumber()) + // The remaining fields are optional. + if (auto maxDistElem = options["maxDistance"]) { + uassert(ErrorCodes::TypeMismatch, + "maxDistance must be a number", + isNumericBSONType(maxDistElem.type())); maxDistance = options["maxDistance"].numberDouble(); + uassert(ErrorCodes::BadValue, "maxDistance must be nonnegative", *maxDistance >= 0); + } - if (options["minDistance"].isNumber()) + if (auto minDistElem = options["minDistance"]) { + uassert(ErrorCodes::TypeMismatch, + "minDistance must be a number", + isNumericBSONType(minDistElem.type())); minDistance = options["minDistance"].numberDouble(); + uassert(ErrorCodes::BadValue, "minDistance must be nonnegative", *minDistance >= 0); + } - if (options["query"].type() == Object) - query = options["query"].embeddedObject().getOwned(); + if (auto distMultElem = options["distanceMultiplier"]) { + uassert(ErrorCodes::TypeMismatch, + "distanceMultiplier must be a number", + isNumericBSONType(distMultElem.type())); + distanceMultiplier = options["distanceMultiplier"].numberDouble(); + uassert(ErrorCodes::BadValue, + "distanceMultiplier must be nonnegative", + *distanceMultiplier >= 0); + } - spherical = options["spherical"].trueValue(); + if (auto queryElem = options["query"]) { + uassert(ErrorCodes::TypeMismatch, + "query must be an object", + queryElem.type() == BSONType::Object); + query = queryElem.embeddedObject().getOwned(); + } - if (options["distanceMultiplier"].isNumber()) - distanceMultiplier = options["distanceMultiplier"].numberDouble(); + spherical = options["spherical"].trueValue(); if (options.hasField("includeLocs")) { uassert(16607, "$geoNear requires that 'includeLocs' option is a String", options["includeLocs"].type() == String); - includeLocs.reset(new FieldPath(options["includeLocs"].str())); + includeLocs = FieldPath(options["includeLocs"].str()); } if (options.hasField("uniqueDocs")) @@ -259,27 +179,64 @@ void DocumentSourceGeoNear::parseOptions(BSONObj options) { << "' must be of type string but found type: " << typeName(keyElt.type()), keyElt.type() == BSONType::String); - keyFieldPath = keyElt.str(); + const auto keyFieldStr = keyElt.valueStringData(); uassert(ErrorCodes::BadValue, str::stream() << "$geoNear parameter '" << DocumentSourceGeoNear::kKeyFieldName << "' cannot be the empty string", - !keyFieldPath.empty()); + !keyFieldStr.empty()); + keyFieldPath = FieldPath(keyFieldStr); } +} - // The collation field is disallowed, even though it is accepted by the geoNear command, since - // the $geoNear operation should respect the collation associated with the entire pipeline. - uassert(40227, - "$geoNear does not accept the 'collation' parameter. Instead, specify a collation " - "for the entire aggregation command.", - !options["collation"]); +BSONObj DocumentSourceGeoNear::asNearQuery(StringData nearFieldName) const { + BSONObjBuilder queryBuilder; + queryBuilder.appendElements(query); + + BSONObjBuilder nearBuilder(queryBuilder.subobjStart(nearFieldName)); + if (spherical) { + if (coordsIsArray) { + nearBuilder.appendArray("$nearSphere", coords); + } else { + nearBuilder.append("$nearSphere", coords); + } + } else { + if (coordsIsArray) { + nearBuilder.appendArray("$near", coords); + } else { + nearBuilder.append("$near", coords); + } + } + if (minDistance) { + nearBuilder.append("$minDistance", *minDistance); + } + if (maxDistance) { + nearBuilder.append("$maxDistance", *maxDistance); + } + nearBuilder.doneFast(); + return queryBuilder.obj(); +} + +bool DocumentSourceGeoNear::needsGeoNearPoint() const { + return static_cast<bool>(includeLocs); +} + +DocumentSource::GetDepsReturn DocumentSourceGeoNear::getDependencies(DepsTracker* deps) const { + // TODO (SERVER-35424): Implement better dependency tracking. For example, 'distanceField' is + // produced by this stage, and we could inform the query system that it need not include it in + // its response. For now, assume that we require the entire document as well as the appropriate + // geoNear metadata. + deps->setNeedsMetadata(DepsTracker::MetadataType::GEO_NEAR_DISTANCE, true); + deps->setNeedsMetadata(DepsTracker::MetadataType::GEO_NEAR_POINT, needsGeoNearPoint()); + + deps->needWholeDocument = true; + return GetDepsReturn::EXHAUSTIVE_FIELDS; } DocumentSourceGeoNear::DocumentSourceGeoNear(const intrusive_ptr<ExpressionContext>& pExpCtx) - : DocumentSource(pExpCtx), - coordsIsArray(false), - limit(DocumentSourceGeoNear::kDefaultLimit), - maxDistance(-1.0), - minDistance(-1.0), - spherical(false), - distanceMultiplier(1.0) {} + : DocumentSource(pExpCtx), coordsIsArray(false), spherical(false) {} + +std::list<boost::intrusive_ptr<DocumentSource>> DocumentSourceGeoNear::getMergeSources() { + return {DocumentSourceSort::create( + pExpCtx, BSON(distanceField->fullPath() << 1 << "$mergePresorted" << true))}; } +} // namespace mongo diff --git a/src/mongo/db/pipeline/document_source_geo_near.h b/src/mongo/db/pipeline/document_source_geo_near.h index fa67ab1c71a..7d77fa5a422 100644 --- a/src/mongo/db/pipeline/document_source_geo_near.h +++ b/src/mongo/db/pipeline/document_source_geo_near.h @@ -35,86 +35,120 @@ namespace mongo { class DocumentSourceGeoNear : public DocumentSource, public NeedsMergerDocumentSource { public: - static const long long kDefaultLimit; - static constexpr StringData kKeyFieldName = "key"_sd; + static constexpr auto kStageName = "$geoNear"; - // virtuals from DocumentSource - GetNextResult getNext() final; - const char* getSourceName() const final; /** - * Attempts to combine with a subsequent limit stage, setting the internal limit field - * as a result. + * Only exposed for testing. */ - Pipeline::SourceContainer::iterator doOptimizeAt(Pipeline::SourceContainer::iterator itr, - Pipeline::SourceContainer* container) final; + static boost::intrusive_ptr<DocumentSourceGeoNear> create( + const boost::intrusive_ptr<ExpressionContext>&); + + const char* getSourceName() const final { + return kStageName; + } StageConstraints constraints(Pipeline::SplitState pipeState) const final { - StageConstraints constraints(StreamType::kStreaming, - PositionRequirement::kFirst, - HostTypeRequirement::kAnyShard, - DiskUseRequirement::kNoDiskUse, - FacetRequirement::kNotAllowed, - TransactionRequirement::kAllowed); - - constraints.requiresInputDocSource = false; - return constraints; + return {StreamType::kStreaming, + PositionRequirement::kFirst, + HostTypeRequirement::kAnyShard, + DiskUseRequirement::kNoDiskUse, + FacetRequirement::kNotAllowed, + TransactionRequirement::kAllowed}; } - Value serialize(boost::optional<ExplainOptions::Verbosity> explain = boost::none) const final; - BSONObjSet getOutputSorts() final { - return SimpleBSONObjComparator::kInstance.makeBSONObjSet( - {BSON(distanceField->fullPath() << -1)}); + /** + * DocumentSourceGeoNear should always be replaced by a DocumentSourceGeoNearCursor before + * executing a pipeline, so this method should never be called. + */ + GetNextResult getNext() final { + MONGO_UNREACHABLE; } - // Virtuals for NeedsMergerDocumentSource - boost::intrusive_ptr<DocumentSource> getShardSource() final; - std::list<boost::intrusive_ptr<DocumentSource>> getMergeSources() final; + Value serialize(boost::optional<ExplainOptions::Verbosity> explain = boost::none) const final; static boost::intrusive_ptr<DocumentSource> createFromBson( BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& pCtx); - static char geoNearName[]; - - long long getLimit() { - return limit; - } - + /** + * A query predicate to apply to the documents in addition to the "near" predicate. + */ BSONObj getQuery() const { return query; }; - // this should only be used for testing - static boost::intrusive_ptr<DocumentSourceGeoNear> create( - const boost::intrusive_ptr<ExpressionContext>& pCtx); + /** + * The field in which the computed distance will be stored. + */ + FieldPath getDistanceField() const { + return *distanceField; + } + + /** + * The field in which the matching point will be stored, if requested. + */ + boost::optional<FieldPath> getLocationField() const { + return includeLocs; + } + + /** + * The field over which to apply the "near" predicate, if specified. + */ + boost::optional<FieldPath> getKeyField() const { + return keyFieldPath; + } + + /** + * A scaling factor to apply to the distance, if specified by the user. + */ + boost::optional<double> getDistanceMultiplier() const { + return distanceMultiplier; + } + + GetDepsReturn getDependencies(DepsTracker* deps) const final; + + /** + * Returns true if the $geoNear specification requires the geoNear point metadata. + */ + bool needsGeoNearPoint() const; + + /** + * Converts this $geoNear aggregation stage into an equivalent $near or $nearSphere query on + * 'nearFieldName'. + */ + BSONObj asNearQuery(StringData nearFieldName) const; + + /** + * This document source is sent as-is to the shards. + */ + boost::intrusive_ptr<DocumentSource> getShardSource() final { + return this; + } + + /** + * In a sharded cluster, this becomes a merge sort by distance, from nearest to furthest. + */ + std::list<boost::intrusive_ptr<DocumentSource>> getMergeSources() final; private: explicit DocumentSourceGeoNear(const boost::intrusive_ptr<ExpressionContext>& pExpCtx); + /** + * Parses the fields in the object 'options', throwing if an error occurs. + */ void parseOptions(BSONObj options); - BSONObj buildGeoNearCmd() const; - void runCommand(); // These fields describe the command to run. - // coords and distanceField are required, rest are optional + // 'coords' and 'distanceField' are required; the rest are optional. BSONObj coords; // "near" option, but near is a reserved keyword on windows bool coordsIsArray; std::unique_ptr<FieldPath> distanceField; // Using unique_ptr because FieldPath can't be empty - long long limit; - double maxDistance; - double minDistance; BSONObj query; bool spherical; - double distanceMultiplier; - std::unique_ptr<FieldPath> includeLocs; - - // The field path over which the command should run, extracted from the 'key' parameter passed - // by the user. Or the empty string the user did not provide a 'key'. - std::string keyFieldPath; - - // these fields are used while processing the results - BSONObj cmdOutput; - std::unique_ptr<BSONObjIterator> resultsIterator; // iterator over cmdOutput["results"] + boost::optional<double> maxDistance; + boost::optional<double> minDistance; + boost::optional<double> distanceMultiplier; + boost::optional<FieldPath> includeLocs; + boost::optional<FieldPath> keyFieldPath; }; - } // namespace mongo diff --git a/src/mongo/db/pipeline/document_source_geo_near_cursor.cpp b/src/mongo/db/pipeline/document_source_geo_near_cursor.cpp new file mode 100644 index 00000000000..4b1a9596fd9 --- /dev/null +++ b/src/mongo/db/pipeline/document_source_geo_near_cursor.cpp @@ -0,0 +1,118 @@ +/** + * Copyright (C) 2018 MongoDB 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/>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the GNU Affero General Public License in all respects + * for all of the code used other than as permitted herein. If you modify + * file(s) with this exception, you may extend this exception to your + * version of the file(s), but you are not obligated to do so. If you do not + * wish to do so, delete this exception statement from your version. If you + * delete this exception statement from all source files in the program, + * then also delete it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/pipeline/document_source_geo_near_cursor.h" + +#include <boost/intrusive_ptr.hpp> +#include <boost/optional.hpp> +#include <list> +#include <memory> + +#include "mongo/base/string_data.h" +#include "mongo/bson/bsonelement.h" +#include "mongo/bson/bsonobj.h" +#include "mongo/bson/simple_bsonobj_comparator.h" +#include "mongo/db/catalog/collection.h" +#include "mongo/db/pipeline/document.h" +#include "mongo/db/pipeline/document_source_cursor.h" +#include "mongo/db/pipeline/document_source_sort.h" +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/pipeline/field_path.h" +#include "mongo/db/query/plan_executor.h" + +namespace mongo { +constexpr const char* DocumentSourceGeoNearCursor::kStageName; + +boost::intrusive_ptr<DocumentSourceGeoNearCursor> DocumentSourceGeoNearCursor::create( + Collection* collection, + std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + FieldPath distanceField, + boost::optional<FieldPath> locationField, + double distanceMultiplier) { + return {new DocumentSourceGeoNearCursor(collection, + std::move(exec), + expCtx, + std::move(distanceField), + std::move(locationField), + distanceMultiplier)}; +} + +DocumentSourceGeoNearCursor::DocumentSourceGeoNearCursor( + Collection* collection, + std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + FieldPath distanceField, + boost::optional<FieldPath> locationField, + double distanceMultiplier) + : DocumentSourceCursor(collection, std::move(exec), expCtx), + _distanceField(std::move(distanceField)), + _locationField(std::move(locationField)), + _distanceMultiplier(distanceMultiplier) { + invariant(_distanceMultiplier >= 0); +} + +const char* DocumentSourceGeoNearCursor::getSourceName() const { + return kStageName; +} + +BSONObjSet DocumentSourceGeoNearCursor::getOutputSorts() { + return SimpleBSONObjComparator::kInstance.makeBSONObjSet( + {BSON(_distanceField.fullPath() << 1)}); +} + +Document DocumentSourceGeoNearCursor::transformBSONObjToDocument(const BSONObj& obj) const { + MutableDocument output(Document::fromBsonWithMetaData(obj)); + + // Scale the distance by the requested factor. + invariant(output.peek().hasGeoNearDistance(), + str::stream() + << "Query returned a document that is unexpectedly missing the geoNear distance: " + << obj.jsonString()); + const auto distance = output.peek().getGeoNearDistance() * _distanceMultiplier; + + output.setNestedField(_distanceField, Value(distance)); + if (_locationField) { + invariant( + output.peek().hasGeoNearPoint(), + str::stream() + << "Query returned a document that is unexpectedly missing the geoNear point: " + << obj.jsonString()); + output.setNestedField(*_locationField, output.peek().getGeoNearPoint()); + } + + // In a cluster, $geoNear will be merged via $sort, so add the sort key. + if (pExpCtx->needsMerge) { + output.setSortKeyMetaField(BSON("" << distance)); + } + + return output.freeze(); +} +} // namespace mongo diff --git a/src/mongo/db/pipeline/document_source_geo_near_cursor.h b/src/mongo/db/pipeline/document_source_geo_near_cursor.h new file mode 100644 index 00000000000..1e95ba5daa6 --- /dev/null +++ b/src/mongo/db/pipeline/document_source_geo_near_cursor.h @@ -0,0 +1,102 @@ +/** + * Copyright (C) 2018 MongoDB 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/>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the GNU Affero General Public License in all respects + * for all of the code used other than as permitted herein. If you modify + * file(s) with this exception, you may extend this exception to your + * version of the file(s), but you are not obligated to do so. If you do not + * wish to do so, delete this exception statement from your version. If you + * delete this exception statement from all source files in the program, + * then also delete it in the license file. + */ + +#pragma once + +#include <boost/intrusive_ptr.hpp> +#include <boost/optional.hpp> +#include <list> +#include <memory> + +#include "mongo/base/string_data.h" +#include "mongo/bson/bsonobj.h" +#include "mongo/db/catalog/collection.h" +#include "mongo/db/pipeline/document.h" +#include "mongo/db/pipeline/document_source_cursor.h" +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/pipeline/field_path.h" +#include "mongo/db/query/plan_executor.h" + +namespace mongo { +/** + * Like DocumentSourceCursor, this stage returns Documents from BSONObjs produced by a PlanExecutor, + * but does extra work to compute distances to satisfy a $near or $nearSphere query. + */ +class DocumentSourceGeoNearCursor final : public DocumentSourceCursor { +public: + /** + * The name of this stage. + */ + static constexpr auto kStageName = "$geoNearCursor"; + + /** + * Create a new DocumentSourceGeoNearCursor. If specified, 'distanceMultiplier' must be + * nonnegative. + */ + static boost::intrusive_ptr<DocumentSourceGeoNearCursor> create( + Collection*, + std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>, + const boost::intrusive_ptr<ExpressionContext>&, + FieldPath distanceField, + boost::optional<FieldPath> locationField = boost::none, + double distanceMultiplier = 1.0); + + const char* getSourceName() const final; + + /** + * $geoNear returns documents ordered from nearest to furthest, which is an ascending sort on + * '_distanceField'. + */ + BSONObjSet getOutputSorts() final; + +private: + DocumentSourceGeoNearCursor(Collection*, + std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>, + const boost::intrusive_ptr<ExpressionContext>&, + FieldPath distanceField, + boost::optional<FieldPath> locationField, + double distanceMultiplier); + + ~DocumentSourceGeoNearCursor() = default; + + /** + * Transforms 'obj' into a Document, calculating the distance. + */ + Document transformBSONObjToDocument(const BSONObj& obj) const final; + + // The output field in which to store the computed distance. + FieldPath _distanceField; + + // The output field to store the point that matched, if specified. + boost::optional<FieldPath> _locationField; + + // A multiplicative factor applied to each distance. For example, you can use this to convert + // radian distances into meters by multiplying by the radius of the Earth. + double _distanceMultiplier = 1.0; +}; +} // namespace mongo diff --git a/src/mongo/db/pipeline/document_source_geo_near_test.cpp b/src/mongo/db/pipeline/document_source_geo_near_test.cpp index e1fdbc6d5e6..f6fe24096f5 100644 --- a/src/mongo/db/pipeline/document_source_geo_near_test.cpp +++ b/src/mongo/db/pipeline/document_source_geo_near_test.cpp @@ -44,45 +44,6 @@ namespace { // This provides access to getExpCtx(), but we'll use a different name for this test suite. using DocumentSourceGeoNearTest = AggregationContextFixture; -TEST_F(DocumentSourceGeoNearTest, ShouldAbsorbSubsequentLimitStage) { - auto geoNear = DocumentSourceGeoNear::create(getExpCtx()); - - Pipeline::SourceContainer container; - container.push_back(geoNear); - - ASSERT_EQUALS(geoNear->getLimit(), DocumentSourceGeoNear::kDefaultLimit); - - container.push_back(DocumentSourceLimit::create(getExpCtx(), 200)); - geoNear->optimizeAt(container.begin(), &container); - - ASSERT_EQUALS(container.size(), 1U); - ASSERT_EQUALS(geoNear->getLimit(), DocumentSourceGeoNear::kDefaultLimit); - - container.push_back(DocumentSourceLimit::create(getExpCtx(), 50)); - geoNear->optimizeAt(container.begin(), &container); - - ASSERT_EQUALS(container.size(), 1U); - ASSERT_EQUALS(geoNear->getLimit(), 50); - - container.push_back(DocumentSourceLimit::create(getExpCtx(), 30)); - geoNear->optimizeAt(container.begin(), &container); - - ASSERT_EQUALS(container.size(), 1U); - ASSERT_EQUALS(geoNear->getLimit(), 30); -} - -TEST_F(DocumentSourceGeoNearTest, ShouldReportOutputsAreSortedByDistanceField) { - BSONObj queryObj = fromjson( - "{geoNear: { near: {type: 'Point', coordinates: [0, 0]}, distanceField: 'dist', " - "maxDistance: 2}}"); - auto geoNear = DocumentSourceGeoNear::createFromBson(queryObj.firstElement(), getExpCtx()); - - BSONObjSet outputSort = geoNear->getOutputSorts(); - - ASSERT_EQUALS(outputSort.count(BSON("dist" << -1)), 1U); - ASSERT_EQUALS(outputSort.size(), 1U); -} - TEST_F(DocumentSourceGeoNearTest, FailToParseIfKeyFieldNotAString) { auto stageObj = fromjson("{$geoNear: {distanceField: 'dist', near: [0, 0], key: 1}}"); ASSERT_THROWS_CODE(DocumentSourceGeoNear::createFromBson(stageObj.firstElement(), getExpCtx()), @@ -97,6 +58,35 @@ TEST_F(DocumentSourceGeoNearTest, FailToParseIfKeyIsTheEmptyString) { ErrorCodes::BadValue); } +TEST_F(DocumentSourceGeoNearTest, FailToParseIfDistanceMultiplierIsNegative) { + auto stageObj = + fromjson("{$geoNear: {distanceField: 'dist', near: [0, 0], distanceMultiplier: -1.0}}"); + ASSERT_THROWS_CODE(DocumentSourceGeoNear::createFromBson(stageObj.firstElement(), getExpCtx()), + AssertionException, + ErrorCodes::BadValue); +} + +TEST_F(DocumentSourceGeoNearTest, FailToParseIfLimitFieldSpecified) { + auto stageObj = fromjson("{$geoNear: {distanceField: 'dist', near: [0, 0], limit: 1}}"); + ASSERT_THROWS_CODE(DocumentSourceGeoNear::createFromBson(stageObj.firstElement(), getExpCtx()), + AssertionException, + 50858); +} + +TEST_F(DocumentSourceGeoNearTest, FailToParseIfNumFieldSpecified) { + auto stageObj = fromjson("{$geoNear: {distanceField: 'dist', near: [0, 0], num: 1}}"); + ASSERT_THROWS_CODE(DocumentSourceGeoNear::createFromBson(stageObj.firstElement(), getExpCtx()), + AssertionException, + 50857); +} + +TEST_F(DocumentSourceGeoNearTest, FailToParseIfStartOptionIsSpecified) { + auto stageObj = fromjson("{$geoNear: {distanceField: 'dist', near: [0, 0], start: {}}}"); + ASSERT_THROWS_CODE(DocumentSourceGeoNear::createFromBson(stageObj.firstElement(), getExpCtx()), + AssertionException, + 50856); +} + TEST_F(DocumentSourceGeoNearTest, CanParseAndSerializeKeyField) { auto stageObj = fromjson("{$geoNear: {distanceField: 'dist', near: [0, 0], key: 'a.b'}}"); auto geoNear = DocumentSourceGeoNear::createFromBson(stageObj.firstElement(), getExpCtx()); @@ -108,12 +98,9 @@ TEST_F(DocumentSourceGeoNearTest, CanParseAndSerializeKeyField) { Value{Document{{"key", "a.b"_sd}, {"near", std::vector<Value>{Value{0}, Value{0}}}, {"distanceField", "dist"_sd}, - {"limit", 100}, {"query", BSONObj()}, - {"spherical", false}, - {"distanceMultiplier", 1}}}}}}; + {"spherical", false}}}}}}; ASSERT_VALUE_EQ(expectedSerialization, serialized[0]); } - } // namespace } // namespace mongo diff --git a/src/mongo/db/pipeline/document_source_group_test.cpp b/src/mongo/db/pipeline/document_source_group_test.cpp index de99452f3d7..8d0ce4a6679 100644 --- a/src/mongo/db/pipeline/document_source_group_test.cpp +++ b/src/mongo/db/pipeline/document_source_group_test.cpp @@ -767,7 +767,7 @@ public: ASSERT_EQUALS(1U, dependencies.fields.count("u")); ASSERT_EQUALS(1U, dependencies.fields.count("v")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } }; diff --git a/src/mongo/db/pipeline/document_source_limit_test.cpp b/src/mongo/db/pipeline/document_source_limit_test.cpp index 6294a897ff2..b776de9789b 100644 --- a/src/mongo/db/pipeline/document_source_limit_test.cpp +++ b/src/mongo/db/pipeline/document_source_limit_test.cpp @@ -98,7 +98,7 @@ TEST_F(DocumentSourceLimitTest, ShouldNotIntroduceAnyDependencies) { ASSERT_EQUALS(DocumentSource::SEE_NEXT, limit->getDependencies(&dependencies)); ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceLimitTest, ShouldPropagatePauses) { diff --git a/src/mongo/db/pipeline/document_source_match.cpp b/src/mongo/db/pipeline/document_source_match.cpp index 47e85a4c6e7..45bc8c8629e 100644 --- a/src/mongo/db/pipeline/document_source_match.cpp +++ b/src/mongo/db/pipeline/document_source_match.cpp @@ -485,7 +485,7 @@ DocumentSource::GetDepsReturn DocumentSourceMatch::getDependencies(DepsTracker* // A $text aggregation field should return EXHAUSTIVE_FIELDS, since we don't necessarily // know what field it will be searching without examining indices. deps->needWholeDocument = true; - deps->setNeedTextScore(true); + deps->setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); return EXHAUSTIVE_FIELDS; } diff --git a/src/mongo/db/pipeline/document_source_match_test.cpp b/src/mongo/db/pipeline/document_source_match_test.cpp index 4e14a990c7a..5f81dc71986 100644 --- a/src/mongo/db/pipeline/document_source_match_test.cpp +++ b/src/mongo/db/pipeline/document_source_match_test.cpp @@ -219,7 +219,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddDependenciesOfAllBranchesOfOrClause) { ASSERT_EQUALS(1U, dependencies.fields.count("x.y")); ASSERT_EQUALS(2U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, TextSearchShouldRequireWholeDocumentAndTextScore) { @@ -227,7 +227,7 @@ TEST_F(DocumentSourceMatchTest, TextSearchShouldRequireWholeDocumentAndTextScore DepsTracker dependencies(DepsTracker::MetadataAvailable::kTextScore); ASSERT_EQUALS(DocumentSource::EXHAUSTIVE_FIELDS, match->getDependencies(&dependencies)); ASSERT_EQUALS(true, dependencies.needWholeDocument); - ASSERT_EQUALS(true, dependencies.getNeedTextScore()); + ASSERT_EQUALS(true, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldOnlyAddOuterFieldAsDependencyOfImplicitEqualityPredicate) { @@ -238,7 +238,7 @@ TEST_F(DocumentSourceMatchTest, ShouldOnlyAddOuterFieldAsDependencyOfImplicitEqu ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldOnlyAddOuterFieldAsDependencyOfClausesWithinElemMatch) { @@ -249,7 +249,7 @@ TEST_F(DocumentSourceMatchTest, ShouldOnlyAddOuterFieldAsDependencyOfClausesWith ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, @@ -266,7 +266,7 @@ TEST_F(DocumentSourceMatchTest, ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, @@ -277,7 +277,7 @@ TEST_F(DocumentSourceMatchTest, ASSERT_EQUALS(DocumentSource::SEE_NEXT, match->getDependencies(&dependencies)); ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(true, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, @@ -288,7 +288,7 @@ TEST_F(DocumentSourceMatchTest, ASSERT_EQUALS(DocumentSource::SEE_NEXT, match->getDependencies(&dependencies1)); ASSERT_EQUALS(0U, dependencies1.fields.size()); ASSERT_EQUALS(true, dependencies1.needWholeDocument); - ASSERT_EQUALS(false, dependencies1.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies1.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); query = fromjson("{a: {$_internalSchemaObjectMatch: {$_internalSchemaMaxProperties: 1}}}"); match = DocumentSourceMatch::create(query, getExpCtx()); @@ -297,7 +297,7 @@ TEST_F(DocumentSourceMatchTest, ASSERT_EQUALS(1U, dependencies2.fields.size()); ASSERT_EQUALS(1U, dependencies2.fields.count("a")); ASSERT_EQUALS(false, dependencies2.needWholeDocument); - ASSERT_EQUALS(false, dependencies2.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies2.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, @@ -310,7 +310,7 @@ TEST_F(DocumentSourceMatchTest, ASSERT_EQUALS(DocumentSource::SEE_NEXT, match->getDependencies(&dependencies)); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(true, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, @@ -321,7 +321,7 @@ TEST_F(DocumentSourceMatchTest, ASSERT_EQUALS(DocumentSource::SEE_NEXT, match->getDependencies(&dependencies)); ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(true, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithInternalSchemaType) { @@ -332,7 +332,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithIntern ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithInternalSchemaCond) { @@ -345,7 +345,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithIntern ASSERT_EQUALS(1U, dependencies.fields.count("b")); ASSERT_EQUALS(1U, dependencies.fields.count("c")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithInternalSchemaXor) { @@ -358,7 +358,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithIntern ASSERT_EQUALS(1U, dependencies.fields.count("b")); ASSERT_EQUALS(1U, dependencies.fields.count("c")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithEmptyJSONSchema) { @@ -368,7 +368,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithEmptyJ ASSERT_EQUALS(DocumentSource::SEE_NEXT, match->getDependencies(&dependencies)); ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithJSONSchemaProperties) { @@ -379,7 +379,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForClausesWithJSONSc ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForMultiplePredicatesWithJSONSchema) { @@ -391,7 +391,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddCorrectDependenciesForMultiplePredicate ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.count("b")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddOuterFieldToDependenciesIfElemMatchContainsNoFieldNames) { @@ -402,7 +402,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddOuterFieldToDependenciesIfElemMatchCont ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddNotClausesFieldAsDependency) { @@ -412,7 +412,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddNotClausesFieldAsDependency) { ASSERT_EQUALS(1U, dependencies.fields.count("b")); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ShouldAddDependenciesOfEachNorClause) { @@ -424,7 +424,7 @@ TEST_F(DocumentSourceMatchTest, ShouldAddDependenciesOfEachNorClause) { ASSERT_EQUALS(1U, dependencies.fields.count("b.c")); ASSERT_EQUALS(2U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, CommentShouldNotAddAnyDependencies) { @@ -433,7 +433,7 @@ TEST_F(DocumentSourceMatchTest, CommentShouldNotAddAnyDependencies) { ASSERT_EQUALS(DocumentSource::SEE_NEXT, match->getDependencies(&dependencies)); ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, ClauseAndedWithCommentShouldAddDependencies) { @@ -444,7 +444,7 @@ TEST_F(DocumentSourceMatchTest, ClauseAndedWithCommentShouldAddDependencies) { ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceMatchTest, MultipleMatchStagesShouldCombineIntoOne) { diff --git a/src/mongo/db/pipeline/document_source_project_test.cpp b/src/mongo/db/pipeline/document_source_project_test.cpp index 5688d343675..c991eebbbc4 100644 --- a/src/mongo/db/pipeline/document_source_project_test.cpp +++ b/src/mongo/db/pipeline/document_source_project_test.cpp @@ -181,7 +181,7 @@ TEST_F(ProjectStageTest, InclusionShouldAddDependenciesOfIncludedAndComputedFiel ASSERT_EQUALS(1U, dependencies.fields.count("c")); ASSERT_EQUALS(1U, dependencies.fields.count("d")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(true, dependencies.getNeedTextScore()); + ASSERT_EQUALS(true, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(ProjectStageTest, ExclusionShouldNotAddDependencies) { @@ -192,7 +192,7 @@ TEST_F(ProjectStageTest, ExclusionShouldNotAddDependencies) { ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(ProjectStageTest, InclusionProjectionReportsIncludedPathsFromGetModifiedPaths) { diff --git a/src/mongo/db/pipeline/document_source_replace_root_test.cpp b/src/mongo/db/pipeline/document_source_replace_root_test.cpp index 321afbc8220..0ac8a9f9db0 100644 --- a/src/mongo/db/pipeline/document_source_replace_root_test.cpp +++ b/src/mongo/db/pipeline/document_source_replace_root_test.cpp @@ -268,7 +268,7 @@ TEST_F(ReplaceRootBasics, OnlyDependentFieldIsNewRoot) { // Should not need any other fields. ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(ReplaceRootBasics, ReplaceRootModifiesAllFields) { diff --git a/src/mongo/db/pipeline/document_source_sort.cpp b/src/mongo/db/pipeline/document_source_sort.cpp index 30ee3a5b2e2..a3374521faf 100644 --- a/src/mongo/db/pipeline/document_source_sort.cpp +++ b/src/mongo/db/pipeline/document_source_sort.cpp @@ -226,7 +226,7 @@ DocumentSource::GetDepsReturn DocumentSourceSort::getDependencies(DepsTracker* d } if (pExpCtx->needsMerge) { // Include the sort key if we will merge several sorted streams later. - deps->setNeedSortKey(true); + deps->setNeedsMetadata(DepsTracker::MetadataType::SORT_KEY, true); } return SEE_NEXT; diff --git a/src/mongo/db/pipeline/document_source_sort_test.cpp b/src/mongo/db/pipeline/document_source_sort_test.cpp index fca4caaaf28..9f1f9606220 100644 --- a/src/mongo/db/pipeline/document_source_sort_test.cpp +++ b/src/mongo/db/pipeline/document_source_sort_test.cpp @@ -169,7 +169,7 @@ TEST_F(DocumentSourceSortTest, Dependencies) { ASSERT_EQUALS(1U, dependencies.fields.count("a")); ASSERT_EQUALS(1U, dependencies.fields.count("b.c")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(DocumentSourceSortTest, OutputSort) { diff --git a/src/mongo/db/pipeline/document_source_unwind_test.cpp b/src/mongo/db/pipeline/document_source_unwind_test.cpp index 8c01c50bde9..af04d436083 100644 --- a/src/mongo/db/pipeline/document_source_unwind_test.cpp +++ b/src/mongo/db/pipeline/document_source_unwind_test.cpp @@ -680,7 +680,7 @@ TEST_F(UnwindStageTest, AddsUnwoundPathToDependencies) { ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(1U, dependencies.fields.count("x.y.z")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(UnwindStageTest, TruncatesOutputSortAtUnwoundPath) { diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp index 9094071c9d0..5cb0122f565 100644 --- a/src/mongo/db/pipeline/expression.cpp +++ b/src/mongo/db/pipeline/expression.cpp @@ -2534,7 +2534,7 @@ Value ExpressionMeta::evaluate(const Document& root) const { void ExpressionMeta::_doAddDependencies(DepsTracker* deps) const { if (_metaType == MetaType::TEXT_SCORE) { - deps->setNeedTextScore(true); + deps->setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); } } diff --git a/src/mongo/db/pipeline/expression_test.cpp b/src/mongo/db/pipeline/expression_test.cpp index 374056289c3..2d7c4da922a 100644 --- a/src/mongo/db/pipeline/expression_test.cpp +++ b/src/mongo/db/pipeline/expression_test.cpp @@ -244,7 +244,7 @@ protected: } ASSERT_BSONOBJ_EQ(expectedDependencies, dependenciesBson.arr()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } void assertContents(const intrusive_ptr<Testable>& expr, const BSONArray& expectedContents) { @@ -1565,7 +1565,7 @@ public: ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(1U, dependencies.fields.count("a.b")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } }; @@ -2071,7 +2071,7 @@ public: expression->addDependencies(&dependencies); ASSERT_EQUALS(0U, dependencies.fields.size()); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } }; @@ -2502,7 +2502,7 @@ public: ASSERT_EQUALS(1U, dependencies.fields.size()); ASSERT_EQUALS(1U, dependencies.fields.count("a.b")); ASSERT_EQUALS(false, dependencies.needWholeDocument); - ASSERT_EQUALS(false, dependencies.getNeedTextScore()); + ASSERT_EQUALS(false, dependencies.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } }; diff --git a/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp b/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp index 348044018e9..33d915a47a9 100644 --- a/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp +++ b/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp @@ -120,7 +120,7 @@ TEST(ExclusionProjection, ShouldNotAddAnyDependencies) { ASSERT_EQ(deps.fields.size(), 0UL); ASSERT_FALSE(deps.needWholeDocument); - ASSERT_FALSE(deps.getNeedTextScore()); + ASSERT_FALSE(deps.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST(ExclusionProjection, ShouldReportExcludedFieldsAsModified) { diff --git a/src/mongo/db/pipeline/pipeline.cpp b/src/mongo/db/pipeline/pipeline.cpp index 79dd962d871..6654b836a3e 100644 --- a/src/mongo/db/pipeline/pipeline.cpp +++ b/src/mongo/db/pipeline/pipeline.cpp @@ -69,6 +69,9 @@ using DiskUseRequirement = DocumentSource::StageConstraints::DiskUseRequirement; using FacetRequirement = DocumentSource::StageConstraints::FacetRequirement; using StreamType = DocumentSource::StageConstraints::StreamType; +constexpr MatchExpressionParser::AllowedFeatureSet Pipeline::kAllowedMatcherFeatures; +constexpr MatchExpressionParser::AllowedFeatureSet Pipeline::kGeoNearMatcherFeatures; + Pipeline::Pipeline(const intrusive_ptr<ExpressionContext>& pTheCtx) : pCtx(pTheCtx) {} Pipeline::Pipeline(SourceContainer stages, const intrusive_ptr<ExpressionContext>& expCtx) @@ -509,12 +512,9 @@ DepsTracker Pipeline::getDependencies(DepsTracker::MetadataAvailable metadataAva } if (!knowAllMeta) { - if (localDeps.getNeedTextScore()) - deps.setNeedTextScore(true); - - if (localDeps.getNeedSortKey()) - deps.setNeedSortKey(true); - + for (auto&& req : localDeps.getAllRequiredMetadataTypes()) { + deps.setNeedsMetadata(req, true); + } knowAllMeta = status & DocumentSource::EXHAUSTIVE_META; } @@ -531,11 +531,12 @@ DepsTracker Pipeline::getDependencies(DepsTracker::MetadataAvailable metadataAva if (metadataAvailable & DepsTracker::MetadataAvailable::kTextScore) { // If there is a text score, assume we need to keep it if we can't prove we don't. If we are // the first half of a pipeline which has been split, future stages might need it. - if (!knowAllMeta) - deps.setNeedTextScore(true); + if (!knowAllMeta) { + deps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); + } } else { // If there is no text score available, then we don't need to ask for it. - deps.setNeedTextScore(false); + deps.setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, false); } return deps; diff --git a/src/mongo/db/pipeline/pipeline.h b/src/mongo/db/pipeline/pipeline.h index 22a8f8a8f88..afcedc14ee6 100644 --- a/src/mongo/db/pipeline/pipeline.h +++ b/src/mongo/db/pipeline/pipeline.h @@ -69,7 +69,7 @@ public: enum class SplitState { kUnsplit, kSplitForShards, kSplitForMerge }; /** - * List of supported match expression features in a pipeline. + * The list of default supported match expression features. */ static constexpr MatchExpressionParser::AllowedFeatureSet kAllowedMatcherFeatures = MatchExpressionParser::AllowedFeatures::kText | @@ -77,6 +77,15 @@ public: MatchExpressionParser::AllowedFeatures::kJSONSchema; /** + * The match expression features allowed when running a pipeline with $geoNear. + */ + static constexpr MatchExpressionParser::AllowedFeatureSet kGeoNearMatcherFeatures = + MatchExpressionParser::AllowedFeatures::kText | + MatchExpressionParser::AllowedFeatures::kExpr | + MatchExpressionParser::AllowedFeatures::kJSONSchema | + MatchExpressionParser::AllowedFeatures::kGeoNear; + + /** * Parses a Pipeline from a vector of BSONObjs. Returns a non-OK status if it failed to parse. * The returned pipeline is not optimized, but the caller may convert it to an optimized * pipeline by calling optimizePipeline(). diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp index 7b5cd468130..9fec47e815d 100644 --- a/src/mongo/db/pipeline/pipeline_d.cpp +++ b/src/mongo/db/pipeline/pipeline_d.cpp @@ -58,6 +58,8 @@ #include "mongo/db/pipeline/document_source.h" #include "mongo/db/pipeline/document_source_change_stream.h" #include "mongo/db/pipeline/document_source_cursor.h" +#include "mongo/db/pipeline/document_source_geo_near.h" +#include "mongo/db/pipeline/document_source_geo_near_cursor.h" #include "mongo/db/pipeline/document_source_match.h" #include "mongo/db/pipeline/document_source_merge_cursors.h" #include "mongo/db/pipeline/document_source_sample.h" @@ -181,7 +183,8 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> attemptToGetExe BSONObj projectionObj, BSONObj sortObj, const AggregationRequest* aggRequest, - const size_t plannerOpts) { + const size_t plannerOpts, + const MatchExpressionParser::AllowedFeatureSet& matcherFeatures) { auto qr = stdx::make_unique<QueryRequest>(nss); qr->setTailableMode(pExpCtx->tailableMode); qr->setOplogReplay(oplogReplay); @@ -206,7 +209,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> attemptToGetExe const ExtensionsCallbackReal extensionsCallback(pExpCtx->opCtx, &nss); auto cq = CanonicalQuery::canonicalize( - opCtx, std::move(qr), pExpCtx, extensionsCallback, Pipeline::kAllowedMatcherFeatures); + opCtx, std::move(qr), pExpCtx, extensionsCallback, matcherFeatures); if (!cq.isOK()) { // Return an error instead of uasserting, since there are cases where the combination of @@ -226,6 +229,50 @@ BSONObj removeSortKeyMetaProjection(BSONObj projectionObj) { } return projectionObj.removeField(Document::metaFieldSortKey); } + +/** + * Examines the indexes in 'collection' and returns the field name of a geo-indexed field suitable + * for use in $geoNear. 2d indexes are given priority over 2dsphere indexes. + * + * The 'collection' is required to exist. Throws if no usable 2d or 2dsphere index could be found. + */ +StringData extractGeoNearFieldFromIndexes(OperationContext* opCtx, Collection* collection) { + invariant(collection); + + std::vector<IndexDescriptor*> idxs; + collection->getIndexCatalog()->findIndexByType(opCtx, IndexNames::GEO_2D, idxs); + uassert(ErrorCodes::IndexNotFound, + str::stream() << "There is more than one 2d index on " << collection->ns().ns() + << "; unsure which to use for $geoNear", + idxs.size() <= 1U); + if (idxs.size() == 1U) { + for (auto&& elem : idxs.front()->keyPattern()) { + if (elem.type() == BSONType::String && elem.valueStringData() == IndexNames::GEO_2D) { + return elem.fieldNameStringData(); + } + } + MONGO_UNREACHABLE; + } + + // If there are no 2d indexes, look for a 2dsphere index. + idxs.clear(); + collection->getIndexCatalog()->findIndexByType(opCtx, IndexNames::GEO_2DSPHERE, idxs); + uassert(ErrorCodes::IndexNotFound, + "$geoNear requires a 2d or 2dsphere index, but none were found", + !idxs.empty()); + uassert(ErrorCodes::IndexNotFound, + str::stream() << "There is more than one 2dsphere index on " << collection->ns().ns() + << "; unsure which to use for $geoNear", + idxs.size() <= 1U); + + invariant(idxs.size() == 1U); + for (auto&& elem : idxs.front()->keyPattern()) { + if (elem.type() == BSONType::String && elem.valueStringData() == IndexNames::GEO_2DSPHERE) { + return elem.fieldNameStringData(); + } + } + MONGO_UNREACHABLE; +} } // namespace void PipelineD::prepareCursorSource(Collection* collection, @@ -260,16 +307,32 @@ void PipelineD::prepareCursorSource(Collection* collection, expCtx, sampleSize, idString, numRecords)); addCursorSource( - collection, pipeline, - expCtx, - std::move(exec), + DocumentSourceCursor::create(collection, std::move(exec), expCtx), pipeline->getDependencies(DepsTracker::MetadataAvailable::kNoMetadata)); return; } } } + // If the first stage is $geoNear, prepare a special DocumentSourceGeoNearCursor stage; + // otherwise, create a generic DocumentSourceCursor. + const auto geoNearStage = + sources.empty() ? nullptr : dynamic_cast<DocumentSourceGeoNear*>(sources.front().get()); + if (geoNearStage) { + prepareGeoNearCursorSource(collection, nss, aggRequest, pipeline); + } else { + prepareGenericCursorSource(collection, nss, aggRequest, pipeline); + } +} + +void PipelineD::prepareGenericCursorSource(Collection* collection, + const NamespaceString& nss, + const AggregationRequest* aggRequest, + Pipeline* pipeline) { + Pipeline::SourceContainer& sources = pipeline->_sources; + auto expCtx = pipeline->getContext(); + // Look for an initial match. This works whether we got an initial query or not. If not, it // results in a "{}" query, which will be what we want in that case. bool oplogReplay = false; @@ -283,7 +346,7 @@ void PipelineD::prepareCursorSource(Collection* collection, sources.pop_front(); } else { // A $geoNear stage, the only other stage that can produce an initial query, is also - // a valid initial stage and will be handled above. + // a valid initial stage. However, we should be in prepareGeoNearCursorSource() instead. MONGO_UNREACHABLE; } } @@ -321,6 +384,7 @@ void PipelineD::prepareCursorSource(Collection* collection, deps, queryObj, aggRequest, + Pipeline::kAllowedMatcherFeatures, &sortObj, &projForQuery)); @@ -335,8 +399,71 @@ void PipelineD::prepareCursorSource(Collection* collection, } } - addCursorSource( - collection, pipeline, expCtx, std::move(exec), deps, queryObj, sortObj, projForQuery); + addCursorSource(pipeline, + DocumentSourceCursor::create(collection, std::move(exec), expCtx), + deps, + queryObj, + sortObj, + projForQuery); +} + +void PipelineD::prepareGeoNearCursorSource(Collection* collection, + const NamespaceString& nss, + const AggregationRequest* aggRequest, + Pipeline* pipeline) { + uassert(ErrorCodes::NamespaceNotFound, + str::stream() << "$geoNear requires a geo index to run, but " << nss.ns() + << " does not exist", + collection); + + Pipeline::SourceContainer& sources = pipeline->_sources; + auto expCtx = pipeline->getContext(); + const auto geoNearStage = dynamic_cast<DocumentSourceGeoNear*>(sources.front().get()); + invariant(geoNearStage); + + auto deps = pipeline->getDependencies(DepsTracker::kAllGeoNearDataAvailable); + + // If the user specified a "key" field, use that field to satisfy the "near" query. Otherwise, + // look for a geo-indexed field in 'collection' that can. + auto nearFieldName = + (geoNearStage->getKeyField() ? geoNearStage->getKeyField()->fullPath() + : extractGeoNearFieldFromIndexes(expCtx->opCtx, collection)) + .toString(); + + // Create a PlanExecutor whose query is the "near" predicate on 'nearFieldName' combined with + // the optional "query" argument in the $geoNear stage. + BSONObj fullQuery = geoNearStage->asNearQuery(nearFieldName); + BSONObj proj = deps.toProjection(); + BSONObj sortFromQuerySystem; + auto exec = uassertStatusOK(prepareExecutor(expCtx->opCtx, + collection, + nss, + pipeline, + expCtx, + false, /* oplogReplay */ + nullptr, /* sortStage */ + deps, + std::move(fullQuery), + aggRequest, + Pipeline::kGeoNearMatcherFeatures, + &sortFromQuerySystem, + &proj)); + + invariant(sortFromQuerySystem.isEmpty(), + str::stream() << "Unexpectedly got the following sort from the query system: " + << sortFromQuerySystem.jsonString()); + + auto geoNearCursor = + DocumentSourceGeoNearCursor::create(collection, + std::move(exec), + expCtx, + geoNearStage->getDistanceField(), + geoNearStage->getLocationField(), + geoNearStage->getDistanceMultiplier().value_or(1.0)); + + // Remove the initial $geoNear; it will be replaced by $geoNearCursor. + sources.pop_front(); + addCursorSource(pipeline, std::move(geoNearCursor), std::move(deps)); } StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prepareExecutor( @@ -350,6 +477,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prep const DepsTracker& deps, const BSONObj& queryObj, const AggregationRequest* aggRequest, + const MatchExpressionParser::AllowedFeatureSet& matcherFeatures, BSONObj* sortObj, BSONObj* projectionObj) { // The query system has the potential to use an index to provide a non-blocking sort and/or to @@ -378,11 +506,11 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prep plannerOpts |= QueryPlannerParams::IS_COUNT; } - // The only way to get a text score or the sort key is to let the query system handle the - // projection. In all other cases, unless the query system can do an index-covered projection - // and avoid going to the raw record at all, it is faster to have ParsedDeps filter the fields - // we need. - if (!deps.getNeedTextScore() && !deps.getNeedSortKey()) { + // The only way to get meta information (e.g. the text score) is to let the query system handle + // the projection. In all other cases, unless the query system can do an index-covered + // projection and avoid going to the raw record at all, it is faster to have ParsedDeps filter + // the fields we need. + if (!deps.getNeedsAnyMetadata()) { plannerOpts |= QueryPlannerParams::NO_UNCOVERED_PROJECTIONS; } @@ -405,7 +533,8 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prep expCtx->needsMerge ? metaSortProjection : emptyProjection, *sortObj, aggRequest, - plannerOpts); + plannerOpts, + matcherFeatures); if (swExecutorSort.isOK()) { // Success! Now see if the query system can also cover the projection. @@ -418,7 +547,8 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prep *projectionObj, *sortObj, aggRequest, - plannerOpts); + plannerOpts, + matcherFeatures); std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec; if (swExecutorSortAndProj.isOK()) { @@ -458,7 +588,9 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prep // sort. dassert(sortObj->isEmpty()); *projectionObj = removeSortKeyMetaProjection(*projectionObj); - if (deps.getNeedSortKey() && !deps.getNeedTextScore()) { + const auto metadataRequired = deps.getAllRequiredMetadataTypes(); + if (metadataRequired.size() == 1 && + metadataRequired.front() == DepsTracker::MetadataType::SORT_KEY) { // A sort key requirement would have prevented us from being able to add this parameter // before, but now we know the query system won't cover the sort, so we will be able to // compute the sort key ourselves during the $sort stage, and thus don't need a query @@ -476,7 +608,8 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prep *projectionObj, *sortObj, aggRequest, - plannerOpts); + plannerOpts, + matcherFeatures); if (swExecutorProj.isOK()) { // Success! We have a covered projection. return std::move(swExecutorProj.getValue()); @@ -499,34 +632,24 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PipelineD::prep *projectionObj, *sortObj, aggRequest, - plannerOpts); + plannerOpts, + matcherFeatures); } -void PipelineD::addCursorSource(Collection* collection, - Pipeline* pipeline, - const intrusive_ptr<ExpressionContext>& expCtx, - unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec, +void PipelineD::addCursorSource(Pipeline* pipeline, + boost::intrusive_ptr<DocumentSourceCursor> cursor, DepsTracker deps, const BSONObj& queryObj, const BSONObj& sortObj, const BSONObj& projectionObj) { - // DocumentSourceCursor expects a yielding PlanExecutor that has had its state saved. - exec->saveState(); - - // Put the PlanExecutor into a DocumentSourceCursor and add it to the front of the pipeline. - intrusive_ptr<DocumentSourceCursor> pSource = - DocumentSourceCursor::create(collection, std::move(exec), expCtx); - - // Note the query, sort, and projection for explain. - pSource->setQuery(queryObj); - pSource->setSort(sortObj); - + cursor->setQuery(queryObj); + cursor->setSort(sortObj); if (deps.hasNoRequirements()) { - pSource->shouldProduceEmptyDocs(); + cursor->shouldProduceEmptyDocs(); } if (!projectionObj.isEmpty()) { - pSource->setProjection(projectionObj, boost::none); + cursor->setProjection(projectionObj, boost::none); } else { // There may be fewer dependencies now if the sort was covered. if (!sortObj.isEmpty()) { @@ -535,9 +658,9 @@ void PipelineD::addCursorSource(Collection* collection, : DepsTracker::MetadataAvailable::kNoMetadata); } - pSource->setProjection(deps.toProjection(), deps.toParsedDeps()); + cursor->setProjection(deps.toProjection(), deps.toParsedDeps()); } - pipeline->addInitialSource(pSource); + pipeline->addInitialSource(std::move(cursor)); } Timestamp PipelineD::getLatestOplogTimestamp(const Pipeline* pipeline) { diff --git a/src/mongo/db/pipeline/pipeline_d.h b/src/mongo/db/pipeline/pipeline_d.h index 4064b38e7c8..f85b13d1001 100644 --- a/src/mongo/db/pipeline/pipeline_d.h +++ b/src/mongo/db/pipeline/pipeline_d.h @@ -35,6 +35,7 @@ #include "mongo/db/dbdirectclient.h" #include "mongo/db/namespace_string.h" #include "mongo/db/pipeline/aggregation_request.h" +#include "mongo/db/pipeline/document_source_cursor.h" #include "mongo/db/pipeline/mongo_process_common.h" #include "mongo/db/query/plan_executor.h" @@ -136,7 +137,7 @@ public: /** * If the first stage in the pipeline does not generate its own output documents, attaches a - * DocumentSourceCursor to the front of the pipeline which will output documents from the + * cursor document source to the front of the pipeline which will output documents from the * collection to feed into the pipeline. * * This method looks for early pipeline stages that can be folded into the underlying @@ -154,6 +155,24 @@ public: Pipeline* pipeline); /** + * Prepare a generic DocumentSourceCursor for 'pipeline'. + */ + static void prepareGenericCursorSource(Collection* collection, + const NamespaceString& nss, + const AggregationRequest* aggRequest, + Pipeline* pipeline); + + /** + * Prepare a special DocumentSourceGeoNearCursor for 'pipeline'. Unlike + * 'prepareGenericCursorSource()', throws if 'collection' does not exist, as the $geoNearCursor + * requires a 2d or 2dsphere index. + */ + static void prepareGeoNearCursorSource(Collection* collection, + const NamespaceString& nss, + const AggregationRequest* aggRequest, + Pipeline* pipeline); + + /** * Injects a MongodInterface into stages which require access to mongod-specific functionality. */ static void injectMongodInterface(Pipeline* pipeline); @@ -187,17 +206,17 @@ private: const DepsTracker& deps, const BSONObj& queryObj, const AggregationRequest* aggRequest, + const MatchExpressionParser::AllowedFeatureSet& matcherFeatures, BSONObj* sortObj, BSONObj* projectionObj); /** - * Creates a DocumentSourceCursor from the given PlanExecutor and adds it to the front of the - * Pipeline. + * Adds 'cursor' to the front of 'pipeline', using 'deps' to inform the cursor of its + * dependencies. If specified, 'queryObj', 'sortObj' and 'projectionObj' are passed to the + * cursor for explain reporting. */ - static void addCursorSource(Collection* collection, - Pipeline* pipeline, - const boost::intrusive_ptr<ExpressionContext>& expCtx, - std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec, + static void addCursorSource(Pipeline* pipeline, + boost::intrusive_ptr<DocumentSourceCursor> cursor, DepsTracker deps, const BSONObj& queryObj = BSONObj(), const BSONObj& sortObj = BSONObj(), diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp index c53262cfc47..a0573014a1b 100644 --- a/src/mongo/db/pipeline/pipeline_test.cpp +++ b/src/mongo/db/pipeline/pipeline_test.cpp @@ -2492,11 +2492,11 @@ TEST_F(PipelineDependenciesTest, EmptyPipelineShouldRequireWholeDocument) { auto depsTracker = pipeline->getDependencies(DepsTracker::MetadataAvailable::kNoMetadata); ASSERT_TRUE(depsTracker.needWholeDocument); - ASSERT_FALSE(depsTracker.getNeedTextScore()); + ASSERT_FALSE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); depsTracker = pipeline->getDependencies(DepsTracker::MetadataAvailable::kTextScore); ASSERT_TRUE(depsTracker.needWholeDocument); - ASSERT_TRUE(depsTracker.getNeedTextScore()); + ASSERT_TRUE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } // @@ -2557,7 +2557,7 @@ public: class DocumentSourceNeedsOnlyTextScore : public DocumentSourceDependencyDummy { public: GetDepsReturn getDependencies(DepsTracker* deps) const final { - deps->setNeedTextScore(true); + deps->setNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE, true); return GetDepsReturn::EXHAUSTIVE_META; } @@ -2586,7 +2586,7 @@ TEST_F(PipelineDependenciesTest, ShouldRequireWholeDocumentIfAnyStageDoesNotSupp auto depsTracker = pipeline->getDependencies(DepsTracker::MetadataAvailable::kNoMetadata); ASSERT_TRUE(depsTracker.needWholeDocument); // The inputs did not have a text score available, so we should not require a text score. - ASSERT_FALSE(depsTracker.getNeedTextScore()); + ASSERT_FALSE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); // Now in the other order. pipeline = unittest::assertGet(Pipeline::create({notSupported, needsASeeNext}, ctx)); @@ -2625,7 +2625,7 @@ TEST_F(PipelineDependenciesTest, ShouldNotAddAnyRequiredFieldsAfterFirstStageWit auto depsTracker = pipeline->getDependencies(DepsTracker::MetadataAvailable::kNoMetadata); ASSERT_FALSE(depsTracker.needWholeDocument); - ASSERT_FALSE(depsTracker.getNeedTextScore()); + ASSERT_FALSE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); // 'needsOnlyB' claims to know all its field dependencies, so we shouldn't add any from // 'needsASeeNext'. @@ -2638,7 +2638,7 @@ TEST_F(PipelineDependenciesTest, ShouldNotRequireTextScoreIfThereIsNoScoreAvaila auto pipeline = unittest::assertGet(Pipeline::create({}, ctx)); auto depsTracker = pipeline->getDependencies(DepsTracker::MetadataAvailable::kNoMetadata); - ASSERT_FALSE(depsTracker.getNeedTextScore()); + ASSERT_FALSE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(PipelineDependenciesTest, ShouldThrowIfTextScoreIsNeededButNotPresent) { @@ -2655,12 +2655,12 @@ TEST_F(PipelineDependenciesTest, ShouldRequireTextScoreIfAvailableAndNoStageRetu auto pipeline = unittest::assertGet(Pipeline::create({}, ctx)); auto depsTracker = pipeline->getDependencies(DepsTracker::MetadataAvailable::kTextScore); - ASSERT_TRUE(depsTracker.getNeedTextScore()); + ASSERT_TRUE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); auto needsASeeNext = DocumentSourceNeedsASeeNext::create(); pipeline = unittest::assertGet(Pipeline::create({needsASeeNext}, ctx)); depsTracker = pipeline->getDependencies(DepsTracker::MetadataAvailable::kTextScore); - ASSERT_TRUE(depsTracker.getNeedTextScore()); + ASSERT_TRUE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } TEST_F(PipelineDependenciesTest, ShouldNotRequireTextScoreIfAvailableButDefinitelyNotNeeded) { @@ -2673,7 +2673,7 @@ TEST_F(PipelineDependenciesTest, ShouldNotRequireTextScoreIfAvailableButDefinite // 'stripsTextScore' claims that no further stage will need metadata information, so we // shouldn't have the text score as a dependency. - ASSERT_FALSE(depsTracker.getNeedTextScore()); + ASSERT_FALSE(depsTracker.getNeedsMetadata(DepsTracker::MetadataType::TEXT_SCORE)); } } // namespace Dependencies diff --git a/src/mongo/embedded/capi_test.cpp b/src/mongo/embedded/capi_test.cpp index 0efaacba25f..2d35c8495bd 100644 --- a/src/mongo/embedded/capi_test.cpp +++ b/src/mongo/embedded/capi_test.cpp @@ -542,7 +542,6 @@ TEST_F(MongodbCAPITest, RunListCommands) { "explain", "find", "findAndModify", - "geoNear", "getLastError", "getMore", "getParameter", diff --git a/src/mongo/s/commands/SConscript b/src/mongo/s/commands/SConscript index 9a6680c10ed..cfb4aca929c 100644 --- a/src/mongo/s/commands/SConscript +++ b/src/mongo/s/commands/SConscript @@ -51,7 +51,6 @@ env.Library( 'cluster_flush_router_config_cmd.cpp', 'cluster_fsync_cmd.cpp', 'cluster_ftdc_commands.cpp', - 'cluster_geo_near_cmd.cpp', 'cluster_get_last_error_cmd.cpp', 'cluster_get_prev_error_cmd.cpp', 'cluster_get_shard_version_cmd.cpp', diff --git a/src/mongo/s/commands/cluster_geo_near_cmd.cpp b/src/mongo/s/commands/cluster_geo_near_cmd.cpp deleted file mode 100644 index 6871be7565e..00000000000 --- a/src/mongo/s/commands/cluster_geo_near_cmd.cpp +++ /dev/null @@ -1,186 +0,0 @@ -/** - * Copyright (C) 2018 MongoDB 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/>. - * - * As a special exception, the copyright holders give permission to link the - * code of portions of this program with the OpenSSL library under certain - * conditions as described in each individual source file and distribute - * linked combinations including the program with the OpenSSL library. You - * must comply with the GNU Affero General Public License in all respects for - * all of the code used other than as permitted herein. If you modify file(s) - * with this exception, you may extend this exception to your version of the - * file(s), but you are not obligated to do so. If you do not wish to do so, - * delete this exception statement from your version. If you delete this - * exception statement from all source files in the program, then also delete - * it in the license file. - */ - -#define MONGO_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kCommand - -#include "mongo/platform/basic.h" - -#include "mongo/db/commands.h" -#include "mongo/rpc/get_status_from_command_result.h" -#include "mongo/s/commands/cluster_commands_helpers.h" -#include "mongo/s/grid.h" -#include "mongo/util/log.h" - -namespace mongo { -namespace { - -class Geo2dFindNearCmd : public BasicCommand { -public: - Geo2dFindNearCmd() : BasicCommand("geoNear") {} - - std::string help() const override { - return "The geoNear command is deprecated. See " - "http://dochub.mongodb.org/core/geoNear-deprecation for more detail on its " - "replacement."; - } - - AllowedOnSecondary secondaryAllowed(ServiceContext*) const override { - return AllowedOnSecondary::kAlways; - } - - bool adminOnly() const override { - return false; - } - - bool supportsWriteConcern(const BSONObj& cmd) const override { - return false; - } - - std::string parseNs(const std::string& dbname, const BSONObj& cmdObj) const override { - return CommandHelpers::parseNsCollectionRequired(dbname, cmdObj).ns(); - } - - void addRequiredPrivileges(const std::string& dbname, - const BSONObj& cmdObj, - std::vector<Privilege>* out) const override { - ActionSet actions; - actions.addAction(ActionType::find); - out->push_back(Privilege(parseResourcePattern(dbname, cmdObj), actions)); - } - - bool run(OperationContext* opCtx, - const std::string& dbName, - const BSONObj& cmdObj, - BSONObjBuilder& result) override { - RARELY { - warning() << "Support for the geoNear command has been deprecated. Please plan to " - "rewrite geoNear commands using the $near query operator, the $nearSphere " - "query operator, or the $geoNear aggregation stage. See " - "http://dochub.mongodb.org/core/geoNear-deprecation."; - } - - const NamespaceString nss(parseNs(dbName, cmdObj)); - - // We support both "num" and "limit" options to control limit - long long limit = 100; - if (cmdObj["num"].isNumber()) - limit = cmdObj["num"].safeNumberLong(); - else if (cmdObj["limit"].isNumber()) - limit = cmdObj["limit"].safeNumberLong(); - - const auto query = extractQuery(cmdObj); - const auto collation = extractCollation(cmdObj); - - const auto routingInfo = - uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss)); - - auto shardResponses = scatterGatherVersionedTargetByRoutingTable( - opCtx, - nss.db(), - nss, - routingInfo, - CommandHelpers::filterCommandRequestForPassthrough(cmdObj), - ReadPreferenceSetting::get(opCtx), - Shard::RetryPolicy::kIdempotent, - query, - collation); - - std::multimap<double, BSONObj> results; - BSONArrayBuilder shardArray; - std::string nearStr; - double time = 0; - double btreelocs = 0; - double nscanned = 0; - double objectsLoaded = 0; - - for (const auto& shardResponse : shardResponses) { - const auto response = uassertStatusOK(shardResponse.swResponse); - uassertStatusOK(getStatusFromCommandResult(response.data)); - - shardArray.append(shardResponse.shardId.toString()); - const auto& shardResult = response.data; - - if (shardResult.hasField("near")) { - nearStr = shardResult["near"].String(); - } - time += shardResult["stats"]["time"].Number(); - if (!shardResult["stats"]["btreelocs"].eoo()) { - btreelocs += shardResult["stats"]["btreelocs"].Number(); - } - nscanned += shardResult["stats"]["nscanned"].Number(); - if (!shardResult["stats"]["objectsLoaded"].eoo()) { - objectsLoaded += shardResult["stats"]["objectsLoaded"].Number(); - } - - BSONForEach(obj, shardResult["results"].embeddedObject()) { - results.insert( - std::make_pair(obj["dis"].Number(), obj.embeddedObject().getOwned())); - } - - // TODO: maybe shrink results if size() > limit - } - - result.append("ns", nss.ns()); - result.append("near", nearStr); - - long long outCount = 0; - double totalDistance = 0; - double maxDistance = 0; - { - BSONArrayBuilder sub(result.subarrayStart("results")); - for (std::multimap<double, BSONObj>::const_iterator it(results.begin()), - end(results.end()); - it != end && outCount < limit; - ++it, ++outCount) { - totalDistance += it->first; - maxDistance = it->first; // guaranteed to be highest so far - - sub.append(it->second); - } - sub.done(); - } - - { - BSONObjBuilder sub(result.subobjStart("stats")); - sub.append("time", time); - sub.append("btreelocs", btreelocs); - sub.append("nscanned", nscanned); - sub.append("objectsLoaded", objectsLoaded); - sub.append("avgDistance", (outCount == 0) ? 0 : (totalDistance / outCount)); - sub.append("maxDistance", maxDistance); - sub.append("shards", shardArray.arr()); - sub.done(); - } - - return true; - } - -} geo2dFindNearCmd; - -} // namespace -} // namespace mongo |