diff options
Diffstat (limited to 'src/mongo/db/query')
69 files changed, 2666 insertions, 1022 deletions
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript index 025fd11d118..c34a713cc1b 100644 --- a/src/mongo/db/query/SConscript +++ b/src/mongo/db/query/SConscript @@ -61,7 +61,6 @@ env.Library( "$BUILD_DIR/mongo/db/commands/server_status_core", "$BUILD_DIR/mongo/db/exec/sbe/query_sbe_plan_stats", "$BUILD_DIR/mongo/db/index/expression_params", - "$BUILD_DIR/mongo/db/index/key_generator", "$BUILD_DIR/mongo/db/index_names", "canonical_query", "query_index_bounds", @@ -70,6 +69,7 @@ env.Library( ], LIBDEPS_PRIVATE=[ '$BUILD_DIR/mongo/db/fts/base_fts', + '$BUILD_DIR/mongo/db/index/index_access_method', "$BUILD_DIR/mongo/db/record_id_helpers", "$BUILD_DIR/mongo/idl/server_parameter", ], @@ -342,7 +342,7 @@ env.Library( '$BUILD_DIR/mongo/util/fail_point', ], LIBDEPS_PRIVATE=[ - '$BUILD_DIR/mongo/db/concurrency/write_conflict_exception', + '$BUILD_DIR/mongo/db/concurrency/exception_util', '$BUILD_DIR/mongo/db/storage/recovery_unit_base', ], ) @@ -356,6 +356,7 @@ env.CppUnitTest( "classic_stage_builder_test.cpp", "count_command_test.cpp", "cursor_response_test.cpp", + "find_common_test.cpp", "get_executor_test.cpp", "getmore_request_test.cpp", "hint_parser_test.cpp", diff --git a/src/mongo/db/query/canonical_query_encoder.cpp b/src/mongo/db/query/canonical_query_encoder.cpp index ccbcce36635..c8a8cd29d48 100644 --- a/src/mongo/db/query/canonical_query_encoder.cpp +++ b/src/mongo/db/query/canonical_query_encoder.cpp @@ -74,32 +74,6 @@ bool isQueryNegatingEqualToNull(const mongo::MatchExpression* tree) { namespace { -// Delimiters for cache key encoding. -const char kEncodeChildrenBegin = '['; -const char kEncodeChildrenEnd = ']'; -const char kEncodeChildrenSeparator = ','; -const char kEncodeCollationSection = '#'; -const char kEncodeProjectionSection = '|'; -const char kEncodeProjectionRequirementSeparator = '-'; -const char kEncodeRegexFlagsSeparator = '/'; -const char kEncodeSortSection = '~'; -const char kEncodeEngineSection = '@'; - -// These special bytes are used in the encoding of auto-parameterized match expressions in the SBE -// plan cache key. - -// Precedes the id number of a parameter marker. -const char kEncodeParamMarker = '?'; -// Precedes the encoding of a constant when that constant has not been auto-paramterized. The -// constant is typically encoded as a BSON type byte followed by a BSON value (without the -// BSONElement's field name). -const char kEncodeConstantLiteralMarker = ':'; -// Precedes a byte which encodes the bounds tightness associated with a predicate. The structure of -// the plan (i.e. presence of filters) is affected by bounds tightness. Therefore, if different -// parameter values can result in different tightnesses, this must be explicitly encoded into the -// plan cache key. -const char kEncodeBoundsTightnessDiscriminator = ':'; - /** * AppendChar provides the compiler with a type for a "appendChar(...)" member function. */ @@ -584,7 +558,7 @@ void encodeKeyForProj(const projection_ast::Projection* proj, StringBuilder* key return; } - std::set<std::string> requiredFields = proj->getRequiredFields(); + auto requiredFields = proj->getRequiredFields(); // If the only requirement is that $sortKey be included with some value, we just act as if the // entire document is needed. diff --git a/src/mongo/db/query/canonical_query_encoder.h b/src/mongo/db/query/canonical_query_encoder.h index 3164ddbec67..cf124655681 100644 --- a/src/mongo/db/query/canonical_query_encoder.h +++ b/src/mongo/db/query/canonical_query_encoder.h @@ -33,6 +33,38 @@ namespace mongo { +// Delimiters for canonical query portion of cache key encoding. +inline constexpr char kEncodeChildrenBegin = '['; +inline constexpr char kEncodeChildrenEnd = ']'; +inline constexpr char kEncodeChildrenSeparator = ','; +inline constexpr char kEncodeCollationSection = '#'; +inline constexpr char kEncodeProjectionSection = '|'; +inline constexpr char kEncodeProjectionRequirementSeparator = '-'; +inline constexpr char kEncodeRegexFlagsSeparator = '/'; +inline constexpr char kEncodeSortSection = '~'; +inline constexpr char kEncodeEngineSection = '@'; + +// These special bytes are used in the encoding of auto-parameterized match expressions in the SBE +// plan cache key. + +// Precedes the id number of a parameter marker. +inline constexpr char kEncodeParamMarker = '?'; +// Precedes the encoding of a constant when that constant has not been auto-paramterized. The +// constant is typically encoded as a BSON type byte followed by a BSON value (without the +// BSONElement's field name). +inline constexpr char kEncodeConstantLiteralMarker = ':'; +// Precedes a byte which encodes the bounds tightness associated with a predicate. The structure of +// the plan (i.e. presence of filters) is affected by bounds tightness. Therefore, if different +// parameter values can result in different tightnesses, this must be explicitly encoded into the +// plan cache key. +inline constexpr char kEncodeBoundsTightnessDiscriminator = ':'; + +// Delimiters for the discriminator portion of the cache key encoding. +inline constexpr char kEncodeDiscriminatorsBegin = '<'; +inline constexpr char kEncodeDiscriminatorsEnd = '>'; +inline constexpr char kEncodeGlobalDiscriminatorsBegin = '('; +inline constexpr char kEncodeGlobalDiscriminatorsEnd = ')'; + /** * Returns true if the query predicate involves a negation of an EQ, LTE, or GTE comparison to * 'null'. diff --git a/src/mongo/db/query/collection_query_info.cpp b/src/mongo/db/query/collection_query_info.cpp index e624389b642..ff4e430f6a1 100644 --- a/src/mongo/db/query/collection_query_info.cpp +++ b/src/mongo/db/query/collection_query_info.cpp @@ -90,12 +90,13 @@ CollectionQueryInfo::PlanCacheState::PlanCacheState(OperationContext* opCtx, // TODO We shouldn't need to include unfinished indexes, but we must here because the index // catalog may be in an inconsistent state. SERVER-18346. - const bool includeUnfinishedIndexes = true; - std::unique_ptr<IndexCatalog::IndexIterator> ii = - collection->getIndexCatalog()->getIndexIterator(opCtx, includeUnfinishedIndexes); + auto ii = collection->getIndexCatalog()->getIndexIterator( + opCtx, IndexCatalog::InclusionPolicy::kReady | IndexCatalog::InclusionPolicy::kUnfinished); while (ii->more()) { const IndexCatalogEntry* ice = ii->next(); - indexCores.emplace_back(indexInfoFromIndexCatalogEntry(*ice)); + if (ice->accessMethod()) { + indexCores.emplace_back(indexInfoFromIndexCatalogEntry(*ice)); + } } planCacheIndexabilityState.updateDiscriminators(indexCores); @@ -117,8 +118,8 @@ const UpdateIndexData& CollectionQueryInfo::getIndexKeys(OperationContext* opCtx void CollectionQueryInfo::computeIndexKeys(OperationContext* opCtx, const CollectionPtr& coll) { _indexedPaths.clear(); - std::unique_ptr<IndexCatalog::IndexIterator> it = - coll->getIndexCatalog()->getIndexIterator(opCtx, true); + auto it = coll->getIndexCatalog()->getIndexIterator( + opCtx, IndexCatalog::InclusionPolicy::kReady | IndexCatalog::InclusionPolicy::kUnfinished); while (it->more()) { const IndexCatalogEntry* entry = it->next(); const IndexDescriptor* descriptor = entry->descriptor(); @@ -175,10 +176,10 @@ void CollectionQueryInfo::computeIndexKeys(OperationContext* opCtx, const Collec // handle partial indexes const MatchExpression* filter = entry->getFilterExpression(); if (filter) { - stdx::unordered_set<std::string> paths; + RelevantFieldIndexMap paths; QueryPlannerIXSelect::getFields(filter, &paths); for (auto it = paths.begin(); it != paths.end(); ++it) { - _indexedPaths.addPath(FieldRef(*it)); + _indexedPaths.addPath(FieldRef(it->first)); } } } @@ -238,9 +239,8 @@ void CollectionQueryInfo::updatePlanCacheIndexEntries(OperationContext* opCtx, } void CollectionQueryInfo::init(OperationContext* opCtx, const CollectionPtr& coll) { - const bool includeUnfinishedIndexes = false; - std::unique_ptr<IndexCatalog::IndexIterator> ii = - coll->getIndexCatalog()->getIndexIterator(opCtx, includeUnfinishedIndexes); + auto ii = + coll->getIndexCatalog()->getIndexIterator(opCtx, IndexCatalog::InclusionPolicy::kReady); while (ii->more()) { const IndexDescriptor* desc = ii->next()->descriptor(); CollectionIndexUsageTrackerDecoration::get(coll->getSharedDecorations()) diff --git a/src/mongo/db/query/datetime/date_time_support.cpp b/src/mongo/db/query/datetime/date_time_support.cpp index 1fe0ecd81d0..93859935c98 100644 --- a/src/mongo/db/query/datetime/date_time_support.cpp +++ b/src/mongo/db/query/datetime/date_time_support.cpp @@ -842,7 +842,7 @@ StringData serializeTimeUnit(TimeUnit unit) { case TimeUnit::millisecond: return "millisecond"_sd; } - MONGO_UNREACHABLE_TASSERT(5339900); + MONGO_UNREACHABLE_TASSERT(5339903); } DayOfWeek parseDayOfWeek(StringData dayOfWeek) { diff --git a/src/mongo/db/query/datetime/date_time_support_test.cpp b/src/mongo/db/query/datetime/date_time_support_test.cpp index d2b2e3b2d44..df6e97c979a 100644 --- a/src/mongo/db/query/datetime/date_time_support_test.cpp +++ b/src/mongo/db/query/datetime/date_time_support_test.cpp @@ -2629,7 +2629,10 @@ TEST(DateAdd, DateAddWithTimezoneDST) { {europeAmsterdamZone.createFromDateParts(2020, 10, 24, 2, 0, 1, 0), TimeUnit::day, 1, - europeAmsterdamZone.createFromDateParts(2020, 10, 25, 2, 0, 1, 0)}, + europeAmsterdamZone.createFromDateParts(2020, 10, 25, 1, 59, 59, 0) + + Milliseconds{2000}}, // as this date is ambiguous (it could in both timezones, with or + // without DST) and the computation is expected to return the + // "with DST" one, obtain it via a computation {europeAmsterdamZone.createFromDateParts(2020, 10, 24, 3, 0, 1, 0), TimeUnit::day, 1, @@ -2721,10 +2724,13 @@ TEST(DateAdd, DateAddWithTimezoneDST) { TimeUnit::day, 1, newYorkZone.createFromDateParts(2020, 11, 2, 1, 30, 0, 0)}, - {newYorkZone.createFromDateParts(2020, 10, 31, 1, 30, 0, 0), + {newYorkZone.createFromDateParts(2020, 10, 31, 1, 0, 1, 0), TimeUnit::day, 1, - newYorkZone.createFromDateParts(2020, 11, 1, 1, 30, 0, 0)}, + newYorkZone.createFromDateParts(2020, 11, 1, 0, 59, 59, 0) + + Milliseconds{2000}}, // as this date is ambiguous (it could in both timezones, with or + // without DST) and the computation is expected to return the + // "with DST" one, obtain it via a computation {newYorkZone.createFromDateParts(2020, 11, 1, 3, 0, 0, 0), TimeUnit::day, -1, @@ -2785,15 +2791,20 @@ TEST(DateAdd, DateAdd_LordHoweTimezoneDST) { auto australiaLordHoweZone = kDefaultTimeZoneDatabase.getTimeZone("Australia/Lord_Howe"); std::vector<TestCase> tests{ // DST to Standard change: 2021-04-04T02:00:00 -> 2021-04-04T01:30:00 Lord Howe timezone. - {australiaLordHoweZone.createFromDateParts(2021, 4, 4, 1, 30, 0, 0), + {australiaLordHoweZone.createFromDateParts(2021, 4, 4, 1, 29, 59, 0) + + Milliseconds{1000}, // as this date is ambiguous (it could in both timezones, with or + // without DST) and the computation is expected to start from the + // "with DST" one, obtain it via a computation TimeUnit::day, 1, australiaLordHoweZone.createFromDateParts(2021, 4, 5, 1, 30, 0, 0)}, - {australiaLordHoweZone.createFromDateParts(2021, 4, 3, 1, 45, 0, 0), + {australiaLordHoweZone.createFromDateParts(2021, 4, 3, 1, 30, 1, 0), TimeUnit::day, 1, - // Computed time falls into the repeated 1/2 hour. - australiaLordHoweZone.createFromDateParts(2021, 4, 4, 1, 45, 0, 0)}, + australiaLordHoweZone.createFromDateParts(2021, 4, 4, 1, 29, 59, 0) + + Milliseconds{2000}}, // as this date is ambiguous (it could in both timezones, with or + // without DST) and the computation is expected to return the + // "with DST" one, obtain it via a computation {australiaLordHoweZone.createFromDateParts(2021, 4, 5, 1, 0, 0, 0), TimeUnit::day, -1, diff --git a/src/mongo/db/query/find_common.cpp b/src/mongo/db/query/find_common.cpp index 57f9f3954b2..0078533cc60 100644 --- a/src/mongo/db/query/find_common.cpp +++ b/src/mongo/db/query/find_common.cpp @@ -133,5 +133,17 @@ std::size_t FindCommon::getBytesToReserveForGetMoreReply(bool isTailable, // command metadata to the reply. return kMaxBytesToReturnToClientAtOnce; } +bool FindCommon::BSONArrayResponseSizeTracker::haveSpaceForNext(const BSONObj& document) { + return FindCommon::haveSpaceForNext(document, _numberOfDocuments, _bsonArraySizeInBytes); +} +void FindCommon::BSONArrayResponseSizeTracker::add(const BSONObj& document) { + dassert(haveSpaceForNext(document)); + ++_numberOfDocuments; + _bsonArraySizeInBytes += (document.objsize() + kPerDocumentOverheadBytesUpperBound); +} +// Upper bound of BSON array element overhead. The overhead is 1 byte/doc for the type + 1 byte/doc +// for the field name's null terminator + 1 byte per digit of the maximum array index value. +const size_t FindCommon::BSONArrayResponseSizeTracker::kPerDocumentOverheadBytesUpperBound{ + 2 + std::to_string(BSONObjMaxUserSize / BSONObj::kMinBSONLength).length()}; } // namespace mongo diff --git a/src/mongo/db/query/find_common.h b/src/mongo/db/query/find_common.h index 45f60d2fd51..d38d580bdd2 100644 --- a/src/mongo/db/query/find_common.h +++ b/src/mongo/db/query/find_common.h @@ -89,7 +89,7 @@ public: // This max may be exceeded by epsilon for output documents that approach the maximum user // document size. That is, if we must return a BSONObjMaxUserSize document, then the total // response size will be BSONObjMaxUserSize plus the amount of size required for the message - // header and the cursor response "envelope". (The envolope contains namespace and cursor id + // header and the cursor response "envelope". (The envelope contains namespace and cursor id // info.) static const size_t kMaxBytesToReturnToClientAtOnce; @@ -148,6 +148,32 @@ public: static std::size_t getBytesToReserveForGetMoreReply(bool isTailable, size_t firstResultSize, size_t batchSize); + + /** + * Tracker of a size of a server response presented as a BSON array. Facilitates limiting the + * server response size to 16MB + certain epsilon. Accounts for array element and it's overhead + * size. Does not account for response "envelope" size. + */ + class BSONArrayResponseSizeTracker { + // Upper bound of BSON array element overhead. + static const size_t kPerDocumentOverheadBytesUpperBound; + + public: + /** + * Returns true only if 'document' can be added to the BSON array without violating the + * overall response size limit or if it is the first document. + */ + bool haveSpaceForNext(const BSONObj& document); + + /** + * Records that 'document' was added to the response. + */ + void add(const BSONObj& document); + + private: + std::size_t _numberOfDocuments{0}; + std::size_t _bsonArraySizeInBytes{0}; + }; }; } // namespace mongo diff --git a/src/mongo/db/query/find_common_test.cpp b/src/mongo/db/query/find_common_test.cpp new file mode 100644 index 00000000000..d7dfc10d950 --- /dev/null +++ b/src/mongo/db/query/find_common_test.cpp @@ -0,0 +1,74 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 <string> + +#include "mongo/bson/bsonobj.h" +#include "mongo/bson/bsonobjbuilder.h" +#include "mongo/db/query/find_common.h" + +#include "mongo/unittest/unittest.h" + +namespace { + +using namespace mongo; + +TEST(BSONArrayResponseSizeTrackerTest, AddLargeNumberOfElements) { + BSONObjBuilder bsonObjBuilder; + { + FindCommon::BSONArrayResponseSizeTracker sizeTracker; + BSONArrayBuilder arrayBuilder{bsonObjBuilder.subarrayStart("a")}; + BSONObj emptyObject; + while (sizeTracker.haveSpaceForNext(emptyObject)) { + sizeTracker.add(emptyObject); + arrayBuilder.append(emptyObject); + } + } + // If the BSON object is successfully constructed, then space accounting was correct. + bsonObjBuilder.obj(); +} +TEST(BSONArrayResponseSizeTrackerTest, CanAddAtLeastOneDocument) { + auto largeObject = BSON("a" << std::string(16 * 1024 * 1024, 'A')); + BSONObj emptyObject; + BSONObjBuilder bsonObjBuilder; + { + FindCommon::BSONArrayResponseSizeTracker sizeTracker; + BSONArrayBuilder arrayBuilder{bsonObjBuilder.subarrayStart("a")}; + // Add an object that is larger than 16MB. + ASSERT(sizeTracker.haveSpaceForNext(largeObject)); + sizeTracker.add(largeObject); + arrayBuilder.append(largeObject); + ASSERT(!sizeTracker.haveSpaceForNext(emptyObject)); + } + // If the BSON object is successfully constructed, then space accounting was correct. + bsonObjBuilder.obj(); +} +} // namespace diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index db04d6a276a..bd22078b7c9 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -216,10 +216,15 @@ IndexEntry indexEntryFromIndexCatalogEntry(OperationContext* opCtx, MultikeyMetadataAccessStats mkAccessStats; if (canonicalQuery) { - stdx::unordered_set<std::string> fields; - QueryPlannerIXSelect::getFields(canonicalQuery->root(), &fields); - const auto projectedFields = projection_executor_utils::applyProjectionToFields( - wildcardProjection->exec(), fields); + RelevantFieldIndexMap fieldIndexProps; + QueryPlannerIXSelect::getFields(canonicalQuery->root(), &fieldIndexProps); + stdx::unordered_set<std::string> projectedFields; + for (auto&& [fieldName, _] : fieldIndexProps) { + if (projection_executor_utils::applyProjectionToOneField( + wildcardProjection->exec(), fieldName)) { + projectedFields.insert(fieldName); + } + } multikeyPathSet = getWildcardMultikeyPathSet(wam, opCtx, projectedFields, &mkAccessStats); @@ -284,7 +289,8 @@ void fillOutIndexEntries(OperationContext* opCtx, const CanonicalQuery* canonicalQuery, const CollectionPtr& collection, std::vector<IndexEntry>& entries) { - auto ii = collection->getIndexCatalog()->getIndexIterator(opCtx, false); + auto ii = collection->getIndexCatalog()->getIndexIterator( + opCtx, IndexCatalog::InclusionPolicy::kReady); while (ii->more()) { const IndexCatalogEntry* ice = ii->next(); @@ -1526,6 +1532,13 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDele expCtx->setIsCappedDelete(); } + // If the parsed delete does not have a user-specified collation, set it from the collection + // default. + if (collection && parsedDelete->getRequest()->getCollation().isEmpty() && + collection->getDefaultCollator()) { + parsedDelete->setCollator(collection->getDefaultCollator()->clone()); + } + if (collection && collection->isCapped() && opCtx->inMultiDocumentTransaction()) { // This check is duplicated from CollectionImpl::deleteDocument() for two reasons: // - Performing a remove on an empty capped collection would not call @@ -1662,9 +1675,9 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDele deleteStageParams->canonicalQuery = cq.get(); const bool batchDelete = - (deleteStageParams->isMulti && !deleteStageParams->fromMigrate && - !deleteStageParams->returnDeleted && deleteStageParams->sort.isEmpty() && - !deleteStageParams->numStatsForDoc) && + (deleteStageParams->isMulti && !opCtx->inMultiDocumentTransaction() && + !deleteStageParams->fromMigrate && !deleteStageParams->returnDeleted && + deleteStageParams->sort.isEmpty() && !deleteStageParams->numStatsForDoc) && ((gInternalBatchUserMultiDeletesForTest.load() && nss.ns() == "__internalBatchedDeletesTesting.Collection0") || (batchDeletesByDefault.shouldFail())); @@ -2393,8 +2406,8 @@ QueryPlannerParams fillOutPlannerParamsForDistinct(OperationContext* opCtx, // If the caller did not request a "strict" distinct scan then we may choose a plan which // unwinds arrays and treats each element in an array as its own key. const bool mayUnwindArrays = !(plannerOptions & QueryPlannerParams::STRICT_DISTINCT_ONLY); - std::unique_ptr<IndexCatalog::IndexIterator> ii = - collection->getIndexCatalog()->getIndexIterator(opCtx, false); + auto ii = collection->getIndexCatalog()->getIndexIterator( + opCtx, IndexCatalog::InclusionPolicy::kReady); auto query = parsedDistinct.getQuery()->getFindCommandRequest().getFilter(); while (ii->more()) { const IndexCatalogEntry* ice = ii->next(); diff --git a/src/mongo/db/query/index_bounds_builder.cpp b/src/mongo/db/query/index_bounds_builder.cpp index 01e13b5f058..3c027bde140 100644 --- a/src/mongo/db/query/index_bounds_builder.cpp +++ b/src/mongo/db/query/index_bounds_builder.cpp @@ -108,13 +108,38 @@ Interval makeNullPointInterval(bool isHashed) { return isHashed ? kHashedNullInterval : IndexBoundsBuilder::kNullPointInterval; } +/** + * This helper updates the query bounds tightness for the limited set of conditions where we see a + * null query that can be covered. + */ +void updateTightnessForNullQuery(const IndexEntry& index, + IndexBoundsBuilder::BoundsTightness* tightnessOut) { + if (index.sparse || index.type == IndexType::INDEX_HASHED) { + // Sparse indexes and hashed indexes require a FETCH stage with a filter for null queries. + *tightnessOut = IndexBoundsBuilder::INEXACT_FETCH; + return; + } + + if (index.multikey) { + // If we have a simple equality null query and our index is multikey, we cannot cover the + // query. This is because null intervals are translated into the null and undefined point + // intervals, and the undefined point interval includes entries for []. In the case of a + // single null interval, [] should not match. + *tightnessOut = IndexBoundsBuilder::INEXACT_FETCH; + return; + } + + // The query may be fully covered by the index if the projection allows it, since the case above + // about the empty array can only become an issue if there is an empty array present, which + // would mark the index as multikey. + *tightnessOut = IndexBoundsBuilder::EXACT_MAYBE_COVERED; +} + void makeNullEqualityBounds(const IndexEntry& index, bool isHashed, OrderedIntervalList* oil, IndexBoundsBuilder::BoundsTightness* tightnessOut) { - // An equality to null predicate cannot be covered because the index does not distinguish - // between the lack of a value and the literal value null. - *tightnessOut = IndexBoundsBuilder::INEXACT_FETCH; + updateTightnessForNullQuery(index, tightnessOut); // There are two values that could possibly be equal to null in an index: undefined and null. oil->intervals.push_back(makeUndefinedPointInterval(isHashed)); @@ -254,7 +279,10 @@ bool IndexBoundsBuilder::canUseCoveredMatching(const MatchExpression* expr, IndexBoundsBuilder::BoundsTightness tightness; OrderedIntervalList oil; translate(expr, BSONElement{}, index, &oil, &tightness, /* iet::Builder */ nullptr); - return tightness >= IndexBoundsBuilder::INEXACT_COVERED; + // We have additional tightness values (MAYBE_COVERED), but we cannot generally cover those + // cases unless we have an appropriate projection. + return tightness == IndexBoundsBuilder::INEXACT_COVERED || + tightness == IndexBoundsBuilder::EXACT; } // static @@ -404,6 +432,61 @@ const Interval IndexBoundsBuilder::kNullPointInterval = const Interval IndexBoundsBuilder::kEmptyArrayPointInterval = IndexBoundsBuilder::makePointInterval(kEmptyArrayElementObj); +bool detectIfEntireNullIntervalMatchesPredicate(const InMatchExpression* ime, + const IndexEntry& index) { + if (!ime->hasNull()) { + // This isn't a null query. + return false; + } + + if (index.sparse || (IndexType::INDEX_HASHED == index.type)) { + // Sparse indexes and hashed indexes still require a FETCH stage with a filter for null + // queries. + return false; + } + + // Given the context of having a null $in query with eligible indexes, we may be able to cover + // some combinations of intervals that we could not cover individually. + if (index.multikey) { + // If the path has multiple components and we have a multikey index, we still need a FETCH + // in order to defend against cases where we have a multikey index on "a". These documents + // will generate null index keys: {"a.b": null} and {a: [1,2,3]}. However, a query like + // {"a.b": {$in: [null, []]}} should not match {a: [1, 2, 3]}. + // TODO SERVER-71021: it may be possible to cover more cases here. + if (ime->fieldRef()->numParts() > 1) { + return false; + } + + // We must have an equality to an empty array for this null query to be covered, otherwise, + // because we generate both null and undefined point intervals for a null query, and because + // a multikey index reuses the same entry for [] and undefined, we will not be able to cover + // the query. + if (!ime->hasEmptyArray()) { + return false; + } + } + + return true; +} + +void IndexBoundsBuilder::_mergeTightness(const BoundsTightness& tightness, + BoundsTightness& tightnessOut) { + // There is a special case where we may have a covered null query (EXACT_MAYBE_COVERED) and a + // regex with inexact bounds that doesn't need a FETCH (INEXACT_COVERED). In this case, we want + // to update the tightness to INEXACT_MAYBE_COVERED, to indicate that we need to check if the + // projection allows us to cover the query, but ensure that we will have a filter on the index + // if it turns out we can. + if (((tightness == BoundsTightness::EXACT_MAYBE_COVERED) && + (tightnessOut == BoundsTightness::INEXACT_COVERED)) || + ((tightness == BoundsTightness::INEXACT_COVERED) && + (tightnessOut == BoundsTightness::EXACT_MAYBE_COVERED))) { + tightnessOut = BoundsTightness::INEXACT_MAYBE_COVERED; + } else if (tightness < tightnessOut) { + // Otherwise, fallback to picking the new tightness if it is looser than the old tightness. + tightnessOut = tightness; + } +} + void IndexBoundsBuilder::_translatePredicate(const MatchExpression* expr, const BSONElement& elt, const IndexEntry& index, @@ -955,51 +1038,45 @@ void IndexBoundsBuilder::_translatePredicate(const MatchExpression* expr, }); const InMatchExpression* ime = static_cast<const InMatchExpression*>(expr); - *tightnessOut = IndexBoundsBuilder::EXACT; // Create our various intervals. IndexBoundsBuilder::BoundsTightness tightness; - bool arrayOrNullPresent = false; + // We check if the $in predicate satisfies conditions to be a covered null predicate on the + // basis of indexes, null intervals, and array intervals. + const bool entireNullIntervalMatchesPredicate = + detectIfEntireNullIntervalMatchesPredicate(ime, index); for (auto&& equality : ime->getEqualities()) { - translateEquality(equality, index, isHashed, oilOut, &tightness); - // The ordering invariant of oil has been violated by the call to translateEquality. - arrayOrNullPresent = arrayOrNullPresent || equality.type() == BSONType::jstNULL || - equality.type() == BSONType::Array; - if (tightness != IndexBoundsBuilder::EXACT) { - *tightnessOut = tightness; + // First, we generate the bounds the same way that we would do for an individual + // equality. This will set tightness to the value it should be if this equality is being + // considered in isolation. + IndexBoundsBuilder::translateEquality(equality, index, isHashed, oilOut, &tightness); + if (entireNullIntervalMatchesPredicate && + (BSONType::jstNULL == equality.type() || + (BSONType::Array == equality.type() && equality.Obj().isEmpty()))) { + // We may have a covered null query. In this case, we update both empty array and + // null interval tightness to EXACT_MAYBE_COVERED, as individually they would have a + // tightness of INEXACT_FETCH. However, we already know we will be able to cover + // these intervals together if we have appropriate projections. Note that any other + // intervals that cannot be covered may still require the query to use a FETCH. + tightness = IndexBoundsBuilder::EXACT_MAYBE_COVERED; } + IndexBoundsBuilder::_mergeTightness(tightness, *tightnessOut); } for (auto&& regex : ime->getRegexes()) { translateRegex(regex.get(), index, oilOut, &tightness); - if (tightness != IndexBoundsBuilder::EXACT) { - *tightnessOut = tightness; - } - } - - if (ime->hasNull()) { - // A null index key does not always match a null query value so we must fetch the - // doc and run a full comparison. See SERVER-4529. - // TODO: Do we already set the tightnessOut by calling translateEquality? - *tightnessOut = INEXACT_FETCH; - } - - if (ime->hasEmptyArray()) { - // Empty arrays are indexed as undefined. - BSONObjBuilder undefinedBob; - undefinedBob.appendUndefined(""); - oilOut->intervals.push_back(makePointInterval(undefinedBob.obj())); - *tightnessOut = IndexBoundsBuilder::INEXACT_FETCH; + IndexBoundsBuilder::_mergeTightness(tightness, *tightnessOut); } // Equalities are already sorted and deduped so unionize is unneccesary if no regexes // are present. Hashed indexes may also cause the bounds to be out-of-order. - // Arrays and nulls introduce multiple elements that neccesitate a sort and deduping. - if (!ime->getRegexes().empty() || index.type == IndexType::INDEX_HASHED || - arrayOrNullPresent) + // Arrays and nulls introduce multiple elements that necessitate a sort and deduping. + if (ime->hasNonScalarOrNonEmptyValues() || index.type == IndexType::INDEX_HASHED) { unionize(oilOut); + } + } else if (MatchExpression::GEO == expr->matchType()) { const GeoMatchExpression* gme = static_cast<const GeoMatchExpression*>(expr); if ("2dsphere" == elt.valueStringDataSafe()) { @@ -1315,6 +1392,7 @@ void IndexBoundsBuilder::translateEquality(const BSONElement& data, } std::sort(oil->intervals.begin(), oil->intervals.end(), IntervalComparison); + *tightnessOut = IndexBoundsBuilder::INEXACT_FETCH; } diff --git a/src/mongo/db/query/index_bounds_builder.h b/src/mongo/db/query/index_bounds_builder.h index b48228328f5..d1067caa561 100644 --- a/src/mongo/db/query/index_bounds_builder.h +++ b/src/mongo/db/query/index_bounds_builder.h @@ -67,16 +67,28 @@ public: * increasing tightness. These values are used when we need to do comparison between two * BoundsTightness values. Such comparisons can answer questions such as "Does predicate * X have tighter or looser bounds than predicate Y?". + * + * These enum values are ordered from loosest to tightest. */ enum BoundsTightness { // Index bounds are inexact, and a fetch is required. INEXACT_FETCH = 0, - // Index bounds are inexact, but no fetch is required - INEXACT_COVERED = 1, + // Index bounds are inexact, and a fetch may be required depending on the projection. + // For example, a count $in query on null + a regex can be covered, but a find query with + // the same filter and no projection cannot. + INEXACT_MAYBE_COVERED = 1, + + // Index bounds are exact, but a fetch may be required depending on the projection. + // For example, a find query on null may be covered, depending on which fields we project + // out. + EXACT_MAYBE_COVERED = 2, + + // Index bounds are inexact, but no fetch is required. + INEXACT_COVERED = 3, // Index bounds are exact. - EXACT = 2 + EXACT = 4 }; /** @@ -301,6 +313,11 @@ private: OrderedIntervalList* oilOut, BoundsTightness* tightnessOut, interval_evaluation_tree::Builder* ietBuilder); + + /** + * Helper method for merging interval tightness for $in expressions. + */ + static void _mergeTightness(const BoundsTightness& tightness, BoundsTightness& tightnessOut); }; } // namespace mongo diff --git a/src/mongo/db/query/index_bounds_builder_eq_null_test.cpp b/src/mongo/db/query/index_bounds_builder_eq_null_test.cpp index af4cdf91303..31e17b6e2c8 100644 --- a/src/mongo/db/query/index_bounds_builder_eq_null_test.cpp +++ b/src/mongo/db/query/index_bounds_builder_eq_null_test.cpp @@ -48,7 +48,7 @@ void assertBoundsRepresentEqualsNull(const OrderedIntervalList& oil) { oil.intervals[1].compare(Interval(fromjson("{'': null, '': null}"), true, true))); } -TEST_F(IndexBoundsBuilderTest, TranslateExprEqualToNullIsInexactFetch) { +TEST_F(IndexBoundsBuilderTest, TranslateExprEqualToNullIsExactMaybeCovered) { BSONObj keyPattern = BSON("a" << 1); BSONElement elt = keyPattern.firstElement(); auto testIndex = buildSimpleIndexEntry(keyPattern); @@ -65,11 +65,11 @@ TEST_F(IndexBoundsBuilderTest, TranslateExprEqualToNullIsInexactFetch) { oil.intervals[0].compare(Interval(fromjson("{'': undefined, '': undefined}"), true, true))); ASSERT_EQUALS(Interval::INTERVAL_EQUALS, oil.intervals[1].compare(Interval(fromjson("{'': null, '': null}"), true, true))); - ASSERT_EQUALS(tightness, IndexBoundsBuilder::INEXACT_FETCH); + ASSERT_EQUALS(tightness, IndexBoundsBuilder::EXACT_MAYBE_COVERED); assertIET(inputParamIdMap, ietBuilder, elt, testIndex, oil); } -TEST_F(IndexBoundsBuilderTest, TranslateEqualsToNullShouldBuildInexactBounds) { +TEST_F(IndexBoundsBuilderTest, TranslateEqualsToNullShouldBuildExactMaybeCoveredBounds) { BSONObj indexPattern = BSON("a" << 1); auto testIndex = buildSimpleIndexEntry(indexPattern); @@ -83,12 +83,12 @@ TEST_F(IndexBoundsBuilderTest, TranslateEqualsToNullShouldBuildInexactBounds) { expr.get(), indexPattern.firstElement(), testIndex, &oil, &tightness, &ietBuilder); ASSERT_EQUALS(oil.name, "a"); - ASSERT_EQUALS(tightness, IndexBoundsBuilder::INEXACT_FETCH); + ASSERT_EQUALS(tightness, IndexBoundsBuilder::EXACT_MAYBE_COVERED); assertBoundsRepresentEqualsNull(oil); assertIET(inputParamIdMap, ietBuilder, indexPattern.firstElement(), testIndex, oil); } -TEST_F(IndexBoundsBuilderTest, TranslateDottedEqualsToNullShouldBuildInexactBounds) { +TEST_F(IndexBoundsBuilderTest, TranslateDottedEqualsToNullShouldBuildExactMaybeCoveredBounds) { BSONObj indexPattern = BSON("a.b" << 1); auto testIndex = buildSimpleIndexEntry(indexPattern); @@ -102,7 +102,9 @@ TEST_F(IndexBoundsBuilderTest, TranslateDottedEqualsToNullShouldBuildInexactBoun expr.get(), indexPattern.firstElement(), testIndex, &oil, &tightness, &ietBuilder); ASSERT_EQUALS(oil.name, "a.b"); - ASSERT_EQUALS(tightness, IndexBoundsBuilder::INEXACT_FETCH); + // Depending on the query projection, this will either be converted to EXACT or to INEXACT_FETCH + // before we build an IXSCAN plan. + ASSERT_EQUALS(tightness, IndexBoundsBuilder::EXACT_MAYBE_COVERED); assertBoundsRepresentEqualsNull(oil); assertIET(inputParamIdMap, ietBuilder, indexPattern.firstElement(), testIndex, oil); } diff --git a/src/mongo/db/query/index_tag.cpp b/src/mongo/db/query/index_tag.cpp index 29c450d6ae2..ab5acea4de0 100644 --- a/src/mongo/db/query/index_tag.cpp +++ b/src/mongo/db/query/index_tag.cpp @@ -120,21 +120,33 @@ void sortUsingTags(MatchExpression* tree) { }); } -// Attaches 'node' to 'target'. If 'target' is an AND, adds 'node' as a child of 'target'. -// Otherwise, creates an AND that is a child of 'targetParent' at position 'targetPosition', and -// adds 'target' and 'node' as its children. Tags 'node' with 'tagData'. +/** + * Attaches 'node' to 'target'. If 'target' is an AND, adds 'node' as a child of 'target'. + * Otherwise, creates an AND that is a child of 'targetParent' at position 'targetPosition', and + * adds 'target' and 'node' as its children. Tags 'node' with 'tagData'. If 'node' appears as a key + * in 'pathsToUpdate', then we set the new path onto the clone. + */ void attachNode(MatchExpression* node, MatchExpression* target, OrMatchExpression* targetParent, size_t targetPosition, - std::unique_ptr<MatchExpression::TagData> tagData) { + std::unique_ptr<MatchExpression::TagData> tagData, + const stdx::unordered_map<MatchExpression*, FieldRef>& pathsToUpdate) { auto clone = node->shallowClone(); if (clone->matchType() == MatchExpression::NOT) { IndexTag* indexTag = static_cast<IndexTag*>(tagData.get()); clone->setTag(new IndexTag(indexTag->index)); clone->getChild(0)->setTag(tagData.release()); + + if (auto it = pathsToUpdate.find(node->getChild(0)); it != pathsToUpdate.end()) { + checked_cast<PathMatchExpression*>(clone->getChild(0)) + ->setPath(it->second.dottedField()); + } } else { clone->setTag(tagData.release()); + if (auto it = pathsToUpdate.find(node); it != pathsToUpdate.end()) { + checked_cast<PathMatchExpression*>(clone.get())->setPath(it->second.dottedField()); + } } if (MatchExpression::AND == target->matchType()) { @@ -164,17 +176,24 @@ stdx::unordered_map<size_t, std::vector<OrPushdownTag::Destination>> partitionCh return childDestinations; } -// Finds the node within 'tree' that is an indexed OR, if one exists. -MatchExpression* getIndexedOr(MatchExpression* tree) { +/** + * Finds the node within 'tree' that is an indexed OR, if one exists. It also returns the subpath in + * which the indexed OR lives. + */ +std::pair<MatchExpression*, FieldRef> getIndexedOr(FieldRef currentPath, MatchExpression* tree) { if (MatchExpression::OR == tree->matchType() && tree->getTag()) { - return tree; + return {tree, std::move(currentPath)}; + } + if (const auto* fieldRef = tree->fieldRef()) { + currentPath = currentPath + *fieldRef; } + for (size_t i = 0; i < tree->numChildren(); ++i) { - if (auto indexedOrChild = getIndexedOr(tree->getChild(i))) { - return indexedOrChild; + if (auto result = getIndexedOr(currentPath, tree->getChild(i)); result.first) { + return result; } } - return nullptr; + return {}; } // Pushes down 'node' along the routes in 'target' specified in 'destinations'. Each value in the @@ -182,7 +201,8 @@ MatchExpression* getIndexedOr(MatchExpression* tree) { // descendant of 'target'. bool pushdownNode(MatchExpression* node, MatchExpression* target, - std::vector<OrPushdownTag::Destination> destinations) { + std::vector<OrPushdownTag::Destination> destinations, + const stdx::unordered_map<MatchExpression*, FieldRef>& pathsToUpdate) { if (MatchExpression::OR == target->matchType()) { OrMatchExpression* orNode = static_cast<OrMatchExpression*>(target); bool moveToAllChildren = true; @@ -206,13 +226,15 @@ bool pushdownNode(MatchExpression* node, orNode->getChild(i), orNode, i, - std::move(childDestinations->second[0].tagData)); + std::move(childDestinations->second[0].tagData), + pathsToUpdate); } else { // This child was specified by a non-trivial route in destinations, so we recur. moveToAllChildren = pushdownNode(node, orNode->getChild(i), - std::move(childDestinations->second)) && + std::move(childDestinations->second), + pathsToUpdate) && moveToAllChildren; } } @@ -221,36 +243,81 @@ bool pushdownNode(MatchExpression* node, } if (MatchExpression::AND == target->matchType()) { - auto indexedOr = getIndexedOr(target); + auto [indexedOr, fieldRef_unused] = getIndexedOr({} /*fieldRef*/, target); invariant(indexedOr); - return pushdownNode(node, indexedOr, std::move(destinations)); + return pushdownNode(node, indexedOr, std::move(destinations), pathsToUpdate); } MONGO_UNREACHABLE_TASSERT(4457014); } -// Populates 'out' with all descendants of 'node' that have OrPushdownTags, assuming the initial -// input is an ELEM_MATCH_OBJECT. -void getElemMatchOrPushdownDescendants(MatchExpression* node, std::vector<MatchExpression*>* out) { +/** + * Populates 'out' with all descendants of 'node' that have OrPushdownTags, assuming the initial + * input is an ELEM_MATCH_OBJECT. The "currentPath" argument is the combined path traversed so far. + * Additionally, we populate a map to keep track of paths to update afterward during cloning. + */ +void getElemMatchOrPushdownDescendants( + const FieldRef& indexedOrPath, + FieldRef currentPath, + MatchExpression* node, + std::vector<MatchExpression*>* out, + stdx::unordered_map<MatchExpression*, FieldRef>* pathsToUpdate) { + const bool updatePath = node->fieldRef() != nullptr; + if (updatePath) { + currentPath = currentPath + *node->fieldRef(); + } + + // Do not do extra pushdown of OR inside $elemmatch. if (node->getTag() && node->getTag()->getType() == TagType::OrPushdownTag) { + if (updatePath) { + // Make sure that we remove the common prefix between the "destination" OR and the + // current expression, as it may be contained within the same $elemmatch. + + const auto prefixSize = indexedOrPath.commonPrefixSize(currentPath); + for (auto i = 0; i < prefixSize; i++) { + currentPath.removeFirstPart(); + } + if (currentPath != *node->fieldRef()) { + pathsToUpdate->emplace(node, std::move(currentPath)); + } + } out->push_back(node); } else if (node->matchType() == MatchExpression::ELEM_MATCH_OBJECT || node->matchType() == MatchExpression::AND) { for (size_t i = 0; i < node->numChildren(); ++i) { - getElemMatchOrPushdownDescendants(node->getChild(i), out); + getElemMatchOrPushdownDescendants( + indexedOrPath, currentPath, node->getChild(i), out, pathsToUpdate); } } else if (node->matchType() == MatchExpression::NOT) { // The immediate child of NOT may be tagged, but there should be no tags deeper than this. auto* childNode = node->getChild(0); if (childNode->getTag() && childNode->getTag()->getType() == TagType::OrPushdownTag) { + if (!childNode->path().empty()) { + // Make sure that we remove the common prefix between the "destination" OR and the + // current expression, as it may be contained within the same $elemmatch. + + currentPath = currentPath + *childNode->fieldRef(); + const auto prefixSize = indexedOrPath.commonPrefixSize(currentPath); + for (auto i = 0; i < prefixSize; i++) { + currentPath.removeFirstPart(); + } + if (currentPath != *childNode->fieldRef()) { + pathsToUpdate->emplace(childNode, std::move(currentPath)); + } + } out->push_back(node); } } } -// Attempts to push the given node down into the 'indexedOr' subtree. Returns true if the predicate -// can subsequently be trimmed from the MatchExpression tree, false otherwise. -bool processOrPushdownNode(MatchExpression* node, MatchExpression* indexedOr) { +/** + * Attempts to push the given node down into the 'indexedOr' subtree. Returns true if the predicate + * can subsequently be trimmed from the MatchExpression tree, false otherwise. Also supplied is a + * map to optionally update the path of the 'node' being pushed down. + */ +bool processOrPushdownNode(MatchExpression* node, + MatchExpression* indexedOr, + const stdx::unordered_map<MatchExpression*, FieldRef>& pathsToUpdate) { // If the node is a negation, then its child is the predicate node that may be tagged. auto* predNode = node->matchType() == MatchExpression::NOT ? node->getChild(0) : node; @@ -267,7 +334,7 @@ bool processOrPushdownNode(MatchExpression* node, MatchExpression* indexedOr) { predNode->setTag(nullptr); // Attempt to push the node into the indexedOr, then re-set its tag to the indexTag. - const bool pushedDown = pushdownNode(node, indexedOr, std::move(destinations)); + const bool pushedDown = pushdownNode(node, indexedOr, std::move(destinations), pathsToUpdate); predNode->setTag(indexTag.release()); // Return true if we can trim the predicate. We could trim the node even if it had an index tag @@ -284,20 +351,26 @@ void resolveOrPushdowns(MatchExpression* tree) { } if (MatchExpression::AND == tree->matchType()) { AndMatchExpression* andNode = static_cast<AndMatchExpression*>(tree); - MatchExpression* indexedOr = getIndexedOr(andNode); + auto [indexedOr, indexedOrPath] = getIndexedOr({} /*fieldRef*/, andNode); for (size_t i = 0; i < andNode->numChildren(); ++i) { auto child = andNode->getChild(i); - // For ELEM_MATCH_OBJECT, we push down all tagged descendants. However, we cannot trim - // any of these predicates, since the $elemMatch filter must be applied in its entirety. + // For ELEM_MATCH_OBJECT, we push down all tagged descendants. However, we cannot + // trim any of these predicates, since the $elemMatch filter must be applied in its + // entirety. if (child->matchType() == MatchExpression::ELEM_MATCH_OBJECT) { std::vector<MatchExpression*> orPushdownDescendants; - getElemMatchOrPushdownDescendants(child, &orPushdownDescendants); + stdx::unordered_map<MatchExpression*, FieldRef> pathsToUpdate; + getElemMatchOrPushdownDescendants(indexedOrPath, + {} /*currentPath*/, + child, + &orPushdownDescendants, + &pathsToUpdate); for (auto descendant : orPushdownDescendants) { - static_cast<void>(processOrPushdownNode(descendant, indexedOr)); + static_cast<void>(processOrPushdownNode(descendant, indexedOr, pathsToUpdate)); } - } else if (processOrPushdownNode(child, indexedOr)) { + } else if (processOrPushdownNode(child, indexedOr, {} /*pathsToUpdate*/)) { // The indexed $or can completely satisfy the child predicate, so we trim it. auto ownedChild = andNode->removeChild(i); --i; diff --git a/src/mongo/db/query/indexability.h b/src/mongo/db/query/indexability.h index 7e543da4a2d..48b9e0d91b4 100644 --- a/src/mongo/db/query/indexability.h +++ b/src/mongo/db/query/indexability.h @@ -55,6 +55,26 @@ public: } /** + * Type bracketing does not apply to internal Expressions. This could cause the use of a sparse + * index return incomplete results. For example, a query {$expr: {$lt: ["$missing", "r"]}} would + * expect a document like, {a: 1}, with field "missing" missing be returned. However, a sparse + * index, {missing: 1} does not index the document. Therefore, we should ban use of any sparse + * index on following expression types. + */ + static bool nodeSupportedBySparseIndex(const MatchExpression* me) { + switch (me->matchType()) { + case MatchExpression::INTERNAL_EXPR_EQ: + case MatchExpression::INTERNAL_EXPR_GT: + case MatchExpression::INTERNAL_EXPR_GTE: + case MatchExpression::INTERNAL_EXPR_LT: + case MatchExpression::INTERNAL_EXPR_LTE: + return false; + default: + return true; + } + } + + /** * This array operator doesn't have any children with fields and can use an index. * * Example: a: {$elemMatch: {$gte: 1, $lte: 1}}. diff --git a/src/mongo/db/query/internal_plans.cpp b/src/mongo/db/query/internal_plans.cpp index 04f70b1d2cc..78820eb8111 100644 --- a/src/mongo/db/query/internal_plans.cpp +++ b/src/mongo/db/query/internal_plans.cpp @@ -123,7 +123,8 @@ CollectionScanParams createCollectionScanParams( boost::optional<RecordId> resumeAfterRecordId, boost::optional<RecordIdBound> minRecord, boost::optional<RecordIdBound> maxRecord, - CollectionScanParams::ScanBoundInclusion boundInclusion) { + CollectionScanParams::ScanBoundInclusion boundInclusion, + bool shouldReturnEofOnFilterMismatch) { const auto& collection = *coll; invariant(collection); @@ -139,6 +140,7 @@ CollectionScanParams createCollectionScanParams( params.direction = CollectionScanParams::BACKWARD; } params.boundInclusion = boundInclusion; + params.shouldReturnEofOnFilterMismatch = shouldReturnEofOnFilterMismatch; return params; } } // namespace @@ -151,7 +153,8 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::collection boost::optional<RecordId> resumeAfterRecordId, boost::optional<RecordIdBound> minRecord, boost::optional<RecordIdBound> maxRecord, - CollectionScanParams::ScanBoundInclusion boundInclusion) { + CollectionScanParams::ScanBoundInclusion boundInclusion, + bool shouldReturnEofOnFilterMismatch) { const auto& collection = *coll; invariant(collection); @@ -167,7 +170,8 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::collection resumeAfterRecordId, minRecord, maxRecord, - boundInclusion); + boundInclusion, + shouldReturnEofOnFilterMismatch); auto cs = _collectionScan(expCtx, ws.get(), &collection, collScanParams); @@ -218,11 +222,19 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWith boost::optional<RecordIdBound> minRecord, boost::optional<RecordIdBound> maxRecord, CollectionScanParams::ScanBoundInclusion boundInclusion, - boost::optional<std::unique_ptr<BatchedDeleteStageBatchParams>> batchParams) { + std::unique_ptr<BatchedDeleteStageBatchParams> batchedDeleteParams, + const MatchExpression* filter, + bool shouldReturnEofOnFilterMismatch) { const auto& collection = *coll; invariant(collection); - auto ws = std::make_unique<WorkingSet>(); + if (shouldReturnEofOnFilterMismatch) { + tassert(7010801, + "MatchExpression filter must be provided when 'shouldReturnEofOnFilterMismatch' is " + "set to true ", + filter); + } + auto ws = std::make_unique<WorkingSet>(); auto expCtx = make_intrusive<ExpressionContext>( opCtx, std::unique_ptr<CollatorInterface>(nullptr), collection->ns()); @@ -237,14 +249,15 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWith boost::none /* resumeAfterId */, minRecord, maxRecord, - boundInclusion); + boundInclusion, + shouldReturnEofOnFilterMismatch); - auto root = _collectionScan(expCtx, ws.get(), &collection, collScanParams); + auto root = _collectionScan(expCtx, ws.get(), &collection, collScanParams, filter); - if (batchParams) { + if (batchedDeleteParams) { root = std::make_unique<BatchedDeleteStage>(expCtx.get(), std::move(params), - std::move(*batchParams), + std::move(batchedDeleteParams), ws.get(), collection, root.release()); @@ -454,12 +467,13 @@ std::unique_ptr<PlanStage> InternalPlanner::_collectionScan( const boost::intrusive_ptr<ExpressionContext>& expCtx, WorkingSet* ws, const CollectionPtr* coll, - const CollectionScanParams& params) { + const CollectionScanParams& params, + const MatchExpression* filter) { const auto& collection = *coll; invariant(collection); - return std::make_unique<CollectionScan>(expCtx.get(), collection, params, ws, nullptr); + return std::make_unique<CollectionScan>(expCtx.get(), collection, params, ws, filter); } std::unique_ptr<PlanStage> InternalPlanner::_indexScan( diff --git a/src/mongo/db/query/internal_plans.h b/src/mongo/db/query/internal_plans.h index ea8de7c0042..46400c4e3fa 100644 --- a/src/mongo/db/query/internal_plans.h +++ b/src/mongo/db/query/internal_plans.h @@ -83,7 +83,8 @@ public: boost::optional<RecordIdBound> minRecord = boost::none, boost::optional<RecordIdBound> maxRecord = boost::none, CollectionScanParams::ScanBoundInclusion boundInclusion = - CollectionScanParams::ScanBoundInclusion::kIncludeBothStartAndEndRecords); + CollectionScanParams::ScanBoundInclusion::kIncludeBothStartAndEndRecords, + bool shouldReturnEofOnFilterMismatch = false); static std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> collectionScan( OperationContext* opCtx, @@ -104,7 +105,9 @@ public: boost::optional<RecordIdBound> maxRecord = boost::none, CollectionScanParams::ScanBoundInclusion boundInclusion = CollectionScanParams::ScanBoundInclusion::kIncludeBothStartAndEndRecords, - boost::optional<std::unique_ptr<BatchedDeleteStageBatchParams>> batchParams = boost::none); + std::unique_ptr<BatchedDeleteStageBatchParams> batchedDeleteParams = nullptr, + const MatchExpression* filter = nullptr, + bool shouldReturnEofOnFilterMismatch = false); /** * Returns an index scan. Caller owns returned pointer. @@ -197,7 +200,8 @@ private: const boost::intrusive_ptr<ExpressionContext>& expCtx, WorkingSet* ws, const CollectionPtr* collection, - const CollectionScanParams& params); + const CollectionScanParams& params, + const MatchExpression* filter = nullptr); /** * Returns a plan stage that is either an index scan or an index scan with a fetch stage. diff --git a/src/mongo/db/query/multiple_collection_accessor.h b/src/mongo/db/query/multiple_collection_accessor.h index 26fc081000e..9d302ce2b28 100644 --- a/src/mongo/db/query/multiple_collection_accessor.h +++ b/src/mongo/db/query/multiple_collection_accessor.h @@ -101,6 +101,17 @@ public: _secondaryColls.clear(); } + void forEach(std::function<void(const CollectionPtr&)> func) const { + if (hasMainCollection()) { + func(getMainCollection()); + } + for (const auto& [name, coll] : getSecondaryCollections()) { + if (coll) { + func(coll); + } + } + } + private: const CollectionPtr* _mainColl{&CollectionPtr::null}; diff --git a/src/mongo/db/query/plan_cache_indexability.cpp b/src/mongo/db/query/plan_cache_indexability.cpp index 4f1b0c176f8..0498d5edc26 100644 --- a/src/mongo/db/query/plan_cache_indexability.cpp +++ b/src/mongo/db/query/plan_cache_indexability.cpp @@ -79,7 +79,6 @@ IndexabilityDiscriminator getCollatedIndexDiscriminator(const CollatorInterface* } return true; } - // The predicate never compares strings so it is not affected by collation. return true; }; @@ -104,14 +103,7 @@ void PlanCacheIndexabilityState::processSparseIndex(const std::string& indexName void PlanCacheIndexabilityState::processPartialIndex(const std::string& indexName, const MatchExpression* filterExpr) { - invariant(filterExpr); - for (size_t i = 0; i < filterExpr->numChildren(); ++i) { - processPartialIndex(indexName, filterExpr->getChild(i)); - } - if (filterExpr->getCategory() != MatchExpression::MatchCategory::kLogical) { - _pathDiscriminatorsMap[filterExpr->path()][indexName].addDiscriminator( - getPartialIndexDiscriminator(filterExpr)); - } + _globalDiscriminatorMap[indexName].addDiscriminator(getPartialIndexDiscriminator(filterExpr)); } void PlanCacheIndexabilityState::processWildcardIndex(const CoreIndexInfo& cii) { @@ -134,7 +126,7 @@ namespace { const IndexToDiscriminatorMap emptyDiscriminators{}; } // namespace -const IndexToDiscriminatorMap& PlanCacheIndexabilityState::getDiscriminators( +const IndexToDiscriminatorMap& PlanCacheIndexabilityState::getPathDiscriminators( StringData path) const { PathDiscriminatorsMap::const_iterator it = _pathDiscriminatorsMap.find(path); if (it == _pathDiscriminatorsMap.end()) { @@ -166,6 +158,7 @@ IndexToDiscriminatorMap PlanCacheIndexabilityState::buildWildcardDiscriminators( void PlanCacheIndexabilityState::updateDiscriminators( const std::vector<CoreIndexInfo>& indexCores) { _pathDiscriminatorsMap = PathDiscriminatorsMap(); + _globalDiscriminatorMap = IndexToDiscriminatorMap(); _wildcardIndexDiscriminators.clear(); for (const auto& idx : indexCores) { diff --git a/src/mongo/db/query/plan_cache_indexability.h b/src/mongo/db/query/plan_cache_indexability.h index 9bc03494865..0aa08359c27 100644 --- a/src/mongo/db/query/plan_cache_indexability.h +++ b/src/mongo/db/query/plan_cache_indexability.h @@ -47,6 +47,7 @@ class ProjectionExecutor; using IndexabilityDiscriminator = std::function<bool(const MatchExpression* me)>; using IndexabilityDiscriminators = std::vector<IndexabilityDiscriminator>; using IndexToDiscriminatorMap = StringMap<CompositeIndexabilityDiscriminator>; +using PathDiscriminatorsMap = StringMap<IndexToDiscriminatorMap>; /** * CompositeIndexabilityDiscriminator holds all indexability discriminators for a particular path, @@ -77,9 +78,14 @@ private: }; /** - * PlanCacheIndexabilityState holds a set of "indexability discriminators" for certain paths. - * An indexability discriminator is a binary predicate function, used to classify match - * expressions based on the data values in the expression. + * PlanCacheIndexabilityState holds a set of "indexability discriminators. An indexability + * discriminator is a binary predicate function, used to classify match expressions based on the + * data values in the expression. + * + * These discriminators are used to distinguish between queries of a similar shape but not the same + * candidate indexes. So each discriminator typically represents a decision like "is this index + * valid?" or "does this piece of the query disqualify it from using this index?". The output of + * these decisions is included in the plan cache key. */ class PlanCacheIndexabilityState { PlanCacheIndexabilityState(const PlanCacheIndexabilityState&) = delete; @@ -95,7 +101,15 @@ public: * The object returned by reference is valid until the next call to updateDiscriminators() or * until destruction of 'this', whichever is first. */ - const IndexToDiscriminatorMap& getDiscriminators(StringData path) const; + const IndexToDiscriminatorMap& getPathDiscriminators(StringData path) const; + + /** + * Returns a map of index name to discriminator set. These discriminators are not + * associated with a particular path of a query and apply to the entire MatchExpression. + */ + const IndexToDiscriminatorMap& getGlobalDiscriminators() const { + return _globalDiscriminatorMap; + } /** * Construct an IndexToDiscriminator map for the given path, only for the wildcard indexes @@ -109,8 +123,6 @@ public: void updateDiscriminators(const std::vector<CoreIndexInfo>& indexCores); private: - using PathDiscriminatorsMap = StringMap<IndexToDiscriminatorMap>; - /** * A $** index may index an infinite number of fields. We cannot just store a discriminator for * every possible field that it indexes, so we have to maintain some special context about the @@ -142,8 +154,8 @@ private: void processSparseIndex(const std::string& indexName, const BSONObj& keyPattern); /** - * Adds partial index discriminators for the partial index with the given filter expression - * to the discriminators for that index in '_pathDiscriminatorsMap'. + * Adds a global discriminator for the partial index with the given filter expression + * to the discriminators for that index in '_globalDiscriminatorMap'. * * A partial index discriminator distinguishes expressions that match a given partial index * predicate from expressions that don't match the partial index predicate. For example, @@ -174,6 +186,10 @@ private: // PathDiscriminatorsMap is a map from field path to index name to IndexabilityDiscriminator. PathDiscriminatorsMap _pathDiscriminatorsMap; + // Map from index name to global discriminators. These are discriminators which do not apply to + // a single path but the entire MatchExpression. + IndexToDiscriminatorMap _globalDiscriminatorMap; + std::vector<WildcardIndexDiscriminatorContext> _wildcardIndexDiscriminators; }; diff --git a/src/mongo/db/query/plan_cache_indexability_test.cpp b/src/mongo/db/query/plan_cache_indexability_test.cpp index af7677a8cd5..49f9fb79285 100644 --- a/src/mongo/db/query/plan_cache_indexability_test.cpp +++ b/src/mongo/db/query/plan_cache_indexability_test.cpp @@ -104,7 +104,7 @@ TEST(PlanCacheIndexabilityTest, SparseIndexSimple) { nullptr, nullptr)}); - auto discriminators = state.getDiscriminators("a"); + auto discriminators = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1") != discriminators.end()); @@ -146,7 +146,7 @@ TEST(PlanCacheIndexabilityTest, SparseIndexCompound) { nullptr)}); { - auto discriminators = state.getDiscriminators("a"); + auto discriminators = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1_b_1") != discriminators.end()); @@ -159,7 +159,7 @@ TEST(PlanCacheIndexabilityTest, SparseIndexCompound) { } { - auto discriminators = state.getDiscriminators("b"); + auto discriminators = state.getPathDiscriminators("b"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1_b_1") != discriminators.end()); @@ -193,12 +193,17 @@ TEST(PlanCacheIndexabilityTest, PartialIndexSimple) { nullptr, nullptr)}); + // The partial index is represented as a global discriminator that applies to the entire + // incoming MatchExpression. { - auto discriminators = state.getDiscriminators("f"); - ASSERT_EQ(1U, discriminators.size()); - ASSERT(discriminators.find("a_1") != discriminators.end()); + auto discriminators = state.getPathDiscriminators("f"); + ASSERT_EQ(0U, discriminators.size()); - auto disc = discriminators["a_1"]; + auto globalDiscriminators = state.getGlobalDiscriminators(); + ASSERT_EQ(1U, globalDiscriminators.size()); + ASSERT(globalDiscriminators.find("a_1") != globalDiscriminators.end()); + + auto disc = globalDiscriminators["a_1"]; ASSERT_EQ(false, disc.isMatchCompatibleWithIndex( parseMatchExpression(BSON("f" << BSON("$gt" << -5))).get())); @@ -208,7 +213,7 @@ TEST(PlanCacheIndexabilityTest, PartialIndexSimple) { } { - auto discriminators = state.getDiscriminators("a"); + auto discriminators = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1") != discriminators.end()); @@ -243,32 +248,52 @@ TEST(PlanCacheIndexabilityTest, PartialIndexAnd) { nullptr, nullptr)}); + // partial index discriminators are global to the entire query, so an individual path should not + // have any discriminators. Also the entire query must be a subset of the partial filter + // expression, not just the leaves. + auto globalDiscriminators = state.getGlobalDiscriminators(); + ASSERT(globalDiscriminators.find("a_1") != globalDiscriminators.end()); + auto globalDisc = globalDiscriminators["a_1"]; + { - auto discriminators = state.getDiscriminators("f"); - ASSERT_EQ(1U, discriminators.size()); - ASSERT(discriminators.find("a_1") != discriminators.end()); + auto discriminators = state.getPathDiscriminators("f"); + ASSERT_EQ(0U, discriminators.size()); - auto disc = discriminators["a_1"]; - ASSERT_EQ(false, - disc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 0)).get())); - ASSERT_EQ(true, - disc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 1)).get())); + ASSERT_EQ( + false, + globalDisc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 0)).get())); + ASSERT_EQ( + false, + globalDisc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 1)).get())); } { - auto discriminators = state.getDiscriminators("g"); - ASSERT_EQ(1U, discriminators.size()); - ASSERT(discriminators.find("a_1") != discriminators.end()); + auto discriminators = state.getPathDiscriminators("g"); + ASSERT_EQ(0U, discriminators.size()); - auto disc = discriminators["a_1"]; + ASSERT_EQ( + false, + globalDisc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("g" << 0)).get())); + ASSERT_EQ( + false, + globalDisc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("g" << 1)).get())); + } + + { + // A match expression which is covered entirely by the partial filter should pass the global + // discriminator. ASSERT_EQ(false, - disc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("g" << 0)).get())); + globalDisc.isMatchCompatibleWithIndex( + parseMatchExpression(BSON("g" << 1 << "f" << 0)).get())); ASSERT_EQ(true, - disc.isMatchCompatibleWithIndex(parseMatchExpression(BSON("g" << 1)).get())); + globalDisc.isMatchCompatibleWithIndex( + parseMatchExpression(BSON("g" << 1 << "f" << 1)).get())); } { - auto discriminators = state.getDiscriminators("a"); + // The path 'a' will still have a discriminator for the collation (even though it's + // defaulted). + auto discriminators = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1") != discriminators.end()); @@ -319,33 +344,44 @@ TEST(PlanCacheIndexabilityTest, MultiplePartialIndexes) { nullptr, nullptr)}); - { - auto discriminators = state.getDiscriminators("f"); - ASSERT_EQ(2U, discriminators.size()); - ASSERT(discriminators.find("a_1") != discriminators.end()); - ASSERT(discriminators.find("b_1") != discriminators.end()); + // partial index discriminators are global to the entire query, so an individual path within the + // partial filter should not have any discriminators. Also the entire query must be a subset of + // the partial filter expression, not just the leaves. + auto globalDiscriminators = state.getGlobalDiscriminators(); + ASSERT(globalDiscriminators.find("a_1") != globalDiscriminators.end()); + ASSERT(globalDiscriminators.find("b_1") != globalDiscriminators.end()); + auto globalDiscA = globalDiscriminators["a_1"]; + auto globalDiscB = globalDiscriminators["b_1"]; - auto discA = discriminators["a_1"]; - auto discB = discriminators["b_1"]; + { + auto discriminators = state.getPathDiscriminators("f"); + ASSERT_EQ(0U, discriminators.size()); - ASSERT_EQ(false, - discA.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 0)).get())); - ASSERT_EQ(false, - discB.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 0)).get())); + ASSERT_EQ( + false, + globalDiscA.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 0)).get())); + ASSERT_EQ( + false, + globalDiscB.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 0)).get())); - ASSERT_EQ(true, - discA.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 1)).get())); - ASSERT_EQ(false, - discB.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 1)).get())); + ASSERT_EQ( + true, + globalDiscA.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 1)).get())); + ASSERT_EQ( + false, + globalDiscB.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 1)).get())); - ASSERT_EQ(false, - discA.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 2)).get())); - ASSERT_EQ(true, - discB.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 2)).get())); + ASSERT_EQ( + false, + globalDiscA.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 2)).get())); + ASSERT_EQ( + true, + globalDiscB.isMatchCompatibleWithIndex(parseMatchExpression(BSON("f" << 2)).get())); } + // The paths 'a' and 'b' will have one discriminator each to capture the collation of the index. { - auto discriminators = state.getDiscriminators("a"); + auto discriminators = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1") != discriminators.end()); @@ -359,7 +395,7 @@ TEST(PlanCacheIndexabilityTest, MultiplePartialIndexes) { } { - auto discriminators = state.getDiscriminators("b"); + auto discriminators = state.getPathDiscriminators("b"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("b_1") != discriminators.end()); @@ -392,7 +428,7 @@ TEST(PlanCacheIndexabilityTest, IndexNeitherSparseNorPartial) { BSONObj(), nullptr, nullptr)}); - auto discriminators = state.getDiscriminators("a"); + auto discriminators = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1") != discriminators.end()); } @@ -421,7 +457,7 @@ TEST(PlanCacheIndexabilityTest, DiscriminatorForCollationIndicatesWhenCollations boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest()); expCtx->setCollator(collator.clone()); - auto discriminators = state.getDiscriminators("a"); + auto discriminators = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminators.size()); ASSERT(discriminators.find("a_1") != discriminators.end()); @@ -506,11 +542,11 @@ TEST(PlanCacheIndexabilityTest, CompoundIndexCollationDiscriminator) { nullptr, nullptr)}); - auto discriminatorsA = state.getDiscriminators("a"); + auto discriminatorsA = state.getPathDiscriminators("a"); ASSERT_EQ(1U, discriminatorsA.size()); ASSERT(discriminatorsA.find("a_1_b_1") != discriminatorsA.end()); - auto discriminatorsB = state.getDiscriminators("b"); + auto discriminatorsB = state.getPathDiscriminators("b"); ASSERT_EQ(1U, discriminatorsB.size()); ASSERT(discriminatorsB.find("a_1_b_1") != discriminatorsB.end()); } @@ -619,13 +655,15 @@ TEST(PlanCacheIndexabilityTest, WildcardPartialIndexDiscriminator) { ASSERT_TRUE(wildcardDiscriminators.isMatchCompatibleWithIndex( parseMatchExpression(fromjson("{b: 6}")).get())); - // The regular (non-wildcard) set of discriminators for the path "a" should reflect whether a - // predicate on "a" is compatible with the partial filter expression. + // The global discriminator for the index "indexName" should reflect whether a MatchExpression + // is compatible with the partial filter expression. { - discriminatorsA = state.getDiscriminators("a"); - auto discriminatorsIt = discriminatorsA.find("indexName"); - ASSERT(discriminatorsIt != discriminatorsA.end()); - auto disc = discriminatorsIt->second; + discriminatorsA = state.getPathDiscriminators("a"); + ASSERT(discriminatorsA.find("indexName") == discriminatorsA.end()); + + auto globalDisc = state.getGlobalDiscriminators(); + ASSERT(globalDisc.find("indexName") != globalDisc.end()); + auto disc = globalDisc["indexName"]; ASSERT_FALSE( disc.isMatchCompatibleWithIndex(parseMatchExpression(fromjson("{a: 0}")).get())); @@ -640,7 +678,7 @@ TEST(PlanCacheIndexabilityTest, WildcardPartialIndexDiscriminator) { // There shouldn't be any regular discriminators associated with path "b". { - auto&& discriminatorsB = state.getDiscriminators("b"); + auto&& discriminatorsB = state.getPathDiscriminators("b"); ASSERT_FALSE(discriminatorsB.count("indexName")); } } diff --git a/src/mongo/db/query/plan_cache_key_factory.cpp b/src/mongo/db/query/plan_cache_key_factory.cpp index d47f1768858..c5dbfbd7ea4 100644 --- a/src/mongo/db/query/plan_cache_key_factory.cpp +++ b/src/mongo/db/query/plan_cache_key_factory.cpp @@ -29,30 +29,29 @@ #include "mongo/db/query/plan_cache_key_factory.h" +#include "mongo/db/query/canonical_query_encoder.h" #include "mongo/db/query/collection_query_info.h" #include "mongo/db/query/planner_ixselect.h" #include "mongo/db/s/operation_sharding_state.h" namespace mongo { namespace plan_cache_detail { -// Delimiters for cache key encoding. -const char kEncodeDiscriminatorsBegin = '<'; -const char kEncodeDiscriminatorsEnd = '>'; void encodeIndexabilityForDiscriminators(const MatchExpression* tree, const IndexToDiscriminatorMap& discriminators, StringBuilder* keyBuilder) { + for (auto&& indexAndDiscriminatorPair : discriminators) { *keyBuilder << indexAndDiscriminatorPair.second.isMatchCompatibleWithIndex(tree); } } -void encodeIndexability(const MatchExpression* tree, - const PlanCacheIndexabilityState& indexabilityState, - StringBuilder* keyBuilder) { +void encodeIndexabilityRecursive(const MatchExpression* tree, + const PlanCacheIndexabilityState& indexabilityState, + StringBuilder* keyBuilder) { if (!tree->path().empty()) { const IndexToDiscriminatorMap& discriminators = - indexabilityState.getDiscriminators(tree->path()); + indexabilityState.getPathDiscriminators(tree->path()); IndexToDiscriminatorMap wildcardDiscriminators = indexabilityState.buildWildcardDiscriminators(tree->path()); if (!discriminators.empty() || !wildcardDiscriminators.empty()) { @@ -72,8 +71,26 @@ void encodeIndexability(const MatchExpression* tree, } for (size_t i = 0; i < tree->numChildren(); ++i) { - encodeIndexability(tree->getChild(i), indexabilityState, keyBuilder); + encodeIndexabilityRecursive(tree->getChild(i), indexabilityState, keyBuilder); + } +} + +void encodeIndexability(const MatchExpression* tree, + const PlanCacheIndexabilityState& indexabilityState, + StringBuilder* keyBuilder) { + // Before encoding the indexability of the leaf MatchExpressions, apply the global + // discriminators to the expression as a whole. This is for cases such as partial indexes which + // must discriminate based on the entire query. + const auto& globalDiscriminators = indexabilityState.getGlobalDiscriminators(); + if (!globalDiscriminators.empty()) { + *keyBuilder << kEncodeGlobalDiscriminatorsBegin; + for (auto&& indexAndDiscriminatorPair : globalDiscriminators) { + *keyBuilder << indexAndDiscriminatorPair.second.isMatchCompatibleWithIndex(tree); + } + *keyBuilder << kEncodeGlobalDiscriminatorsEnd; } + + encodeIndexabilityRecursive(tree, indexabilityState, keyBuilder); } PlanCacheKeyInfo makePlanCacheKeyInfo(const CanonicalQuery& query, @@ -111,8 +128,8 @@ boost::optional<Timestamp> computeNewestVisibleIndexTimestamp(OperationContext* Timestamp currentNewestVisible = Timestamp::min(); - std::unique_ptr<IndexCatalog::IndexIterator> ii = - collection->getIndexCatalog()->getIndexIterator(opCtx, /*includeUnfinishedIndexes*/ true); + auto ii = collection->getIndexCatalog()->getIndexIterator( + opCtx, IndexCatalog::InclusionPolicy::kReady | IndexCatalog::InclusionPolicy::kUnfinished); while (ii->more()) { const IndexCatalogEntry* ice = ii->next(); auto minVisibleSnapshot = ice->getMinimumVisibleSnapshot(); diff --git a/src/mongo/db/query/plan_cache_key_info_test.cpp b/src/mongo/db/query/plan_cache_key_info_test.cpp index 7235386e7f4..a13616e12c4 100644 --- a/src/mongo/db/query/plan_cache_key_info_test.cpp +++ b/src/mongo/db/query/plan_cache_key_info_test.cpp @@ -194,6 +194,107 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyPartialIndex) { makeKey(*cqGtZero, indexCores)); } +TEST(PlanCacheKeyInfoTest, ComputeKeyPartialIndexConjunction) { + BSONObj filterObj = fromjson("{f: {$gt: 0, $lt: 10}}"); + unique_ptr<MatchExpression> filterExpr(parseMatchExpression(filterObj)); + + const auto keyPattern = BSON("a" << 1); + const std::vector<CoreIndexInfo> indexCores = { + CoreIndexInfo(keyPattern, + IndexNames::nameToType(IndexNames::findPluginName(keyPattern)), + false, // sparse + IndexEntry::Identifier{""}, // name + filterExpr.get())}; // filterExpr + + unique_ptr<CanonicalQuery> satisfySinglePredicate(canonicalize("{f: {$gt: 0}}")); + ASSERT_EQ(makeKey(*satisfySinglePredicate, indexCores).getIndexabilityDiscriminators(), "(0)"); + + unique_ptr<CanonicalQuery> satisfyBothPredicates(canonicalize("{f: {$eq: 5}}")); + ASSERT_EQ(makeKey(*satisfyBothPredicates, indexCores).getIndexabilityDiscriminators(), "(1)"); + + unique_ptr<CanonicalQuery> conjSingleField(canonicalize("{f: {$gt: 2, $lt: 9}}")); + ASSERT_EQ(makeKey(*conjSingleField, indexCores).getIndexabilityDiscriminators(), "(1)"); + + unique_ptr<CanonicalQuery> conjSingleFieldNoMatch(canonicalize("{f: {$gt: 2, $lt: 11}}")); + ASSERT_EQ(makeKey(*conjSingleFieldNoMatch, indexCores).getIndexabilityDiscriminators(), "(0)"); + + // Note that these queries get optimized to a single $in over 'f'. + unique_ptr<CanonicalQuery> disjSingleFieldBothSatisfy( + canonicalize("{$or: [{f: {$eq: 2}}, {f: {$eq: 3}}]}")); + ASSERT_EQ(makeKey(*disjSingleFieldBothSatisfy, indexCores).getIndexabilityDiscriminators(), + "(1)"); + + unique_ptr<CanonicalQuery> disjSingleFieldNotSubset( + canonicalize("{$or: [{f: {$eq: 2}}, {f: {$eq: 11}}]}")); + ASSERT_EQ(makeKey(*disjSingleFieldNotSubset, indexCores).getIndexabilityDiscriminators(), + "(0)"); +} + +TEST(PlanCacheKeyInfoTest, ComputeKeyPartialIndexDisjunction) { + BSONObj filterObj = fromjson("{$or: [{f: {$gt: 10}}, {f: {$lt: 0}}]}"); + unique_ptr<MatchExpression> filterExpr(parseMatchExpression(filterObj)); + + const auto keyPattern = BSON("a" << 1); + const std::vector<CoreIndexInfo> indexCores = { + CoreIndexInfo(keyPattern, + IndexNames::nameToType(IndexNames::findPluginName(keyPattern)), + false, // sparse + IndexEntry::Identifier{""}, // name + filterExpr.get())}; // filterExpr + + unique_ptr<CanonicalQuery> satisfySinglePredicate(canonicalize("{f: {$eq: 11}}")); + ASSERT_EQ(makeKey(*satisfySinglePredicate, indexCores).getIndexabilityDiscriminators(), "(1)"); + + unique_ptr<CanonicalQuery> satisfyNeither(canonicalize("{f: {$eq: 5}}")); + ASSERT_EQ(makeKey(*satisfyNeither, indexCores).getIndexabilityDiscriminators(), "(0)"); + + unique_ptr<CanonicalQuery> conjSingleFieldMatch(canonicalize("{f: {$lt: 20, $gt: 10}}")); + ASSERT_EQ(makeKey(*conjSingleFieldMatch, indexCores).getIndexabilityDiscriminators(), "(1)"); + + unique_ptr<CanonicalQuery> conjSingleFieldNoMatch(canonicalize("{f: {$gt: 2, $lt: 10}}")); + ASSERT_EQ(makeKey(*conjSingleFieldNoMatch, indexCores).getIndexabilityDiscriminators(), "(0)"); + + unique_ptr<CanonicalQuery> conjSingleFieldOverlap(canonicalize("{f: {$gt: 2, $lt: 12}}")); + ASSERT_EQ(makeKey(*conjSingleFieldOverlap, indexCores).getIndexabilityDiscriminators(), "(0)"); + + // Although this query is technically a subset of the partial filter, the logic to determine + // such ('isSubsetOf' in the code) is conservative in how it compares certain shapes of + // expression trees. + unique_ptr<CanonicalQuery> disjSingleFieldBothSatisfy( + canonicalize("{$or: [{f: {$eq: -1}}, {f: {$gt: 10}}]}")); + ASSERT_EQ(makeKey(*disjSingleFieldBothSatisfy, indexCores).getIndexabilityDiscriminators(), + "(0)"); + + unique_ptr<CanonicalQuery> disjSingleFieldNotSubset( + canonicalize("{$or: [{f: {$eq: 2}}, {f: {$eq: 11}}]}")); + ASSERT_EQ(makeKey(*disjSingleFieldNotSubset, indexCores).getIndexabilityDiscriminators(), + "(0)"); +} + +TEST(PlanCacheKeyInfoTest, ComputeKeyPartialIndexNestedDisjunction) { + BSONObj filterObj = fromjson(R"( + {$and: [ + {$or: [{f: {$gt: 10}}, {f: {$lt: 0}}]}, + {$or: [{f: {$gt: 11}}, {f: {$lt: 1}}]} + ]})"); + unique_ptr<MatchExpression> filterExpr(parseMatchExpression(filterObj)); + + const auto keyPattern = BSON("a" << 1); + const std::vector<CoreIndexInfo> indexCores = { + CoreIndexInfo(keyPattern, + IndexNames::nameToType(IndexNames::findPluginName(keyPattern)), + false, // sparse + IndexEntry::Identifier{""}, // name + filterExpr.get())}; // filterExpr + + + unique_ptr<CanonicalQuery> satisfySinglePredicate(canonicalize("{f: {$eq: 11}}")); + ASSERT_EQ(makeKey(*satisfySinglePredicate, indexCores).getIndexabilityDiscriminators(), "(0)"); + + unique_ptr<CanonicalQuery> notCompat(canonicalize("{f: {$eq: 12}}")); + ASSERT_EQ(makeKey(*notCompat, indexCores).getIndexabilityDiscriminators(), "(1)"); +} + // Query shapes should get the same plan cache key if they have the same collation indexability. TEST(PlanCacheKeyInfoTest, ComputeKeyCollationIndex) { CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString); @@ -366,8 +467,8 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyWildcardDiscriminatesCorrectlyBasedOnPartia // The discriminator strings have the format "<xx>". That is, there are two discriminator // bits for the "x" predicate, the first pertaining to the partialFilterExpression and the // second around applicability to the wildcard index. - ASSERT_EQ(compatibleKey.getIndexabilityDiscriminators(), "<11>"); - ASSERT_EQ(incompatibleKey.getIndexabilityDiscriminators(), "<01>"); + ASSERT_EQ(compatibleKey.getIndexabilityDiscriminators(), "(1)<1>"); + ASSERT_EQ(incompatibleKey.getIndexabilityDiscriminators(), "(0)<1>"); } // The partialFilterExpression should lead to a discriminator over field 'x', but not over 'y'. @@ -382,8 +483,8 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyWildcardDiscriminatesCorrectlyBasedOnPartia // The discriminator strings have the format "<xx><y>". That is, there are two discriminator // bits for the "x" predicate (the first pertaining to the partialFilterExpression, the // second around applicability to the wildcard index) and one discriminator bit for "y". - ASSERT_EQ(compatibleKey.getIndexabilityDiscriminators(), "<11><1>"); - ASSERT_EQ(incompatibleKey.getIndexabilityDiscriminators(), "<01><1>"); + ASSERT_EQ(compatibleKey.getIndexabilityDiscriminators(), "(1)<1><1>"); + ASSERT_EQ(incompatibleKey.getIndexabilityDiscriminators(), "(0)<1><1>"); } // $eq:null predicates cannot be assigned to a wildcard index. Make sure that this is @@ -398,8 +499,8 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyWildcardDiscriminatesCorrectlyBasedOnPartia // The discriminator strings have the format "<xx><y>". That is, there are two discriminator // bits for the "x" predicate (the first pertaining to the partialFilterExpression, the // second around applicability to the wildcard index) and one discriminator bit for "y". - ASSERT_EQ(compatibleKey.getIndexabilityDiscriminators(), "<11><1>"); - ASSERT_EQ(incompatibleKey.getIndexabilityDiscriminators(), "<11><0>"); + ASSERT_EQ(compatibleKey.getIndexabilityDiscriminators(), "(1)<1><1>"); + ASSERT_EQ(incompatibleKey.getIndexabilityDiscriminators(), "(1)<1><0>"); } // Test that the discriminators are correct for an $eq:null predicate on 'x'. This predicate is @@ -408,7 +509,7 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyWildcardDiscriminatesCorrectlyBasedOnPartia // result in two "0" bits inside the discriminator string. { auto key = makeKey(*canonicalize("{x: {$eq: null}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<00>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(0)<0>"); } } @@ -449,11 +550,11 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyWildcardDiscriminatesCorrectlyWithPartialFi const std::vector<CoreIndexInfo> indexCores = {indexInfo}; { - // The discriminators should have the format <xx><yy><z>. The 'z' predicate has just one - // discriminator because it is not referenced in the partial filter expression. All + // TODO update The discriminators should have the format <xx><yy><z>. The 'z' predicate has + // just one discriminator because it is not referenced in the partial filter expression. All // predicates are compatible. auto key = makeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, z: {$eq: 3}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<11><11><1>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(1)<1><1><1>"); } { @@ -461,7 +562,7 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyWildcardDiscriminatesCorrectlyWithPartialFi // compatible with the partial filter expression, leading to one of the 'y' bits being set // to zero. auto key = makeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: -2}, z: {$eq: 3}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<11><01><1>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(0)<1><1><1>"); } } @@ -480,20 +581,20 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyDiscriminatesCorrectlyWithPartialFilterAndW // the predicate is compatible with the partial filter expression, whereas the disciminator // for 'y' is about compatibility with the wildcard index. auto key = makeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, z: {$eq: 3}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<1><1>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(1)<1>"); } { // Similar to the previous case, except with an 'x' predicate that is incompatible with the // partial filter expression. auto key = makeKey(*canonicalize("{x: {$eq: -1}, y: {$eq: 2}, z: {$eq: 3}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<0><1>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(0)<1>"); } { // Case where the 'y' predicate is not compatible with the wildcard index. auto key = makeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: null}, z: {$eq: 3}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<1><0>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(1)<0>"); } } @@ -511,14 +612,14 @@ TEST(PlanCacheKeyInfoTest, ComputeKeyWildcardDiscriminatesCorrectlyWithPartialFi // The discriminators have the format <x><(x.y)(x.y)<y>. All predicates are compatible auto key = makeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, 'x.y': {$eq: 3}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<1><11><1>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(1)<1><1><1>"); } { // Here, the predicate on "x.y" is not compatible with the partial filter expression. auto key = makeKey(*canonicalize("{x: {$eq: 1}, y: {$eq: 2}, 'x.y': {$eq: -3}}"), indexCores); - ASSERT_EQ(key.getIndexabilityDiscriminators(), "<1><01><1>"); + ASSERT_EQ(key.getIndexabilityDiscriminators(), "(0)<1><1><1>"); } } diff --git a/src/mongo/db/query/plan_enumerator.cpp b/src/mongo/db/query/plan_enumerator.cpp index a11bc9b415a..96fc70fb99c 100644 --- a/src/mongo/db/query/plan_enumerator.cpp +++ b/src/mongo/db/query/plan_enumerator.cpp @@ -1689,8 +1689,8 @@ bool PlanEnumerator::LockstepOrAssignment::allIdentical() const { return true; } -bool PlanEnumerator::LockstepOrAssignment::shouldResetBeforeProceeding( - size_t totalEnumerated) const { +bool PlanEnumerator::LockstepOrAssignment::shouldResetBeforeProceeding(size_t totalEnumerated, + size_t orLimit) const { if (totalEnumerated == 0 || !exhaustedLockstepIteration) { return false; } @@ -1700,7 +1700,12 @@ bool PlanEnumerator::LockstepOrAssignment::shouldResetBeforeProceeding( if (!subnode.maxIterCount) { return false; // Haven't yet looped over this child entirely, not ready yet. } - totalPossibleEnumerations *= subnode.maxIterCount.get(); + totalPossibleEnumerations *= subnode.maxIterCount.value(); + // If 'totalPossibleEnumerations' reaches the limit, we can just shortcut it. Otherwise, + // 'totalPossibleEnumerations' could overflow if we have a large $or. + if (totalPossibleEnumerations >= orLimit) { + return false; + } } // If we're able to compute a total number expected enumerations, we must have already cycled @@ -1737,7 +1742,7 @@ bool PlanEnumerator::_nextMemoForLockstepOrAssignment( } // Edge case: if every child has only one option available, we are already finished // enumerating. - if (assignment->shouldResetBeforeProceeding(assignment->totalEnumerated)) { + if (assignment->shouldResetBeforeProceeding(assignment->totalEnumerated, _orLimit)) { assignment->exhaustedLockstepIteration = false; return true; // We're back at the beginning, no need to reset. } @@ -1776,7 +1781,7 @@ bool PlanEnumerator::_nextMemoForLockstepOrAssignment( // This special ordering is tricky to reset. Because it iterates the sub nodes in such a // unique order, it can be difficult to know when it has actually finished iterating. Our // strategy is just to compute a total and go back to the beginning once we hit that total. - if (!assignment->shouldResetBeforeProceeding(assignment->totalEnumerated)) { + if (!assignment->shouldResetBeforeProceeding(assignment->totalEnumerated, _orLimit)) { return false; } // Reset! diff --git a/src/mongo/db/query/plan_enumerator.h b/src/mongo/db/query/plan_enumerator.h index 60344f0c9ee..b82b738c57b 100644 --- a/src/mongo/db/query/plan_enumerator.h +++ b/src/mongo/db/query/plan_enumerator.h @@ -244,7 +244,7 @@ private: * Returns true if 'totalEnumerated' matches the total number of expected plans for this * assignment. */ - bool shouldResetBeforeProceeding(size_t totalEnumerated) const; + bool shouldResetBeforeProceeding(size_t totalEnumerated, size_t orLimit) const; /** * Returns true if each sub node is at the same iterationCount. diff --git a/src/mongo/db/query/plan_executor.cpp b/src/mongo/db/query/plan_executor.cpp index ee41d15d84c..99b2fd8fefa 100644 --- a/src/mongo/db/query/plan_executor.cpp +++ b/src/mongo/db/query/plan_executor.cpp @@ -38,6 +38,10 @@ namespace { MONGO_FAIL_POINT_DEFINE(planExecutorAlwaysFails); } // namespace +const OperationContext::Decoration<boost::optional<SharedSemiFuture<void>>> + planExecutorShardingCriticalSectionFuture = + OperationContext::declareDecoration<boost::optional<SharedSemiFuture<void>>>(); + std::string PlanExecutor::stateToStr(ExecState execState) { switch (execState) { case PlanExecutor::ADVANCED: diff --git a/src/mongo/db/query/plan_executor.h b/src/mongo/db/query/plan_executor.h index 33fbd075b93..30ba3d69c63 100644 --- a/src/mongo/db/query/plan_executor.h +++ b/src/mongo/db/query/plan_executor.h @@ -52,8 +52,21 @@ class RecordId; * 'clientsLastKnownCommittedOpTime' represents the time passed to the getMore command. * If the replication coordinator ever reports a higher committed op time, we should stop waiting * for inserts and return immediately to speed up the propagation of commit level changes. + * + * A boost::none value opts out of the commit point propagation. A null optime compares less than + * any non-null optimes and thus will always trigger an empty batch for commit point propagation. + */ +extern const OperationContext::Decoration<boost::optional<repl::OpTime>> + clientsLastKnownCommittedOpTime; + +/** + * If a plan yielded because it encountered a sharding critical section, + * 'planExecutorShardingCriticalSectionFuture' will be set to a future that becomes ready when the + * critical section ends. This future can be waited on to hold off resuming the plan execution while + * the critical section is still active. */ -extern const OperationContext::Decoration<repl::OpTime> clientsLastKnownCommittedOpTime; +extern const OperationContext::Decoration<boost::optional<SharedSemiFuture<void>>> + planExecutorShardingCriticalSectionFuture; /** * A PlanExecutor is the abstraction that knows how to crank a tree of stages into execution. diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp index 6691a52fb8a..a7198fdcec1 100644 --- a/src/mongo/db/query/plan_executor_impl.cpp +++ b/src/mongo/db/query/plan_executor_impl.cpp @@ -37,6 +37,7 @@ #include "mongo/bson/simple_bsonobj_comparator.h" #include "mongo/db/catalog/collection.h" +#include "mongo/db/concurrency/exception_util.h" #include "mongo/db/concurrency/write_conflict_exception.h" #include "mongo/db/curop.h" #include "mongo/db/exec/cached_plan.h" @@ -61,6 +62,7 @@ #include "mongo/db/query/plan_yield_policy_impl.h" #include "mongo/db/query/yield_policy_callbacks_impl.h" #include "mongo/db/repl/replication_coordinator.h" +#include "mongo/db/s/operation_sharding_state.h" #include "mongo/db/service_context.h" #include "mongo/logv2/log.h" #include "mongo/util/fail_point.h" @@ -74,8 +76,8 @@ using std::string; using std::unique_ptr; using std::vector; -const OperationContext::Decoration<repl::OpTime> clientsLastKnownCommittedOpTime = - OperationContext::declareDecoration<repl::OpTime>(); +const OperationContext::Decoration<boost::optional<repl::OpTime>> clientsLastKnownCommittedOpTime = + OperationContext::declareDecoration<boost::optional<repl::OpTime>>(); // This failpoint is also accessed by the SBE executor so we define it outside of an anonymous // namespace. @@ -359,8 +361,25 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob // 2) some stage requested a yield, or // 3) we need to yield and retry due to a WriteConflictException. // In all cases, the actual yielding happens here. + + const auto whileYieldingFn = [&]() { + // If we yielded because we encountered a sharding critical section, wait for the + // critical section to end before continuing. By waiting for the critical section to be + // exited we avoid busy spinning immediately and encountering the same critical section + // again. It is important that this wait happens after having released the lock + // hierarchy -- otherwise deadlocks could happen, or the very least, locks would be + // unnecessarily held while waiting. + const auto& shardingCriticalSection = planExecutorShardingCriticalSectionFuture(_opCtx); + if (shardingCriticalSection) { + OperationShardingState::waitForCriticalSectionToComplete(_opCtx, + *shardingCriticalSection) + .ignore(); + planExecutorShardingCriticalSectionFuture(_opCtx).reset(); + } + }; + if (_yieldPolicy->shouldYieldOrInterrupt(_opCtx)) { - uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx)); + uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx, whileYieldingFn)); } WorkingSetID id = WorkingSet::INVALID_ID; @@ -427,8 +446,7 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob CurOp::get(_opCtx)->debug().additiveMetrics.incrementWriteConflicts(1); writeConflictsInARow++; - WriteConflictException::logAndBackoff( - writeConflictsInARow, "plan execution", _nss.ns()); + logWriteConflictAndBackoff(writeConflictsInARow, "plan execution", _nss.ns()); // If we're allowed to, we will yield next time through the loop. if (_yieldPolicy->canAutoYield()) { diff --git a/src/mongo/db/query/plan_explainer_impl.cpp b/src/mongo/db/query/plan_explainer_impl.cpp index 1e34cf73ebe..c36bc330826 100644 --- a/src/mongo/db/query/plan_explainer_impl.cpp +++ b/src/mongo/db/query/plan_explainer_impl.cpp @@ -668,6 +668,9 @@ void PlanExplainerImpl::getSummaryStats(PlanSummaryStats* statsOut) const { statsOut->totalKeysExamined = 0; statsOut->totalDocsExamined = 0; + statsOut->indexesUsed.clear(); + statsOut->collectionScans = 0; + statsOut->collectionScansNonTailable = 0; for (size_t i = 0; i < stages.size(); i++) { statsOut->totalKeysExamined += diff --git a/src/mongo/db/query/plan_explainer_sbe.cpp b/src/mongo/db/query/plan_explainer_sbe.cpp index 2f8f5b43b08..25ef5435572 100644 --- a/src/mongo/db/query/plan_explainer_sbe.cpp +++ b/src/mongo/db/query/plan_explainer_sbe.cpp @@ -371,9 +371,10 @@ void PlanExplainerSBE::getSummaryStats(PlanSummaryStats* statsOut) const { // Use the pre-computed summary stats instead of traversing the QuerySolution tree. const auto& indexesUsed = _debugInfo->mainStats.indexesUsed; + statsOut->indexesUsed.clear(); statsOut->indexesUsed.insert(indexesUsed.begin(), indexesUsed.end()); - statsOut->collectionScans += _debugInfo->mainStats.collectionScans; - statsOut->collectionScansNonTailable += _debugInfo->mainStats.collectionScansNonTailable; + statsOut->collectionScans = _debugInfo->mainStats.collectionScans; + statsOut->collectionScansNonTailable = _debugInfo->mainStats.collectionScansNonTailable; } void PlanExplainerSBE::getSecondarySummaryStats(std::string secondaryColl, diff --git a/src/mongo/db/query/plan_insert_listener.cpp b/src/mongo/db/query/plan_insert_listener.cpp index 0d86c76c9e7..1b7270eef76 100644 --- a/src/mongo/db/query/plan_insert_listener.cpp +++ b/src/mongo/db/query/plan_insert_listener.cpp @@ -65,9 +65,10 @@ bool shouldWaitForInserts(OperationContext* opCtx, // coordinator's lastCommittedOpTime has progressed past the client's lastCommittedOpTime. // In that case, we will return early so that we can inform the client of the new // lastCommittedOpTime immediately. - if (!clientsLastKnownCommittedOpTime(opCtx).isNull()) { + if (clientsLastKnownCommittedOpTime(opCtx)) { auto replCoord = repl::ReplicationCoordinator::get(opCtx); - return clientsLastKnownCommittedOpTime(opCtx) >= replCoord->getLastCommittedOpTime(); + return clientsLastKnownCommittedOpTime(opCtx).value() >= + replCoord->getLastCommittedOpTime(); } return true; } diff --git a/src/mongo/db/query/plan_yield_policy.cpp b/src/mongo/db/query/plan_yield_policy.cpp index 58064f76d6e..b19edc274f1 100644 --- a/src/mongo/db/query/plan_yield_policy.cpp +++ b/src/mongo/db/query/plan_yield_policy.cpp @@ -32,6 +32,7 @@ #include "mongo/db/query/plan_yield_policy.h" #include "mongo/db/catalog/collection.h" +#include "mongo/db/concurrency/exception_util.h" #include "mongo/db/concurrency/write_conflict_exception.h" #include "mongo/db/operation_context.h" #include "mongo/util/scopeguard.h" @@ -130,7 +131,7 @@ Status PlanYieldPolicy::yieldOrInterrupt(OperationContext* opCtx, if (_callbacks) { _callbacks->handledWriteConflict(opCtx); } - WriteConflictException::logAndBackoff(attempt, "query yield", ""_sd); + logWriteConflictAndBackoff(attempt, "query yield", ""_sd); // Retry the yielding process. } catch (...) { // Errors other than write conflicts don't get retried, and should instead result in diff --git a/src/mongo/db/query/planner_access.cpp b/src/mongo/db/query/planner_access.cpp index e05dfd8a5d1..ef0956dc56f 100644 --- a/src/mongo/db/query/planner_access.cpp +++ b/src/mongo/db/query/planner_access.cpp @@ -228,27 +228,37 @@ bool affectedByCollator(const BSONElement& element) { } } -void setMinRecord(CollectionScanNode* collScan, const BSONObj& min) { - const auto newMinRecord = record_id_helpers::keyForObj(min); - if (!collScan->minRecord || newMinRecord > collScan->minRecord->recordId()) { - collScan->minRecord = RecordIdBound(newMinRecord, min); +// Set 'curr' to 'newMin' if 'newMin' < 'curr' +void setLowestRecord(boost::optional<RecordIdBound>& curr, const RecordIdBound& newMin) { + if (!curr || newMin.recordId() < curr->recordId()) { + curr = newMin; } } -void setMaxRecord(CollectionScanNode* collScan, const BSONObj& max) { - const auto newMaxRecord = record_id_helpers::keyForObj(max); - if (!collScan->maxRecord || newMaxRecord < collScan->maxRecord->recordId()) { - collScan->maxRecord = RecordIdBound(newMaxRecord, max); +// Set 'curr' to 'newMax' if 'newMax' > 'curr' +void setHighestRecord(boost::optional<RecordIdBound>& curr, const RecordIdBound& newMax) { + if (!curr || newMax.recordId() > curr->recordId()) { + curr = newMax; } } +// Set 'curr' to 'newMin' if 'newMin' < 'curr' +void setLowestRecord(boost::optional<RecordIdBound>& curr, const BSONObj& newMin) { + setLowestRecord(curr, RecordIdBound(record_id_helpers::keyForObj(newMin), newMin)); +} + +// Set 'curr' to 'newMax' if 'newMax' > 'curr' +void setHighestRecord(boost::optional<RecordIdBound>& curr, const BSONObj& newMax) { + setHighestRecord(curr, RecordIdBound(record_id_helpers::keyForObj(newMax), newMax)); +} + // Returns whether element is not affected by collators or query and collection collators are // compatible. bool compatibleCollator(const QueryPlannerParams& params, const CollatorInterface* queryCollator, const BSONElement& element) { auto const collCollator = params.clusteredCollectionCollator; - bool compatible = !queryCollator || (collCollator && *queryCollator == *collCollator); + bool compatible = CollatorInterface::collatorsMatch(queryCollator, collCollator); return compatible || !affectedByCollator(element); } @@ -281,47 +291,106 @@ void handleRIDRangeMinMax(const CanonicalQuery& query, // Assumes clustered collection scans are only supported with the forward direction. collScan->boundInclusion = CollectionScanParams::ScanBoundInclusion::kIncludeStartRecordOnly; - setMaxRecord(collScan, IndexBoundsBuilder::objFromElement(maxObj.firstElement(), collator)); + setLowestRecord(collScan->maxRecord, + IndexBoundsBuilder::objFromElement(maxObj.firstElement(), collator)); } if (!minObj.isEmpty() && compatibleCollator(params, collator, minObj.firstElement())) { // The min() is inclusive as are bounded collection scans by default. - setMinRecord(collScan, IndexBoundsBuilder::objFromElement(minObj.firstElement(), collator)); + setHighestRecord(collScan->minRecord, + IndexBoundsBuilder::objFromElement(minObj.firstElement(), collator)); } } /** * Helper function to add an RID range to collection scans. - * If the query solution tree contains a collection scan node with a suitable comparison - * predicate on '_id', we add a minRecord and maxRecord on the collection node. + * If the query solution tree contains a collection scan node with a suitable comparison predicate + * on '_id', we add a minRecord and maxRecord on the collection node. + * + * Returns true if the MatchExpression is a comparison against the cluster key which either: + * 1) is guaranteed to exclude values of the cluster key which are affected by collation or + * 2) may return values of the cluster key which are affected by collation, but the query and + * collection collations match. + * Otherwise, returns false. + * + * For example, assuming the cluster key is "_id": + * Given {a: {$eq: 2}}, we return false, because the comparison is not against the cluster key. + * Given {_id: {$gte: 5}}, we return true, because this comparison against the cluster key excludes + * keys which are affected by collations. + * Given {_id: {$eq: "str"}}, we return true only if the query and collection collations match. + * */ -void handleRIDRangeScan(const MatchExpression* conjunct, - CollectionScanNode* collScan, - const QueryPlannerParams& params, - const CollatorInterface* collator) { +[[nodiscard]] bool handleRIDRangeScan(const MatchExpression* conjunct, + CollectionScanNode* collScan, + const QueryPlannerParams& params, + const CollatorInterface* collator) { invariant(params.clusteredInfo); if (conjunct == nullptr) { - return; + return false; } auto* andMatchPtr = dynamic_cast<const AndMatchExpression*>(conjunct); if (andMatchPtr != nullptr) { + bool atLeastOneConjunctCompatibleCollation = false; for (size_t index = 0; index < andMatchPtr->numChildren(); index++) { - handleRIDRangeScan(andMatchPtr->getChild(index), collScan, params, collator); + if (handleRIDRangeScan(andMatchPtr->getChild(index), collScan, params, collator)) { + atLeastOneConjunctCompatibleCollation = true; + } } - return; + + // If one of the conjuncts excludes values of the cluster key which are affected by + // collation, then the entire $and will also exclude those values. + return atLeastOneConjunctCompatibleCollation; } if (conjunct->path() != clustered_util::getClusterKeyFieldName(params.clusteredInfo->getIndexSpec())) { // No match on the cluster key. - return; + return false; + } + + // TODO SERVER-62707: Allow $in with regex to use a clustered index. + auto inMatch = dynamic_cast<const InMatchExpression*>(conjunct); + if (inMatch && !inMatch->hasRegex()) { + // Iterate through the $in equalities to find the min/max values. The min/max bounds for the + // collscan need to be loose enough to cover all of these values. + boost::optional<RecordIdBound> minBound; + boost::optional<RecordIdBound> maxBound; + + bool allEltsCollationCompatible = true; + for (const auto& element : inMatch->getEqualities()) { + if (compatibleCollator(params, collator, element)) { + const auto collated = IndexBoundsBuilder::objFromElement(element, collator); + setLowestRecord(minBound, collated); + setHighestRecord(maxBound, collated); + } else { + // Set coarse min/max bounds based on type when we can't set tight bounds. + allEltsCollationCompatible = false; + + BSONObjBuilder bMin; + bMin.appendMinForType("", element.type()); + setLowestRecord(minBound, bMin.obj()); + + BSONObjBuilder bMax; + bMax.appendMaxForType("", element.type()); + setHighestRecord(maxBound, bMax.obj()); + } + } + + // Finally, tighten the collscan bounds with the min/max bounds for the $in. + if (minBound) { + setHighestRecord(collScan->minRecord, *minBound); + } + if (maxBound) { + setLowestRecord(collScan->maxRecord, *maxBound); + } + return allEltsCollationCompatible; } auto match = dynamic_cast<const ComparisonMatchExpression*>(conjunct); if (match == nullptr) { - return; // Not a comparison match expression. + return false; // Not a comparison match expression. } const auto& element = match->getData(); @@ -329,32 +398,32 @@ void handleRIDRangeScan(const MatchExpression* conjunct, // Set coarse min/max bounds based on type in case we can't set tight bounds. BSONObjBuilder minb; minb.appendMinForType("", element.type()); - setMinRecord(collScan, minb.obj()); + setHighestRecord(collScan->minRecord, minb.obj()); BSONObjBuilder maxb; maxb.appendMaxForType("", element.type()); - setMaxRecord(collScan, maxb.obj()); + setLowestRecord(collScan->maxRecord, maxb.obj()); bool compatible = compatibleCollator(params, collator, element); if (!compatible) { - return; // Collator affects probe and it's not compatible with collection's collator. + return false; // Collator affects probe and it's not compatible with collection's collator. } // Even if the collations don't match at this point, it's fine, // because the bounds exclude values that use it - collScan->hasCompatibleCollation = true; - const auto collated = IndexBoundsBuilder::objFromElement(element, collator); if (dynamic_cast<const EqualityMatchExpression*>(match)) { - setMinRecord(collScan, collated); - setMaxRecord(collScan, collated); + setHighestRecord(collScan->minRecord, collated); + setLowestRecord(collScan->maxRecord, collated); } else if (dynamic_cast<const LTMatchExpression*>(match) || dynamic_cast<const LTEMatchExpression*>(match)) { - setMaxRecord(collScan, collated); + setLowestRecord(collScan->maxRecord, collated); } else if (dynamic_cast<const GTMatchExpression*>(match) || dynamic_cast<const GTEMatchExpression*>(match)) { - setMinRecord(collScan, collated); + setHighestRecord(collScan->minRecord, collated); } + + return true; } } // namespace @@ -447,13 +516,18 @@ std::unique_ptr<QuerySolutionNode> QueryPlannerAccess::makeCollectionScan( auto queryCollator = query.getCollator(); auto collCollator = params.clusteredCollectionCollator; - csn->hasCompatibleCollation = - !queryCollator || (collCollator && *queryCollator == *collCollator); + csn->hasCompatibleCollation = CollatorInterface::collatorsMatch(queryCollator, collCollator); if (params.clusteredInfo && !csn->resumeAfterRecordId) { // This is a clustered collection. Attempt to perform an efficient, bounded collection scan - // via minRecord and maxRecord if applicable. - handleRIDRangeScan(csn->filter.get(), csn.get(), params, queryCollator); + // via minRecord and maxRecord if applicable. During this process, we will check if the + // query is guaranteed to exclude values of the cluster key which are affected by collation. + // If so, then even if the query and collection collations differ, the collation difference + // won't affect the query results. In that case, we can say hasCompatibleCollation is true. + bool compatibleCollation = + handleRIDRangeScan(csn->filter.get(), csn.get(), params, queryCollator); + csn->hasCompatibleCollation |= compatibleCollation; + handleRIDRangeMinMax(query, csn.get(), params, queryCollator); } @@ -1114,47 +1188,13 @@ std::vector<std::unique_ptr<QuerySolutionNode>> QueryPlannerAccess::collapseEqui } /** - * Returns true if this is a null query that can retrieve all the information it needs directly from - * the index, and so does not need a FETCH stage on top of it. Returns false otherwise. + * This helper determines if a query can be covered depending on the query projection. */ -bool isCoveredNullQuery(const CanonicalQuery& query, - MatchExpression* root, - IndexTag* tag, - const vector<IndexEntry>& indices, - const QueryPlannerParams& params) { - // Sparse indexes and hashed indexes should not use this optimization as they will require a - // FETCH stage with a filter. - if (indices[tag->index].sparse || indices[tag->index].type == IndexType::INDEX_HASHED) { - return false; - } - - // When the index is not multikey, we can support a query on an indexed field searching for null - // values. This optimization can only be done when the index is not multikey, otherwise empty - // arrays in the collection will be treated as null/undefined by the index. When the index is - // multikey, we can support a query searching for both null and empty array values. - const auto multikeyIndex = indices[tag->index].multikey; - if (root->matchType() == MatchExpression::MatchType::MATCH_IN) { - // Check that the query matches null values, if the index is not multikey, or null and empty - // array values, if the index is multikey. Note that the query may match values other than - // null (and empty array). - const auto node = static_cast<const InMatchExpression*>(root); - if (!node->hasNull() || (multikeyIndex && !node->hasEmptyArray())) { - return false; - } - } else if (ComparisonMatchExpressionBase::isEquality(root->matchType()) && !multikeyIndex) { - // Check that the query matches null values. - const auto node = static_cast<const ComparisonMatchExpressionBase*>(root); - if (node->getData().type() != BSONType::jstNULL) { - return false; - } - } else { - return false; - } - +bool projNeedsFetch(const CanonicalQuery& query, const QueryPlannerParams& params) { // If nothing is being projected, the query is fully covered without a fetch. // This is trivially true for a count query. if (params.options & QueryPlannerParams::Options::IS_COUNT) { - return true; + return false; } // This optimization can only be used for find when the index covers the projection completely. @@ -1163,7 +1203,7 @@ bool isCoveredNullQuery(const CanonicalQuery& query, // in the multikey case). Hence, only find queries projecting _id are covered. auto proj = query.getProj(); if (!proj) { - return false; + return true; } // We can cover projections on _id and generated fields and expressions depending only on _id. @@ -1175,10 +1215,38 @@ bool isCoveredNullQuery(const CanonicalQuery& query, // Note that it is not possible to project onto dotted paths of _id here, since they may be // null or missing, and the index cannot differentiate between the two cases, so we would // still need a FETCH stage. - return projFields.size() == 1 && *projFields.begin() == "_id"; + if (projFields.size() == 1 && *projFields.begin() == "_id") { + return false; + } } - return false; + return true; +} + +/** + * This helper updates a MAYBE_COVERED query tightness to one of EXACT, INEXACT_COVERED, or + * INEXACT_FETCH, depending on whether we need a FETCH/filter to answer the query projection. + */ +void refineTightnessForMaybeCoveredQuery(const CanonicalQuery& query, + const QueryPlannerParams& params, + IndexBoundsBuilder::BoundsTightness& tightnessOut) { + // We need to refine the tightness in case we have a "MAYBE_COVERED" tightness bound which + // depends on the query's projection. We will not have information about the projection + // later on in order to make this determination, so we do it here. + const bool noFetchNeededForProj = !projNeedsFetch(query, params); + if (tightnessOut == IndexBoundsBuilder::EXACT_MAYBE_COVERED) { + if (noFetchNeededForProj) { + tightnessOut = IndexBoundsBuilder::EXACT; + } else { + tightnessOut = IndexBoundsBuilder::INEXACT_FETCH; + } + } else if (tightnessOut == IndexBoundsBuilder::INEXACT_MAYBE_COVERED) { + if (noFetchNeededForProj) { + tightnessOut = IndexBoundsBuilder::INEXACT_COVERED; + } else { + tightnessOut = IndexBoundsBuilder::INEXACT_FETCH; + } + } } bool QueryPlannerAccess::processIndexScans(const CanonicalQuery& query, @@ -1222,11 +1290,6 @@ bool QueryPlannerAccess::processIndexScans(const CanonicalQuery& query, // If we're here, we now know that 'child' can use an index directly and the index is // over the child's field. - // We need to track if this is a covered null query so that we can have this information - // at hand when handling the filter on an indexed AND. - scanState.isCoveredNullQuery = - isCoveredNullQuery(query, child, scanState.ixtag, indices, params); - // If 'child' is a NOT, then the tag we're interested in is on the NOT's // child node. if (MatchExpression::NOT == child->matchType()) { @@ -1259,6 +1322,7 @@ bool QueryPlannerAccess::processIndexScans(const CanonicalQuery& query, verify(scanState.currentIndexNumber == scanState.ixtag->index); scanState.tightness = IndexBoundsBuilder::INEXACT_FETCH; mergeWithLeafNode(child, &scanState); + refineTightnessForMaybeCoveredQuery(query, params, scanState.tightness); handleFilter(&scanState); } else { if (nullptr != scanState.currentScan.get()) { @@ -1278,6 +1342,7 @@ bool QueryPlannerAccess::processIndexScans(const CanonicalQuery& query, &scanState.tightness, scanState.getCurrentIETBuilder()); + refineTightnessForMaybeCoveredQuery(query, params, scanState.tightness); handleFilter(&scanState); } } @@ -1693,6 +1758,12 @@ std::unique_ptr<QuerySolutionNode> QueryPlannerAccess::_buildIndexedDataAccess( return soln; } + // We may be able to avoid adding an extra fetch stage even though the bounds are + // inexact, for instance if the query is counting null values on an indexed field + // without projecting that field. We therefore convert "MAYBE_COVERED" bounds into + // either EXACT or INEXACT, depending on the query projection. + refineTightnessForMaybeCoveredQuery(query, params, tightness); + // If the bounds are exact, the set of documents that satisfy the predicate is // exactly equal to the set of documents that the scan provides. // @@ -1700,11 +1771,7 @@ std::unique_ptr<QuerySolutionNode> QueryPlannerAccess::_buildIndexedDataAccess( // superset of documents that satisfy the predicate, and we must check the // predicate. - // We may also be able to avoid adding an extra fetch stage even though the bounds are - // inexact because the query is counting null values on an indexed field without - // projecting that field. - if (tightness == IndexBoundsBuilder::EXACT || - isCoveredNullQuery(query, root, tag, indices, params)) { + if (tightness == IndexBoundsBuilder::EXACT) { return soln; } else if (tightness == IndexBoundsBuilder::INEXACT_COVERED && !indices[tag->index].multikey) { @@ -1850,10 +1917,9 @@ void QueryPlannerAccess::handleFilterAnd(ScanBuildingState* scanState) { // should always be affixed as a filter. We keep 'curChild' in the $and // for affixing later. ++scanState->curChild; - } else if (scanState->tightness == IndexBoundsBuilder::EXACT || scanState->isCoveredNullQuery) { - // The tightness of the bounds is exact or we are dealing with a covered null query. - // Either way, we want to remove this child so that when control returns to handleIndexedAnd - // we know that we don't need it to create a FETCH stage. + } else if (scanState->tightness == IndexBoundsBuilder::EXACT) { + // The tightness of the bounds is exact. We want to remove this child so that when control + // returns to handleIndexedAnd we know that we don't need it to create a FETCH stage. root->getChildVector()->erase(root->getChildVector()->begin() + scanState->curChild); } else if (scanState->tightness == IndexBoundsBuilder::INEXACT_COVERED && (INDEX_TEXT == index.type || !index.multikey)) { diff --git a/src/mongo/db/query/planner_access.h b/src/mongo/db/query/planner_access.h index 6ea44830415..a5bfc04bb51 100644 --- a/src/mongo/db/query/planner_access.h +++ b/src/mongo/db/query/planner_access.h @@ -106,7 +106,7 @@ public: static std::unique_ptr<QuerySolutionNode> makeCollectionScan(const CanonicalQuery& query, bool tailable, const QueryPlannerParams& params, - int direction = 1); + int direction); /** * Return a plan that uses the provided index as a proxy for a collection scan. @@ -126,7 +126,7 @@ public: const BSONObj& endKey); /** - * Consructs a data access plan for 'query' which answers the predicate contained in 'root'. + * Constructs a data access plan for 'query' which answers the predicate contained in 'root'. * Assumes the presence of the passed in indices. Planning behavior is controlled by the * settings in 'params'. */ @@ -145,11 +145,9 @@ private: struct ScanBuildingState { ScanBuildingState(MatchExpression* theRoot, const std::vector<IndexEntry>& indexList, - bool inArrayOp, - bool isCoveredNull = false) + bool inArrayOp) : root(theRoot), inArrayOperator(inArrayOp), - isCoveredNullQuery(isCoveredNull), indices(indexList), currentScan(nullptr), curChild(0), @@ -188,9 +186,6 @@ private: // Are we inside an array operator such as $elemMatch or $all? bool inArrayOperator; - // Is this a covered null query? - bool isCoveredNullQuery; - // A list of relevant indices which 'root' may be tagged to use. const std::vector<IndexEntry>& indices; diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp index 40d8d7b0d0d..ee86d426560 100644 --- a/src/mongo/db/query/planner_analysis.cpp +++ b/src/mongo/db/query/planner_analysis.cpp @@ -337,7 +337,7 @@ void geoSkipValidationOn(const std::set<StringData>& twoDSphereFields, /** * If any field is missing from the list of fields the projection wants, we are not covered. */ -auto providesAllFields(const std::set<std::string>& fields, const QuerySolutionNode& solnRoot) { +auto providesAllFields(const OrderedPathSet& fields, const QuerySolutionNode& solnRoot) { for (auto&& field : fields) { if (!solnRoot.hasField(field)) return false; @@ -580,6 +580,10 @@ void removeProjectSimpleBelowGroupRecursive(QuerySolutionNode* solnRoot) { if (solnRoot->getType() == StageType::STAGE_GROUP) { auto groupNode = static_cast<GroupNode*>(solnRoot); + if (groupNode->needWholeDocument) { + // The sub expression needs the whole document. + return; + } auto projectNodeCandidate = groupNode->children[0]; if (projectNodeCandidate->getType() == StageType::STAGE_GROUP) { // Multiple $group stages may be pushed down. So, if the child is a GROUP, then recurse. diff --git a/src/mongo/db/query/planner_ixselect.cpp b/src/mongo/db/query/planner_ixselect.cpp index 518da370750..21edb7c864f 100644 --- a/src/mongo/db/query/planner_ixselect.cpp +++ b/src/mongo/db/query/planner_ixselect.cpp @@ -245,7 +245,7 @@ static bool boundsGeneratingNodeContainsComparisonToType(MatchExpression* node, // static void QueryPlannerIXSelect::getFields(const MatchExpression* node, string prefix, - stdx::unordered_set<string>* out) { + RelevantFieldIndexMap* out) { // Do not traverse tree beyond a NOR negation node MatchExpression::MatchType exprtype = node->matchType(); if (exprtype == MatchExpression::NOR) { @@ -254,16 +254,13 @@ void QueryPlannerIXSelect::getFields(const MatchExpression* node, // Leaf nodes with a path and some array operators. if (Indexability::nodeCanUseIndexOnOwnField(node)) { - out->insert(prefix + node->path().toString()); - } else if (Indexability::arrayUsesIndexOnChildren(node)) { + bool supportSparse = Indexability::nodeSupportedBySparseIndex(node); + (*out)[prefix + node->path().toString()] = {supportSparse}; + } else if (Indexability::arrayUsesIndexOnChildren(node) && !node->path().empty()) { // If the array uses an index on its children, it's something like // {foo : {$elemMatch: {bar: 1}}}, in which case the predicate is really over foo.bar. - // - // When we have {foo: {$all: [{$elemMatch: {a: 1}}], the path of the embedded elemMatch - // is empty. We don't want to append a dot in that case as the field would be foo..a. - if (!node->path().empty()) { - prefix += node->path().toString() + "."; - } + // Note we skip empty path components since they are not allowed in index key patterns. + prefix += node->path().toString() + "."; for (size_t i = 0; i < node->numChildren(); ++i) { getFields(node->getChild(i), prefix, out); @@ -275,8 +272,7 @@ void QueryPlannerIXSelect::getFields(const MatchExpression* node, } } -void QueryPlannerIXSelect::getFields(const MatchExpression* node, - stdx::unordered_set<string>* out) { +void QueryPlannerIXSelect::getFields(const MatchExpression* node, RelevantFieldIndexMap* out) { getFields(node, "", out); } @@ -316,26 +312,40 @@ std::vector<IndexEntry> QueryPlannerIXSelect::findIndexesByHint( // static std::vector<IndexEntry> QueryPlannerIXSelect::findRelevantIndices( - const stdx::unordered_set<std::string>& fields, const std::vector<IndexEntry>& allIndices) { + const RelevantFieldIndexMap& fields, const std::vector<IndexEntry>& allIndices) { std::vector<IndexEntry> out; - for (auto&& entry : allIndices) { - BSONObjIterator it(entry.keyPattern); + for (auto&& index : allIndices) { + BSONObjIterator it(index.keyPattern); BSONElement elt = it.next(); - if (fields.end() != fields.find(elt.fieldName())) { - out.push_back(entry); + const std::string fieldName = elt.fieldNameStringData().toString(); + + // If the index is non-sparse we can use the field regardless its sparsity, otherwise we + // should find the field that can be answered by a sparse index. + if (fields.contains(fieldName) && + (!index.sparse || fields.find(fieldName)->second.isSparse)) { + out.push_back(index); } } return out; } -std::vector<IndexEntry> QueryPlannerIXSelect::expandIndexes( - const stdx::unordered_set<std::string>& fields, std::vector<IndexEntry> relevantIndices) { +std::vector<IndexEntry> QueryPlannerIXSelect::expandIndexes(const RelevantFieldIndexMap& fields, + std::vector<IndexEntry> relevantIndices, + bool indexHinted) { std::vector<IndexEntry> out; + // Filter out fields that cannot be answered by any sparse index. We know wildcard indexes are + // sparse, so we don't want to expand the wildcard index based on such fields. + stdx::unordered_set<std::string> sparseIncompatibleFields; + for (auto&& [fieldName, idxProperty] : fields) { + if (idxProperty.isSparse || indexHinted) { + sparseIncompatibleFields.insert(fieldName); + } + } for (auto&& entry : relevantIndices) { if (entry.type == IndexType::INDEX_WILDCARD) { - wcp::expandWildcardIndexEntry(entry, fields, &out); + wcp::expandWildcardIndexEntry(entry, sparseIncompatibleFields, &out); } else { out.push_back(std::move(entry)); } @@ -780,7 +790,8 @@ void QueryPlannerIXSelect::_rateIndices(MatchExpression* node, childRt->path = rt->path; node->getChild(0)->setTag(childRt); } - } else if (Indexability::arrayUsesIndexOnChildren(node)) { + } else if (Indexability::arrayUsesIndexOnChildren(node) && !node->path().empty()) { + // Note we skip empty path components since they are not allowed in index key patterns. const auto newPath = prefix + node->path().toString(); ElemMatchContext newContext; // Note this StringData is unowned and references the string declared on the stack here. @@ -791,12 +802,7 @@ void QueryPlannerIXSelect::_rateIndices(MatchExpression* node, // If the array uses an index on its children, it's something like // {foo: {$elemMatch: {bar: 1}}}, in which case the predicate is really over foo.bar. - // - // When we have {foo: {$all: [{$elemMatch: {a: 1}}], the path of the embedded elemMatch - // is empty. We don't want to append a dot in that case as the field would be foo..a. - if (!node->path().empty()) { - prefix += node->path().toString() + "."; - } + prefix += node->path().toString() + "."; for (size_t i = 0; i < node->numChildren(); ++i) { _rateIndices(node->getChild(i), prefix, indices, collator, newContext); } diff --git a/src/mongo/db/query/planner_ixselect.h b/src/mongo/db/query/planner_ixselect.h index 0ef2d480953..15f1e135d5e 100644 --- a/src/mongo/db/query/planner_ixselect.h +++ b/src/mongo/db/query/planner_ixselect.h @@ -38,16 +38,24 @@ namespace mongo { class CollatorInterface; +struct IndexProperties { + bool isSparse = false; // 'true' if a sparse index can answer the field. +}; + +// A relevant field to index requirement map. +using RelevantFieldIndexMap = stdx::unordered_map<std::string, IndexProperties>; + /** * Methods for determining what fields and predicates can use indices. */ class QueryPlannerIXSelect { public: /** - * Return all the fields in the tree rooted at 'node' that we can use an index on - * in order to answer the query. + * Return all the fields in the tree rooted at 'node' that we can use an index to answer the + * query. The output, 'RelevantFieldIndexMap', contains the requirements of the index that can + * answer the field. e.g. Some fields can be supported only by a non-sparse index. */ - static void getFields(const MatchExpression* node, stdx::unordered_set<std::string>* out); + static void getFields(const MatchExpression* node, RelevantFieldIndexMap* out); /** * Similar to other getFields() method, but with 'prefix' argument which is a path prefix to be @@ -57,7 +65,7 @@ public: */ static void getFields(const MatchExpression* node, std::string prefix, - stdx::unordered_set<std::string>* out); + RelevantFieldIndexMap* out); /** * Finds all indices that correspond to the hinted index. Matches the index both by name and by @@ -70,8 +78,8 @@ public: * Finds all indices prefixed by fields we have predicates over. Only these indices are * useful in answering the query. */ - static std::vector<IndexEntry> findRelevantIndices( - const stdx::unordered_set<std::string>& fields, const std::vector<IndexEntry>& allIndices); + static std::vector<IndexEntry> findRelevantIndices(const RelevantFieldIndexMap& fields, + const std::vector<IndexEntry>& allIndices); /** * Determine how useful all of our relevant 'indices' are to all predicates in the subtree @@ -130,9 +138,12 @@ public: /** * Given a list of IndexEntries and fields used by a query's match expression, return a list * "expanded" indexes (where the $** indexes in the given list have been expanded). + * 'hintedIndexBson' indicates that the indexes in 'relevantIndices' are the results of the + * user's hint. */ - static std::vector<IndexEntry> expandIndexes(const stdx::unordered_set<std::string>& fields, - std::vector<IndexEntry> relevantIndices); + static std::vector<IndexEntry> expandIndexes(const RelevantFieldIndexMap& fields, + std::vector<IndexEntry> relevantIndices, + bool hintedIndexBson = false); /** * Check if this match expression is a leaf and is supported by a wildcard index. diff --git a/src/mongo/db/query/planner_ixselect_test.cpp b/src/mongo/db/query/planner_ixselect_test.cpp index 93c4f12c821..1df4d714e67 100644 --- a/src/mongo/db/query/planner_ixselect_test.cpp +++ b/src/mongo/db/query/planner_ixselect_test.cpp @@ -65,6 +65,20 @@ unique_ptr<MatchExpression> parseMatchExpression(const BSONObj& obj) { return std::move(status.getValue()); } +using FieldIter = RelevantFieldIndexMap::iterator; +string toString(FieldIter begin, FieldIter end) { + str::stream ss; + ss << "["; + for (FieldIter i = begin; i != end; i++) { + if (i != begin) { + ss << " "; + } + ss << i->first; + } + ss << "]"; + return ss; +} + /** * Utility function to join elements in iterator range with comma */ @@ -88,10 +102,13 @@ string toString(Iter begin, Iter end) { * to QueryPlannerIXSelect::getFields() * Results are compared with expected fields (parsed from expectedFieldsStr) */ -void testGetFields(const char* query, const char* prefix, const char* expectedFieldsStr) { +void testGetFields(const char* query, + const char* prefix, + const char* expectedFieldsStr, + bool sparseSupported = true) { BSONObj obj = fromjson(query); unique_ptr<MatchExpression> expr(parseMatchExpression(obj)); - stdx::unordered_set<string> fields; + RelevantFieldIndexMap fields; QueryPlannerIXSelect::getFields(expr.get(), prefix, &fields); // Verify results @@ -99,7 +116,7 @@ void testGetFields(const char* query, const char* prefix, const char* expectedFi vector<string> expectedFields = StringSplitter::split(expectedFieldsStr, ","); for (vector<string>::const_iterator i = expectedFields.begin(); i != expectedFields.end(); i++) { - if (fields.find(*i) == fields.end()) { + if (fields[*i].isSparse != sparseSupported) { str::stream ss; ss << "getFields(query=" << query << ", prefix=" << prefix << "): unable to find " << *i << " in result: " << toString(fields.begin(), fields.end()); @@ -159,6 +176,12 @@ TEST(QueryPlannerIXSelectTest, GetFieldsArrayNegation) { testGetFields("{a: {$all: [{$elemMatch: {b: {$ne: 1}}}]}}", "", "a.b"); } +TEST(QueryPlannerIXSelectTest, GetFieldsInternalExpr) { + testGetFields("{$expr: {$lt: ['$a', 'r']}}", "", "", false /* sparse supported */); + testGetFields("{$expr: {$eq: ['$a', null]}}", "", "", false /* sparse supported */); + testGetFields("{$expr: {$eq: ['$a', 1]}}", "", "", false /* sparse supported */); +} + /** * Performs a pre-order traversal of expression tree. Validates * that all tagged nodes contain an instance of RelevantTag. @@ -1153,25 +1176,6 @@ TEST(QueryPlannerIXSelectTest, InternalExprEqCanUseTextIndexSuffix) { "{a: {$_internalExprEq: 1}}", "", kSimpleCollator, indices, "a", expectedIndices); } -TEST(QueryPlannerIXSelectTest, InternalExprEqCanUseSparseIndexWithComparisonToNull) { - auto entry = buildSimpleIndexEntry(BSON("a" << 1)); - entry.sparse = true; - std::vector<IndexEntry> indices; - indices.push_back(entry); - std::set<size_t> expectedIndices = {0}; - testRateIndices( - "{a: {$_internalExprEq: null}}", "", kSimpleCollator, indices, "a", expectedIndices); -} - -TEST(QueryPlannerIXSelectTest, InternalExprEqCanUseSparseIndexWithComparisonToNonNull) { - auto entry = buildSimpleIndexEntry(BSON("a" << 1)); - entry.sparse = true; - std::vector<IndexEntry> indices; - indices.push_back(entry); - std::set<size_t> expectedIndices = {0}; - testRateIndices( - "{a: {$_internalExprEq: 1}}", "", kSimpleCollator, indices, "a", expectedIndices); -} TEST(QueryPlannerIXSelectTest, NotEqualsNullCanUseIndex) { auto entry = buildSimpleIndexEntry(BSON("a" << 1)); std::set<size_t> expectedIndices = {0}; @@ -1357,18 +1361,17 @@ TEST(QueryPlannerIXSelectTest, ExpandWildcardIndices) { const auto indexEntry = makeIndexEntry(BSON("$**" << 1), {}); // Case where no fields are specified. - std::vector<IndexEntry> result = - QueryPlannerIXSelect::expandIndexes(stdx::unordered_set<string>(), {indexEntry.first}); + std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes({}, {indexEntry.first}); ASSERT_TRUE(result.empty()); - stdx::unordered_set<string> fields = {"fieldA", "fieldB"}; + RelevantFieldIndexMap fields = {{"fieldA", {true}}, {"fieldB", {true}}}; result = QueryPlannerIXSelect::expandIndexes(fields, {indexEntry.first}); std::vector<BSONObj> expectedKeyPatterns = {BSON("fieldA" << 1), BSON("fieldB" << 1)}; ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result)); const auto wildcardIndexWithSubpath = makeIndexEntry(BSON("a.b.$**" << 1), {}); - fields = {"a.b", "a.b.c", "a.d"}; + fields = {{"a.b", {true}}, {"a.b.c", {true}}, {"a.d", {true}}}; result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexWithSubpath.first}); expectedKeyPatterns = {BSON("a.b" << 1), BSON("a.b.c" << 1)}; ASSERT_TRUE(indexEntryKeyPatternsMatch(&expectedKeyPatterns, &result)); @@ -1380,7 +1383,8 @@ TEST(QueryPlannerIXSelectTest, ExpandWildcardIndicesInPresenceOfOtherIndices) { auto bIndexEntry = makeIndexEntry(BSON("fieldB" << 1), {}); auto abIndexEntry = makeIndexEntry(BSON("fieldA" << 1 << "fieldB" << 1), {}); - const stdx::unordered_set<string> fields = {"fieldA", "fieldB", "fieldC"}; + const RelevantFieldIndexMap fields = { + {"fieldA", {true}}, {"fieldB", {true}}, {"fieldC", {true}}}; std::vector<BSONObj> expectedKeyPatterns = { BSON("fieldA" << 1), BSON("fieldA" << 1), BSON("fieldB" << 1), BSON("fieldC" << 1)}; @@ -1418,7 +1422,7 @@ TEST(QueryPlannerIXSelectTest, ExpandWildcardIndicesInPresenceOfOtherIndices) { TEST(QueryPlannerIXSelectTest, ExpandedIndexEntriesAreCorrectlyMarkedAsMultikeyOrNonMultikey) { auto wildcardIndexEntry = makeIndexEntry(BSON("$**" << 1), {}, {FieldRef{"a"}}); - const stdx::unordered_set<string> fields = {"a.b", "c.d"}; + RelevantFieldIndexMap fields = {{"a.b", {true}}, {"c.d", {true}}}; std::vector<BSONObj> expectedKeyPatterns = {BSON("a.b" << 1), BSON("c.d" << 1)}; auto result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1442,7 +1446,7 @@ TEST(QueryPlannerIXSelectTest, ExpandedIndexEntriesAreCorrectlyMarkedAsMultikeyO TEST(QueryPlannerIXSelectTest, WildcardIndexExpansionExcludesIdField) { const auto indexEntry = makeIndexEntry(BSON("$**" << 1), {}); - stdx::unordered_set<string> fields = {"_id", "abc", "def"}; + RelevantFieldIndexMap fields = {{"_id", {true}}, {"abc", {true}}, {"def", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {indexEntry.first}); @@ -1454,7 +1458,7 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesExpandedEntryHasCorrectProperties) auto wildcardIndexEntry = makeIndexEntry(BSON("$**" << 1), {}); wildcardIndexEntry.first.identifier = IndexEntry::Identifier("someIndex"); - stdx::unordered_set<string> fields = {"abc", "def"}; + RelevantFieldIndexMap fields = {{"abc", {true}}, {"def", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1484,7 +1488,11 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesExpandedEntryHasCorrectProperties) TEST(QueryPlannerIXSelectTest, WildcardIndicesExcludeNonMatchingKeySubpath) { auto wildcardIndexEntry = makeIndexEntry(BSON("subpath.$**" << 1), {}); - stdx::unordered_set<string> fields = {"abc", "def", "subpath.abc", "subpath.def", "subpath"}; + RelevantFieldIndexMap fields = {{"abc", {true}}, + {"def", {true}}, + {"subpath.abc", {true}}, + {"subpath.def", {true}}, + {"subpath", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1500,7 +1508,11 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesExcludeNonMatchingPathsWithInclusi {}, BSON("wildcardProjection" << BSON("abc" << 1 << "subpath.abc" << 1))); - stdx::unordered_set<string> fields = {"abc", "def", "subpath.abc", "subpath.def", "subpath"}; + RelevantFieldIndexMap fields = {{"abc", {true}}, + {"def", {true}}, + {"subpath.abc", {true}}, + {"subpath.def", {true}}, + {"subpath", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1515,7 +1527,11 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesExcludeNonMatchingPathsWithExclusi {}, BSON("wildcardProjection" << BSON("abc" << 0 << "subpath.abc" << 0))); - stdx::unordered_set<string> fields = {"abc", "def", "subpath.abc", "subpath.def", "subpath"}; + RelevantFieldIndexMap fields = {{"abc", {true}}, + {"def", {true}}, + {"subpath.abc", {true}}, + {"subpath.def", {true}}, + {"subpath", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1531,8 +1547,12 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesWithInclusionProjectionAllowIdExcl {}, BSON("wildcardProjection" << BSON("_id" << 0 << "abc" << 1 << "subpath.abc" << 1))); - stdx::unordered_set<string> fields = { - "_id", "abc", "def", "subpath.abc", "subpath.def", "subpath"}; + RelevantFieldIndexMap fields = {{"_id", {true}}, + {"abc", {true}}, + {"def", {true}}, + {"subpath.abc", {true}}, + {"subpath.def", {true}}, + {"subpath", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1547,8 +1567,12 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesWithInclusionProjectionAllowIdIncl {}, BSON("wildcardProjection" << BSON("_id" << 1 << "abc" << 1 << "subpath.abc" << 1))); - stdx::unordered_set<string> fields = { - "_id", "abc", "def", "subpath.abc", "subpath.def", "subpath"}; + RelevantFieldIndexMap fields = {{"_id", {true}}, + {"abc", {true}}, + {"def", {true}}, + {"subpath.abc", {true}}, + {"subpath.def", {true}}, + {"subpath", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1564,8 +1588,12 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesWithExclusionProjectionAllowIdIncl {}, BSON("wildcardProjection" << BSON("_id" << 1 << "abc" << 0 << "subpath.abc" << 0))); - stdx::unordered_set<string> fields = { - "_id", "abc", "def", "subpath.abc", "subpath.def", "subpath"}; + RelevantFieldIndexMap fields = {{"_id", {true}}, + {"abc", {true}}, + {"def", {true}}, + {"subpath.abc", {true}}, + {"subpath.def", {true}}, + {"subpath", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); @@ -1578,8 +1606,12 @@ TEST(QueryPlannerIXSelectTest, WildcardIndicesIncludeMatchingInternalNodes) { auto wildcardIndexEntry = makeIndexEntry( BSON("$**" << 1), {}, {}, BSON("wildcardProjection" << BSON("_id" << 1 << "subpath" << 1))); - stdx::unordered_set<string> fields = { - "_id", "abc", "def", "subpath.abc", "subpath.def", "subpath"}; + RelevantFieldIndexMap fields = {{"_id", {true}}, + {"abc", {true}}, + {"def", {true}}, + {"subpath.abc", {true}}, + {"subpath.def", {true}}, + {"subpath", {true}}}; std::vector<IndexEntry> result = QueryPlannerIXSelect::expandIndexes(fields, {wildcardIndexEntry.first}); diff --git a/src/mongo/db/query/projection.cpp b/src/mongo/db/query/projection.cpp index af5fedfe780..b55fc03602a 100644 --- a/src/mongo/db/query/projection.cpp +++ b/src/mongo/db/query/projection.cpp @@ -49,7 +49,7 @@ struct DepsAnalysisData { fieldDependencyTracker.fields.insert(fieldName); } - std::set<std::string> requiredFields() const { + OrderedPathSet requiredFields() const { return fieldDependencyTracker.fields; } }; diff --git a/src/mongo/db/query/projection.h b/src/mongo/db/query/projection.h index 914567e87ed..97b8d1e0d30 100644 --- a/src/mongo/db/query/projection.h +++ b/src/mongo/db/query/projection.h @@ -49,7 +49,7 @@ struct ProjectionDependencies { bool containsElemMatch = false; // Which fields are necessary to perform the projection, or boost::none if all are required. - boost::optional<std::set<std::string>> requiredFields; + boost::optional<OrderedPathSet> requiredFields; bool hasDottedPath = false; @@ -95,7 +95,7 @@ public: * Return which fields are required to compute the projection, assuming the entire document is * not needed. */ - const std::set<std::string>& getRequiredFields() const { + const OrderedPathSet& getRequiredFields() const { invariant(_type == ProjectType::kInclusion); return *_deps.requiredFields; } diff --git a/src/mongo/db/query/query_feature_flags.idl b/src/mongo/db/query/query_feature_flags.idl index 5ac4a341a2c..81f89c97997 100644 --- a/src/mongo/db/query/query_feature_flags.idl +++ b/src/mongo/db/query/query_feature_flags.idl @@ -145,3 +145,20 @@ feature_flags: description: "Feature flag to enable using SBE for a larger number of queries" cpp_varname: gFeatureFlagSbeFull default: false + + featureFlagShardedSearchCustomSort: + description: "Feature flag to enable user specified sort for sharded $search queries." + cpp_varname: gFeatureFlagShardedSearchCustomSort + default: true + version: 6.0 + + featureFlagSearchBatchSizeLimit: + description: "Feature flag to enable the search batchsize and limit optimization." + cpp_varname: gFeatureFlagSearchBatchSizeLimit + default: true + version: 6.0 + + featureFlagVectorSearchPublicPreview: + description: "Feature flag to enable vector search for public preview." + cpp_varname: gFeatureFlagVectorSearchPublicPreview + default: false diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl index f894629037f..8aad5969289 100644 --- a/src/mongo/db/query/query_knobs.idl +++ b/src/mongo/db/query/query_knobs.idl @@ -871,6 +871,25 @@ server_parameters: default: expr: false + internalQueryGlobalProfilingFilter: + description: "Enables the setProfilingFilterGlobally command." + set_at: [ startup ] + cpp_varname: internalQueryGlobalProfilingFilter + cpp_vartype: AtomicWord<bool> + default: false + + + internalQueryDocumentSourceWriterBatchExtraReservedBytes: + description: "Space to reserve in document source writer batches for miscellaneous metadata" + set_at: [ startup, runtime ] + cpp_vartype: AtomicWord<int> + cpp_varname: internalQueryDocumentSourceWriterBatchExtraReservedBytes + validator: + gte: 0 + lte: + expr: 8 * 1024 * 1024 # 8MB + default: 0 + # Note for adding additional query knobs: # # When adding a new query knob, you should consider whether or not you need to add an 'on_update' diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp index c258c9e6867..e04516a6330 100644 --- a/src/mongo/db/query/query_planner.cpp +++ b/src/mongo/db/query/query_planner.cpp @@ -252,6 +252,11 @@ void tryToAddColumnScan(const QueryPlannerParams& params, // collection scan. Add that solution. out.push_back(QueryPlannerAnalysis::analyzeDataAccess(query, params, std::move(columnScan))); } + +bool collscanIsBounded(const CollectionScanNode* collscan) { + return collscan->minRecord || collscan->maxRecord; +} + } // namespace using std::numeric_limits; @@ -469,13 +474,57 @@ static BSONObj finishMaxObj(const IndexEntry& indexEntry, } } +bool providesSort(const CanonicalQuery& query, const BSONObj& kp) { + return query.getFindCommandRequest().getSort().isPrefixOf( + kp, SimpleBSONElementComparator::kInstance); +} + +/** + * Determine whether this query has a sort that can be provided by the clustered index, if so, which + * direction the scan should be. If the collection is not clustered, or the sort cannot be provided, + * returns 'boost::none'. + */ +boost::optional<int> determineClusteredScanDirection(const CanonicalQuery& query, + const QueryPlannerParams& params) { + if (params.clusteredInfo && query.getSortPattern() && + CollatorInterface::collatorsMatch(params.clusteredCollectionCollator, + query.getCollator())) { + auto kp = clustered_util::getSortPattern(params.clusteredInfo->getIndexSpec()); + if (providesSort(query, kp)) { + return 1; + } else if (providesSort(query, QueryPlannerCommon::reverseSortObj(kp))) { + return -1; + } + } + + return boost::none; +} + +/** + * Determine the direction of the scan needed for the query. Defaults to 1 unless this is a + * clustered collection and we have a sort that can be provided by the clustered index. + */ +int determineCollscanDirection(const CanonicalQuery& query, const QueryPlannerParams& params) { + return determineClusteredScanDirection(query, params).value_or(1); +} + +std::pair<std::unique_ptr<QuerySolution>, const CollectionScanNode*> buildCollscanSolnWithNode( + const CanonicalQuery& query, + bool tailable, + const QueryPlannerParams& params, + boost::optional<int> direction = boost::none) { + std::unique_ptr<QuerySolutionNode> solnRoot(QueryPlannerAccess::makeCollectionScan( + query, tailable, params, direction.value_or(determineCollscanDirection(query, params)))); + const auto* collscanNode = checked_cast<const CollectionScanNode*>(solnRoot.get()); + return std::make_pair( + QueryPlannerAnalysis::analyzeDataAccess(query, params, std::move(solnRoot)), collscanNode); +} + std::unique_ptr<QuerySolution> buildCollscanSoln(const CanonicalQuery& query, bool tailable, const QueryPlannerParams& params, - int direction = 1) { - std::unique_ptr<QuerySolutionNode> solnRoot( - QueryPlannerAccess::makeCollectionScan(query, tailable, params, direction)); - return QueryPlannerAnalysis::analyzeDataAccess(query, params, std::move(solnRoot)); + boost::optional<int> direction = boost::none) { + return buildCollscanSolnWithNode(query, tailable, params, direction).first; } std::unique_ptr<QuerySolution> buildWholeIXSoln( @@ -491,11 +540,6 @@ std::unique_ptr<QuerySolution> buildWholeIXSoln( return QueryPlannerAnalysis::analyzeDataAccess(query, params, std::move(solnRoot)); } -bool providesSort(const CanonicalQuery& query, const BSONObj& kp) { - return query.getFindCommandRequest().getSort().isPrefixOf( - kp, SimpleBSONElementComparator::kInstance); -} - StatusWith<std::unique_ptr<PlanCacheIndexTree>> QueryPlanner::cacheDataFromTaggedTree( const MatchExpression* const taggedTree, const vector<IndexEntry>& relevantIndices) { if (!taggedTree) { @@ -661,7 +705,7 @@ StatusWith<std::unique_ptr<QuerySolution>> QueryPlanner::planFromCache( } else if (SolutionCacheData::COLLSCAN_SOLN == winnerCacheData.solnType) { // The cached solution is a collection scan. We don't cache collscans // with tailable==true, hence the false below. - auto soln = buildCollscanSoln(query, false, params); + auto soln = buildCollscanSoln(query, false, params, winnerCacheData.wholeIXSolnDir); if (!soln) { return Status(ErrorCodes::NoQueryExecutionPlans, "plan cache error: collection scan soln"); @@ -683,10 +727,11 @@ StatusWith<std::unique_ptr<QuerySolution>> QueryPlanner::planFromCache( "filter"_attr = redact(clone->debugString()), "cacheData"_attr = redact(winnerCacheData.toString())); - stdx::unordered_set<string> fields; + RelevantFieldIndexMap fields; QueryPlannerIXSelect::getFields(query.root(), &fields); + // We will not cache queries with 'hint'. std::vector<IndexEntry> expandedIndexes = - QueryPlannerIXSelect::expandIndexes(fields, params.indices); + QueryPlannerIXSelect::expandIndexes(fields, params.indices, false /* indexHinted */); // Map from index name to index number. map<IndexEntry::Identifier, size_t> indexMap; @@ -888,13 +933,14 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan( } // Figure out what fields we care about. - stdx::unordered_set<string> fields; + RelevantFieldIndexMap fields; QueryPlannerIXSelect::getFields(query.root(), &fields); for (auto&& field : fields) { - LOGV2_DEBUG(20970, 5, "Predicate over field", "field"_attr = field); + LOGV2_DEBUG(20970, 5, "Predicate over field", "field"_attr = field.first); } - fullIndexList = QueryPlannerIXSelect::expandIndexes(fields, std::move(fullIndexList)); + fullIndexList = QueryPlannerIXSelect::expandIndexes( + fields, std::move(fullIndexList), !hintedIndex.isEmpty()); std::vector<IndexEntry> relevantIndices; if (!hintedIndexEntry) { @@ -1257,37 +1303,6 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan( } } } - - // The base index is sorted on some key, so it's possible we might want to use - // a collection scan to provide the sort requested - if (params.clusteredInfo) { - if (CollatorInterface::collatorsMatch(params.clusteredCollectionCollator, - query.getCollator())) { - auto kp = clustered_util::getSortPattern(params.clusteredInfo->getIndexSpec()); - int direction = 0; - if (providesSort(query, kp)) { - direction = 1; - } else if (providesSort(query, QueryPlannerCommon::reverseSortObj(kp))) { - direction = -1; - } - - if (direction != 0) { - auto soln = buildCollscanSoln(query, isTailable, params, direction); - if (soln) { - LOGV2_DEBUG(6082401, - 5, - "Planner: outputting soln that uses clustered index to " - "provide sort"); - SolutionCacheData* scd = new SolutionCacheData(); - scd->solnType = SolutionCacheData::COLLSCAN_SOLN; - scd->wholeIXSolnDir = direction; - - soln->cacheData.reset(scd); - out.push_back(std::move(soln)); - } - } - } - } } // If a projection exists, there may be an index that allows for a covered plan, even if @@ -1339,6 +1354,8 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan( "No indexed plans available, and running with 'notablescan'"); } + bool clusteredCollection = params.clusteredInfo.has_value(); + // geoNear and text queries *require* an index. // Also, if a hint is specified it indicates that we MUST use it. bool possibleToCollscan = @@ -1348,21 +1365,34 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan( return Status(ErrorCodes::NoQueryExecutionPlans, "No query solutions"); } - if (possibleToCollscan && (collscanRequested || collScanRequired)) { - auto collscan = buildCollscanSoln(query, isTailable, params); - if (!collscan && collScanRequired) { + if (possibleToCollscan && (collscanRequested || collScanRequired || clusteredCollection)) { + auto clusteredScanDirection = determineClusteredScanDirection(query, params); + auto direction = clusteredScanDirection.value_or(1); + auto [collscanSoln, collscanNode] = + buildCollscanSolnWithNode(query, isTailable, params, direction); + if (!collscanSoln && collScanRequired) { return Status(ErrorCodes::NoQueryExecutionPlans, "Failed to build collection scan soln"); } - if (collscan) { + + // We consider collection scan in the following cases: + // 1. collScanRequested - specifically requested by caller. + // 2. collScanRequired - there are no other possible plans, so we fallback to full scan. + // 3. collscanIsBounded - collection is clustered and clustered index is used. + // 4. clusteredScanDirection - collection is clustered and sort, provided by clustered + // index, is used + if (collscanSoln && + (collscanRequested || collScanRequired || collscanIsBounded(collscanNode) || + clusteredScanDirection)) { LOGV2_DEBUG(20984, 5, "Planner: outputting a collection scan", - "collectionScan"_attr = redact(collscan->toString())); + "collectionScan"_attr = redact(collscanSoln->toString())); SolutionCacheData* scd = new SolutionCacheData(); scd->solnType = SolutionCacheData::COLLSCAN_SOLN; - collscan->cacheData.reset(scd); - out.push_back(std::move(collscan)); + scd->wholeIXSolnDir = direction; + collscanSoln->cacheData.reset(scd); + out.push_back(std::move(collscanSoln)); } } diff --git a/src/mongo/db/query/query_planner_array_test.cpp b/src/mongo/db/query/query_planner_array_test.cpp index 19dfc99dbcd..ff39e98c3d0 100644 --- a/src/mongo/db/query/query_planner_array_test.cpp +++ b/src/mongo/db/query/query_planner_array_test.cpp @@ -2110,13 +2110,13 @@ TEST_F(QueryPlannerTest, CanHoistNegatedPredFromElemMatchIntoSiblingOrWithMultik "{fetch: {filter: {arr: {$elemMatch: {a: {$ne: 1}, b: {$in: [2, 3]}}}}," "node: {" " or: {nodes: [" - " {fetch: {filter: {a: {$ne: 1}}," + " {fetch: {filter: {'arr.a': {$ne: 1}}," " node: {ixscan: {pattern: {'arr.a': 1, 'arr.b': 1, c: 1, d: 1}," " bounds: {'arr.a': [['MinKey', 1, true, false], [1, 'MaxKey', false, true]]," " 'arr.b': [[2, 2, true, true], [3, 3, true, true]]," " c: [[4, 4, true, true]]," " d: [[5, 5, true, true]]}}}}}," - " {fetch: {filter: {a: {$ne: 1}}," + " {fetch: {filter: {'arr.a': {$ne: 1}}," " node: {ixscan: {pattern: {'arr.a': 1, 'arr.b': 1, c: 1, d: 1}," " bounds: {'arr.a': [['MinKey', 1, true, false],[1, 'MaxKey', false, true]]," " 'arr.b': [[2, 2, true, true], [3, 3, true, true]]," diff --git a/src/mongo/db/query/query_planner_index_test.cpp b/src/mongo/db/query/query_planner_index_test.cpp index 045fa35dbe1..24631b9091e 100644 --- a/src/mongo/db/query/query_planner_index_test.cpp +++ b/src/mongo/db/query/query_planner_index_test.cpp @@ -198,26 +198,20 @@ TEST_F(QueryPlannerTest, SparseIndexForQuery) { "{filter: null, pattern: {a: 1}}}}}"); } -TEST_F(QueryPlannerTest, ExprEqCanUseSparseIndex) { +TEST_F(QueryPlannerTest, ExprEqCannotUseSparseIndex) { params.options &= ~QueryPlannerParams::INCLUDE_COLLSCAN; addIndex(fromjson("{a: 1}"), false, true); runQuery(fromjson("{a: {$_internalExprEq: 1}}")); - assertNumSolutions(1U); - assertSolutionExists( - "{fetch: {filter: null, node: {ixscan: " - "{filter: null, pattern: {a: 1}, bounds: {a: [[1,1,true,true]]}}}}}"); + assertHasOnlyCollscan(); } -TEST_F(QueryPlannerTest, ExprEqCanUseSparseIndexForEqualityToNull) { +TEST_F(QueryPlannerTest, ExprEqCannotUseSparseIndexForEqualityToNull) { params.options &= ~QueryPlannerParams::INCLUDE_COLLSCAN; addIndex(fromjson("{a: 1}"), false, true); runQuery(fromjson("{a: {$_internalExprEq: null}}")); - assertNumSolutions(1U); - assertSolutionExists( - "{fetch: {filter: {a: {$_internalExprEq: null}}, node: {ixscan: {filter: null, pattern: " - "{a: 1}, bounds: {a: [[undefined,undefined,true,true], [null,null,true,true]]}}}}}"); + assertHasOnlyCollscan(); } TEST_F(QueryPlannerTest, NegationCannotUseSparseIndex) { diff --git a/src/mongo/db/query/query_planner_tree_test.cpp b/src/mongo/db/query/query_planner_tree_test.cpp index 9d403989792..853c02388b0 100644 --- a/src/mongo/db/query/query_planner_tree_test.cpp +++ b/src/mongo/db/query/query_planner_tree_test.cpp @@ -2464,6 +2464,32 @@ TEST_F(QueryPlannerTest, LockstepOrEnumerationSanityCheckTwoChildrenTwoIndexesEa "{ixscan: {pattern: {a: 1, c: 1}}}}}}}"); } +TEST_F(QueryPlannerTest, TotalPossibleLockstepOrEnumerationReachesTheOrLimit) { + params.options = + QueryPlannerParams::NO_TABLE_SCAN | QueryPlannerParams::ENUMERATE_OR_CHILDREN_LOCKSTEP; + addIndex(BSON("a" << 1 << "b" << 1)); + addIndex(BSON("a" << 1 << "c" << 1)); + + BSONArrayBuilder orBuilder; + // This max number has a value of 65 in order to potentillay triger any overflow of the possible + // enumeration count, because each predicate in $or has two possible indexes, allowing for 2^65 + // possible enumerations. + const int maxPredicates = 65; + for (int i = 0; i < maxPredicates; i++) { + orBuilder.append(BSON("b" << i << "c" << i)); + } + + auto cmd = BSON("find" + << "testns" + << "filter" << BSON("a" << 1 << "$or" << orBuilder.arr())); + + // Ensure that the query runs fine. + runQueryAsCommand(cmd); + + // internalQueryMaxOrSolutions.load() + 2. + assertNumSolutions(12U); +} + // Test that we enumerate the expected plans with the special parameter set. In this test we have // two branches of an $or, each with one possible indexed solution. TEST_F(QueryPlannerTest, LockstepOrEnumerationSanityCheckTwoChildrenOneIndexEach) { diff --git a/src/mongo/db/query/query_planner_wildcard_index_test.cpp b/src/mongo/db/query/query_planner_wildcard_index_test.cpp index cd943126a17..eca110ffac0 100644 --- a/src/mongo/db/query/query_planner_wildcard_index_test.cpp +++ b/src/mongo/db/query/query_planner_wildcard_index_test.cpp @@ -56,7 +56,7 @@ protected: } void addWildcardIndex(BSONObj keyPattern, - const std::set<std::string>& multikeyPathSet = {}, + const OrderedPathSet& multikeyPathSet = {}, BSONObj wildcardProjection = BSONObj{}, MatchExpression* partialFilterExpr = nullptr, CollatorInterface* collator = nullptr, @@ -418,25 +418,18 @@ TEST_F(QueryPlannerWildcardTest, EqualityIndexScanOverNestedField) { "bounds: {'$_path': [['a.b','a.b',true,true]], 'a.b': [[5,5,true,true]]}}}}}"); } -TEST_F(QueryPlannerWildcardTest, ExprEqCanUseIndex) { +TEST_F(QueryPlannerWildcardTest, ExprEqCannotUseIndex) { addWildcardIndex(BSON("$**" << 1)); runQuery(fromjson("{a: {$_internalExprEq: 1}}")); - assertNumSolutions(1U); - assertSolutionExists( - "{fetch: {filter: null, node: {ixscan: {pattern: {'$_path': 1, a: 1}," - "bounds: {'$_path': [['a','a',true,true]], a: [[1,1,true,true]]}}}}}"); + assertHasOnlyCollscan(); } -TEST_F(QueryPlannerWildcardTest, ExprEqCanUseSparseIndexForEqualityToNull) { +TEST_F(QueryPlannerWildcardTest, ExprEqCannotUseSparseIndexForEqualityToNull) { addWildcardIndex(BSON("$**" << 1)); runQuery(fromjson("{a: {$_internalExprEq: null}}")); - assertNumSolutions(1U); - assertSolutionExists( - "{fetch: {filter: {a: {$_internalExprEq: null}}, node: {ixscan: {pattern: {'$_path': 1, a: " - "1}, bounds: {'$_path': [['a','a',true,true]], a: [[undefined,undefined,true,true], " - "[null,null,true,true]]}}}}}"); + assertHasOnlyCollscan(); } TEST_F(QueryPlannerWildcardTest, PrefixRegex) { diff --git a/src/mongo/db/query/query_request_helper.cpp b/src/mongo/db/query/query_request_helper.cpp index 410c05fcaf0..b713b9d8000 100644 --- a/src/mongo/db/query/query_request_helper.cpp +++ b/src/mongo/db/query/query_request_helper.cpp @@ -85,6 +85,35 @@ Status validateGetMoreCollectionName(StringData collectionName) { return Status::OK(); } +Status validateResumeAfter(const mongo::BSONObj& resumeAfter, bool isClusteredCollection) { + if (resumeAfter.isEmpty()) { + return Status::OK(); + } + + BSONType recordIdType = resumeAfter["$recordId"].type(); + if (resumeAfter.nFields() != 1 || + (recordIdType != BSONType::NumberLong && recordIdType != BSONType::BinData && + recordIdType != BSONType::jstNULL)) { + return Status(ErrorCodes::BadValue, + "Malformed resume token: the '_resumeAfter' object must contain" + " exactly one field named '$recordId', of type NumberLong, BinData " + "or jstNULL."); + } + + // Clustered collections can only have accept '$_resumeAfter' parameter of type + // BinData. Non clustered collections should only accept '$_resumeAfter' of type + // Long. + if ((isClusteredCollection && recordIdType == BSONType::NumberLong) || + (!isClusteredCollection && recordIdType == BSONType::BinData)) { + return Status(ErrorCodes::Error(7738600), + "The '$_resumeAfter parameter must match collection type. Clustered " + "collections only have BinData recordIds, and all other collections" + "have Long recordId."); + } + + return Status::OK(); +} + Status validateFindCommandRequest(const FindCommandRequest& findCommand) { // Min and Max objects must have the same fields. if (!findCommand.getMin().isEmpty() && !findCommand.getMax().isEmpty()) { @@ -127,17 +156,8 @@ Status validateFindCommandRequest(const FindCommandRequest& findCommand) { return Status(ErrorCodes::BadValue, "sort must be unset or {$natural:1} if 'requestResumeToken' is enabled"); } - if (!findCommand.getResumeAfter().isEmpty()) { - if (findCommand.getResumeAfter().nFields() != 1 || - (findCommand.getResumeAfter()["$recordId"].type() != BSONType::NumberLong && - findCommand.getResumeAfter()["$recordId"].type() != BSONType::BinData && - findCommand.getResumeAfter()["$recordId"].type() != BSONType::jstNULL)) { - return Status(ErrorCodes::BadValue, - "Malformed resume token: the '_resumeAfter' object must contain" - " exactly one field named '$recordId', of type NumberLong, BinData " - "or jstNULL."); - } - } + // The $_resumeAfter parameter is checked in 'validateResumeAfter()'. + } else if (!findCommand.getResumeAfter().isEmpty()) { return Status(ErrorCodes::BadValue, "'requestResumeToken' must be true if 'resumeAfter' is" diff --git a/src/mongo/db/query/query_request_helper.h b/src/mongo/db/query/query_request_helper.h index 4d3ec6143c8..c0d8968f563 100644 --- a/src/mongo/db/query/query_request_helper.h +++ b/src/mongo/db/query/query_request_helper.h @@ -63,6 +63,12 @@ static constexpr auto kNaturalSortField = "$natural"; Status validateGetMoreCollectionName(StringData collectionName); /** + * Returns a non-OK status if '$_resumeAfter' is set to an unexpected value, or the wrong type + * determined by the collection type. + */ +Status validateResumeAfter(const mongo::BSONObj& resumeAfter, bool isClusteredCollection); + +/** * Returns a non-OK status if any property of the QR has a bad value (e.g. a negative skip * value) or if there is a bad combination of options (e.g. awaitData is illegal without * tailable). diff --git a/src/mongo/db/query/query_request_test.cpp b/src/mongo/db/query/query_request_test.cpp index 80ed0325802..1493c352b28 100644 --- a/src/mongo/db/query/query_request_test.cpp +++ b/src/mongo/db/query/query_request_test.cpp @@ -278,10 +278,13 @@ TEST(QueryRequestTest, InvalidResumeAfterWrongRecordIdType) { findCommand.setRequestResumeToken(true); // Hint must be explicitly set for the query request to validate. findCommand.setHint(fromjson("{$natural: 1}")); - ASSERT_NOT_OK(query_request_helper::validateFindCommandRequest(findCommand)); + ASSERT_NOT_OK(query_request_helper::validateResumeAfter(findCommand.getResumeAfter(), + false /* isClusteredCollection */)); resumeAfter = BSON("$recordId" << 1LL); findCommand.setResumeAfter(resumeAfter); ASSERT_OK(query_request_helper::validateFindCommandRequest(findCommand)); + ASSERT_OK(query_request_helper::validateResumeAfter(findCommand.getResumeAfter(), + false /* isClusteredCollection */)); } TEST(QueryRequestTest, InvalidResumeAfterExtraField) { @@ -291,7 +294,8 @@ TEST(QueryRequestTest, InvalidResumeAfterExtraField) { findCommand.setRequestResumeToken(true); // Hint must be explicitly set for the query request to validate. findCommand.setHint(fromjson("{$natural: 1}")); - ASSERT_NOT_OK(query_request_helper::validateFindCommandRequest(findCommand)); + ASSERT_NOT_OK(query_request_helper::validateResumeAfter(findCommand.getResumeAfter(), + false /* isClusteredCollection */)); } TEST(QueryRequestTest, ResumeAfterWithHint) { @@ -314,6 +318,8 @@ TEST(QueryRequestTest, ResumeAfterWithSort) { // Hint must be explicitly set for the query request to validate. findCommand.setHint(fromjson("{$natural: 1}")); ASSERT_OK(query_request_helper::validateFindCommandRequest(findCommand)); + ASSERT_OK(query_request_helper::validateResumeAfter(findCommand.getResumeAfter(), + false /* isClusteredCollection */)); findCommand.setSort(fromjson("{a: 1}")); ASSERT_NOT_OK(query_request_helper::validateFindCommandRequest(findCommand)); findCommand.setSort(fromjson("{$natural: 1}")); @@ -329,6 +335,8 @@ TEST(QueryRequestTest, ResumeNoSpecifiedRequestResumeToken) { ASSERT_NOT_OK(query_request_helper::validateFindCommandRequest(findCommand)); findCommand.setRequestResumeToken(true); ASSERT_OK(query_request_helper::validateFindCommandRequest(findCommand)); + ASSERT_OK(query_request_helper::validateResumeAfter(findCommand.getResumeAfter(), + false /* isClusteredCollection */)); } TEST(QueryRequestTest, ExplicitEmptyResumeAfter) { @@ -340,6 +348,8 @@ TEST(QueryRequestTest, ExplicitEmptyResumeAfter) { ASSERT_OK(query_request_helper::validateFindCommandRequest(findCommand)); findCommand.setRequestResumeToken(true); ASSERT_OK(query_request_helper::validateFindCommandRequest(findCommand)); + ASSERT_OK(query_request_helper::validateResumeAfter(findCommand.getResumeAfter(), + false /* isClusteredCollection */)); } // diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp index dc24bc3653a..7b24c1c08a3 100644 --- a/src/mongo/db/query/query_solution.cpp +++ b/src/mongo/db/query/query_solution.cpp @@ -1083,8 +1083,8 @@ bool IndexScanNode::operator==(const IndexScanNode& other) const { // ColumnIndexScanNode // ColumnIndexScanNode::ColumnIndexScanNode(ColumnIndexEntry indexEntry, - std::set<std::string> outputFieldsIn, - std::set<std::string> matchFieldsIn, + OrderedPathSet outputFieldsIn, + OrderedPathSet matchFieldsIn, StringMap<std::unique_ptr<MatchExpression>> filtersByPath, std::unique_ptr<MatchExpression> postAssemblyFilter) : indexEntry(std::move(indexEntry)), diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h index f318884280b..9a50b62c71b 100644 --- a/src/mongo/db/query/query_solution.h +++ b/src/mongo/db/query/query_solution.h @@ -515,8 +515,8 @@ struct CollectionScanNode : public QuerySolutionNodeWithSortSet { struct ColumnIndexScanNode : public QuerySolutionNode { ColumnIndexScanNode(ColumnIndexEntry, - std::set<std::string> outputFields, - std::set<std::string> matchFields, + OrderedPathSet outputFields, + OrderedPathSet matchFields, StringMap<std::unique_ptr<MatchExpression>> filtersByPath, std::unique_ptr<MatchExpression> postAssemblyFilter); @@ -556,11 +556,11 @@ struct ColumnIndexScanNode : public QuerySolutionNode { ColumnIndexEntry indexEntry; // The fields we need to output. Dot separated path names. - std::set<std::string> outputFields; + OrderedPathSet outputFields; // The fields which are referenced by any and all filters - either in 'filtersByPath' or // 'postAssemblyFilter'. - std::set<std::string> matchFields; + OrderedPathSet matchFields; // A column scan can apply a filter to the columns directly while scanning, or to a document // assembled from the scanned columns. @@ -575,7 +575,7 @@ struct ColumnIndexScanNode : public QuerySolutionNode { // A cached copy of the union of the above two field sets which we expect to be frequently asked // for. - std::set<std::string> allFields; + OrderedPathSet allFields; }; /** @@ -1394,11 +1394,13 @@ struct GroupNode : public QuerySolutionNode { for (auto& groupByExprField : groupByExpression->getDependencies().fields) { requiredFields.insert(groupByExprField); } + needWholeDocument = groupByExpression->getDependencies().needWholeDocument; for (auto&& acc : accumulators) { auto argExpr = acc.expr.argument; for (auto& argExprField : argExpr->getDependencies().fields) { requiredFields.insert(argExprField); } + needWholeDocument |= argExpr->getDependencies().needWholeDocument; } } @@ -1434,6 +1436,7 @@ struct GroupNode : public QuerySolutionNode { // the fields in the 'groupByExpressions' and the fields in the input Expressions of the // 'accumulators'. StringSet requiredFields; + bool needWholeDocument = false; // If set to true, generated SBE plan will produce result as BSON object. If false, // 'sbe::Object' is produced instead. diff --git a/src/mongo/db/query/sbe_multi_planner.cpp b/src/mongo/db/query/sbe_multi_planner.cpp index 5f2104a0668..a0d6e975694 100644 --- a/src/mongo/db/query/sbe_multi_planner.cpp +++ b/src/mongo/db/query/sbe_multi_planner.cpp @@ -32,7 +32,6 @@ #include "mongo/db/query/sbe_multi_planner.h" -#include "mongo/db/exec/multi_plan.h" #include "mongo/db/exec/sbe/expressions/expression.h" #include "mongo/db/exec/sbe/values/bson.h" #include "mongo/db/query/collection_query_info.h" diff --git a/src/mongo/db/query/sbe_plan_cache.cpp b/src/mongo/db/query/sbe_plan_cache.cpp index b699387cc47..c8944057639 100644 --- a/src/mongo/db/query/sbe_plan_cache.cpp +++ b/src/mongo/db/query/sbe_plan_cache.cpp @@ -43,8 +43,8 @@ const auto sbePlanCacheDecoration = ServiceContext::declareDecoration<std::unique_ptr<sbe::PlanCache>>(); size_t convertToSizeInBytes(const plan_cache_util::PlanCacheSizeParameter& param) { - constexpr size_t kBytesInMB = 1014 * 1024; - constexpr size_t kMBytesInGB = 1014; + constexpr size_t kBytesInMB = 1024 * 1024; + constexpr size_t kMBytesInGB = 1024; double sizeInMB = param.size; diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp index d35fcc774b0..7a2b2a9bdef 100644 --- a/src/mongo/db/query/sbe_stage_builder.cpp +++ b/src/mongo/db/query/sbe_stage_builder.cpp @@ -769,12 +769,9 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder } // If the slots necessary for performing an index consistency check were not requested in - // 'reqs', then don't pass a pointer to 'iamMap' so 'generateIndexScan' doesn't generate the - // necessary slots. - auto iamMap = &_data.iamMap; - if (!(reqs.has(kSnapshotId) && reqs.has(kIndexId) && reqs.has(kIndexKey))) { - iamMap = nullptr; - } + // 'reqs', then set 'doIndexConsistencyCheck' to false to avoid generating unnecessary logic. + bool doIndexConsistencyCheck = + reqs.has(kSnapshotId) && reqs.has(kIndexId) && reqs.has(kIndexKey); const auto generateIndexScanFunc = ixn->iets.empty() ? generateIndexScan : generateIndexScanWithDynamicBounds; @@ -783,7 +780,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder ixn, indexKeyBitset, _yieldPolicy, - iamMap, + doIndexConsistencyCheck, reqs.has(kIndexKeyPattern)); if (reqs.has(PlanStageSlots::kReturnKey)) { @@ -925,7 +922,6 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder auto [stage, outputs] = build(fn->children[0], childReqs); - auto iamMap = _data.iamMap; uassert(4822880, "RecordId slot is not defined", outputs.has(kRecordId)); uassert( 4953600, "ReturnKey slot is not defined", !reqs.has(kReturnKey) || outputs.has(kReturnKey)); @@ -953,7 +949,6 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder outputs.get(kIndexKey), outputs.get(kIndexKeyPattern), getCurrentCollection(reqs), - std::move(iamMap), root->nodeId(), std::move(relevantSlots), _slotIdGenerator); @@ -2002,27 +1997,17 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder auto outerChild = andSortedNode->children[0]; auto innerChild = andSortedNode->children[1]; - auto [outerStage, outerOutputs] = build(outerChild, childReqs); + auto outerChildReqs = childReqs.copy() + .clear(kSnapshotId) + .clear(kIndexId) + .clear(kIndexKey) + .clear(kIndexKeyPattern); + auto [outerStage, outerOutputs] = build(outerChild, outerChildReqs); auto outerIdSlot = outerOutputs.get(kRecordId); auto outerResultSlot = outerOutputs.get(kResult); auto outerKeySlots = sbe::makeSV(outerIdSlot); auto outerProjectSlots = sbe::makeSV(outerResultSlot); - if (outerOutputs.has(kSnapshotId)) { - outerProjectSlots.push_back(outerOutputs.get(kSnapshotId)); - } - - if (outerOutputs.has(kIndexId)) { - outerProjectSlots.push_back(outerOutputs.get(kIndexId)); - } - - if (outerOutputs.has(kIndexKey)) { - outerProjectSlots.push_back(outerOutputs.get(kIndexKey)); - } - - if (outerOutputs.has(kIndexKeyPattern)) { - outerProjectSlots.push_back(outerOutputs.get(kIndexKeyPattern)); - } auto [innerStage, innerOutputs] = build(innerChild, childReqs); tassert(5073707, "innerOutputs must contain kRecordId slot", innerOutputs.has(kRecordId)); @@ -2100,44 +2085,104 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder namespace { template <typename F> -struct FieldPathVisitor : public SelectiveConstExpressionVisitorBase { +struct FieldPathAndCondPreVisitor : public SelectiveConstExpressionVisitorBase { // To avoid overloaded-virtual warnings. using SelectiveConstExpressionVisitorBase::visit; - FieldPathVisitor(const F& fn) : _fn(fn) {} + FieldPathAndCondPreVisitor(const F& fn, int32_t& nestedCondLevel) + : _fn(fn), _nestedCondLevel(nestedCondLevel) {} void visit(const ExpressionFieldPath* expr) final { - _fn(expr); + _fn(expr, _nestedCondLevel); + } + + void visit(const ExpressionCond* expr) final { + ++_nestedCondLevel; + } + + void visit(const ExpressionSwitch* expr) final { + ++_nestedCondLevel; + } + + void visit(const ExpressionIfNull* expr) final { + ++_nestedCondLevel; + } + + void visit(const ExpressionAnd* expr) final { + ++_nestedCondLevel; + } + + void visit(const ExpressionOr* expr) final { + ++_nestedCondLevel; } F _fn; + // Tracks the number of conditional expressions like $cond or $ifNull that are above us in the + // tree. + int32_t& _nestedCondLevel; +}; + +struct CondPostVisitor : public SelectiveConstExpressionVisitorBase { + // To avoid overloaded-virtual warnings. + using SelectiveConstExpressionVisitorBase::visit; + + CondPostVisitor(int32_t& nestedCondLevel) : _nestedCondLevel(nestedCondLevel) {} + + void visit(const ExpressionCond* expr) final { + --_nestedCondLevel; + } + + void visit(const ExpressionSwitch* expr) final { + --_nestedCondLevel; + } + + void visit(const ExpressionIfNull* expr) final { + --_nestedCondLevel; + } + + void visit(const ExpressionAnd* expr) final { + --_nestedCondLevel; + } + + void visit(const ExpressionOr* expr) final { + --_nestedCondLevel; + } + + int32_t& _nestedCondLevel; }; /** * Walks through the 'expr' expression tree and whenever finds an 'ExpressionFieldPath', calls * the 'fn' function. Type requirement for 'fn' is it must have a const 'ExpressionFieldPath' - * pointer parameter. + * pointer parameter and 'nestedCondLevel' parameter. */ template <typename F> void walkAndActOnFieldPaths(Expression* expr, const F& fn) { - FieldPathVisitor<F> visitor(fn); - ExpressionWalker walker(&visitor, nullptr /*inVisitor*/, nullptr /*postVisitor*/); + int32_t nestedCondLevel = 0; + FieldPathAndCondPreVisitor<F> preVisitor(fn, nestedCondLevel); + CondPostVisitor postVisitor(nestedCondLevel); + ExpressionWalker walker(&preVisitor, nullptr /*inVisitor*/, &postVisitor); expression_walker::walk(expr, &walker); } /** * Checks whether all field paths in 'idExpr' and all accumulator expressions are top-level ones. */ -bool checkAllFieldPathsAreTopLevel(const boost::intrusive_ptr<Expression>& idExpr, - const std::vector<AccumulationStatement>& accStmts) { - auto areAllTopLevelFields = true; +bool areAllFieldPathsOptimizable(const boost::intrusive_ptr<Expression>& idExpr, + const std::vector<AccumulationStatement>& accStmts) { + auto areFieldPathsOptimizable = true; - auto checkFieldPath = [&](const ExpressionFieldPath* fieldExpr) { + auto checkFieldPath = [&](const ExpressionFieldPath* fieldExpr, int32_t nestedCondLevel) { // We optimize neither a field path for the top-level document itself (getPathLength() == 1) // nor a field path that refers to a variable. We can optimize only top-level fields // (getPathLength() == 2). - if (fieldExpr->getFieldPath().getPathLength() != 2 || fieldExpr->isVariableReference()) { - areAllTopLevelFields = false; + // + // The 'nestedCondLevel' being > 0 means that a field path is refered to below conditional + // expressions at the parent $group node, when we cannot optimize field path access and + // therefore, cannot avoid materialization. + if (nestedCondLevel > 0 || fieldExpr->getFieldPath().getPathLength() != 2 || + fieldExpr->isVariableReference()) { + areFieldPathsOptimizable = false; return; } }; @@ -2150,7 +2195,7 @@ bool checkAllFieldPathsAreTopLevel(const boost::intrusive_ptr<Expression>& idExp walkAndActOnFieldPaths(accStmt.expr.argument.get(), checkFieldPath); } - return areAllTopLevelFields; + return areFieldPathsOptimizable; } /** @@ -2176,7 +2221,7 @@ EvalStage optimizeFieldPaths(StageBuilderState& state, auto searchInChildOutputs = !optionalRootSlot.has_value(); auto retEvalStage = std::move(childEvalStage); - walkAndActOnFieldPaths(expr.get(), [&](const ExpressionFieldPath* fieldExpr) { + walkAndActOnFieldPaths(expr.get(), [&](const ExpressionFieldPath* fieldExpr, int32_t) { // We optimize neither a field path for the top-level document itself nor a field path that // refers to a variable instead of calling getField(). if (fieldExpr->getFieldPath().getPathLength() == 1 || fieldExpr->isVariableReference()) { @@ -2307,20 +2352,12 @@ std::tuple<sbe::value::SlotVector, EvalStage, std::unique_ptr<sbe::EExpression>> nodeId, slotIdGenerator); + // The group-by field may end up being 'Nothing' and in that case _id: null will be + // returned. Calling 'makeFillEmptyNull' for the group-by field takes care of that. + auto fillEmptyNullExpr = makeFillEmptyNull(groupByEvalExpr.extractExpr()); sbe::value::SlotId slot; - if (auto isConstIdExpr = dynamic_cast<ExpressionConstant*>(idExpr.get()) != nullptr; - isConstIdExpr) { - std::tie(slot, retEvalStage) = projectEvalExpr( - std::move(groupByEvalExpr), std::move(groupByEvalStage), nodeId, slotIdGenerator); - } else { - // The group-by field may end up being 'Nothing' and in that case _id: null will be - // returned. Calling 'makeFillEmptyNull' for the group-by field takes care of that. - std::tie(slot, retEvalStage) = - projectEvalExpr(makeFillEmptyNull(groupByEvalExpr.extractExpr()), - std::move(groupByEvalStage), - nodeId, - slotIdGenerator); - } + std::tie(slot, retEvalStage) = projectEvalExpr( + std::move(fillEmptyNullExpr), std::move(groupByEvalStage), nodeId, slotIdGenerator); return {sbe::value::SlotVector{slot}, std::move(retEvalStage), nullptr}; } @@ -2332,7 +2369,7 @@ std::tuple<sbe::value::SlotVector, EvalStage> generateAccumulator( const PlanStageSlots& childOutputs, PlanNodeId nodeId, sbe::value::SlotIdGenerator* slotIdGenerator, - sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>>& accSlotToExprMap) { + sbe::SlotExprPairVector& accSlotExprPairs) { // Input fields may need field traversal which ends up being a complex tree. auto evalStage = optimizeFieldPaths( state, accStmt.expr.argument, std::move(childEvalStage), childOutputs, nodeId); @@ -2343,17 +2380,54 @@ std::tuple<sbe::value::SlotVector, EvalStage> generateAccumulator( // One accumulator may be translated to multiple accumulator expressions. For example, The // $avg will have two accumulators expressions, a sum(..) and a count which is implemented // as sum(1). - auto [accExprs, accProjEvalStage] = stage_builder::buildAccumulator( - state, accStmt, std::move(accArgEvalStage), std::move(argExpr), nodeId); + auto collatorSlot = state.data->env->getSlotIfExists("collator"_sd); + auto accExprs = stage_builder::buildAccumulator( + accStmt, std::move(argExpr), collatorSlot, *state.frameIdGenerator); sbe::value::SlotVector aggSlots; for (auto& accExpr : accExprs) { auto slot = slotIdGenerator->generate(); aggSlots.push_back(slot); - accSlotToExprMap.emplace(slot, std::move(accExpr)); + accSlotExprPairs.push_back({slot, std::move(accExpr)}); } - return {std::move(aggSlots), std::move(accProjEvalStage)}; + return {std::move(aggSlots), std::move(accArgEvalStage)}; +} + +/** + * Generate a vector of (inputSlot, mergingExpression) pairs. The slot (whose id is allocated by + * this function) will be used to store spilled partial aggregate values that have been recovered + * from disk and deserialized. The merging expression is an agg function which combines these + * partial aggregates. + * + * Usually the returned vector will be of length 1, but in some cases the MQL accumulation statement + * is implemented by calculating multiple separate aggregates in the SBE plan, which are finalized + * by a subsequent project stage to produce the ultimate value. + */ +sbe::SlotExprPairVector generateMergingExpressions(StageBuilderState& state, + const AccumulationStatement& accStmt, + int numInputSlots) { + tassert(7039555, "'numInputSlots' must be positive", numInputSlots > 0); + auto slotIdGenerator = state.slotIdGenerator; + tassert(7039556, "expected non-null 'slotIdGenerator' pointer", slotIdGenerator); + auto frameIdGenerator = state.frameIdGenerator; + tassert(7039557, "expected non-null 'frameIdGenerator' pointer", frameIdGenerator); + + auto spillSlots = slotIdGenerator->generateMultiple(numInputSlots); + auto collatorSlot = state.data->env->getSlotIfExists("collator"_sd); + auto mergingExprs = + buildCombinePartialAggregates(accStmt, spillSlots, collatorSlot, *frameIdGenerator); + + // Zip the slot vector and expression vector into a vector of pairs. + tassert(7039550, + "expected same number of slots and input exprs", + spillSlots.size() == mergingExprs.size()); + sbe::SlotExprPairVector result; + result.reserve(spillSlots.size()); + for (size_t i = 0; i < spillSlots.size(); ++i) { + result.push_back({spillSlots[i], std::move(mergingExprs[i])}); + } + return result; } std::tuple<std::vector<std::string>, sbe::value::SlotVector, EvalStage> generateGroupFinalStage( @@ -2390,13 +2464,11 @@ std::tuple<std::vector<std::string>, sbe::value::SlotVector, EvalStage> generate auto finalSlots{sbe::value::SlotVector{finalGroupBySlot}}; std::vector<std::string> fieldNames{"_id"}; - auto groupFinalEvalStage = std::move(groupEvalStage); size_t idxAccFirstSlot = dedupedGroupBySlots.size(); for (size_t idxAcc = 0; idxAcc < accStmts.size(); ++idxAcc) { // Gathers field names for the output object from accumulator statements. fieldNames.push_back(accStmts[idxAcc].fieldName); - auto [finalExpr, tempEvalStage] = stage_builder::buildFinalize( - state, accStmts[idxAcc], aggSlotsVec[idxAcc], std::move(groupFinalEvalStage), nodeId); + auto finalExpr = stage_builder::buildFinalize(state, accStmts[idxAcc], aggSlotsVec[idxAcc]); // The final step may not return an expression if it's trivial. For example, $first and // $last's final steps are trivial. @@ -2411,15 +2483,13 @@ std::tuple<std::vector<std::string>, sbe::value::SlotVector, EvalStage> generate // Some accumulator(s) like $avg generate multiple expressions and slots. So, need to // advance this index by the number of those slots for each accumulator. idxAccFirstSlot += aggSlotsVec[idxAcc].size(); - - groupFinalEvalStage = std::move(tempEvalStage); } // Gathers all accumulator results. If there're no project expressions, does not add a project // stage. auto retEvalStage = prjSlotToExprMap.empty() - ? std::move(groupFinalEvalStage) - : makeProject(std::move(groupFinalEvalStage), std::move(prjSlotToExprMap), nodeId); + ? std::move(groupEvalStage) + : makeProject(std::move(groupEvalStage), std::move(prjSlotToExprMap), nodeId); return {std::move(fieldNames), std::move(finalSlots), std::move(retEvalStage)}; } @@ -2478,10 +2548,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder const auto& accStmts = groupNode->accumulators; auto childStageType = childNode->getType(); - auto areAllTopLevelFields = checkAllFieldPathsAreTopLevel(idExpr, accStmts); - - auto childReqs = reqs.copy(); - if (childStageType == StageType::STAGE_GROUP && areAllTopLevelFields) { + auto childReqs = reqs.copy().set(kResult); + if (childStageType == StageType::STAGE_GROUP && areAllFieldPathsOptimizable(idExpr, accStmts)) { // Does not ask the GROUP child for the result slot to avoid unnecessary materialization if // all fields are top-level fields. See the end of this function. For example, GROUP - GROUP // - COLLSCAN case. @@ -2504,17 +2572,28 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder // Translates accumulators which are executed inside the group stage and gets slots for // accumulators. stage_builder::EvalStage accProjEvalStage = std::move(groupByEvalStage); - sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> accSlotToExprMap; + sbe::SlotExprPairVector accSlotExprPairs; std::vector<sbe::value::SlotVector> aggSlotsVec; + // Since partial accumulator state may be spilled to disk and then merged, we must construct not + // only the basic agg expressions for each accumulator, but also agg expressions that are used + // to combine partial aggregates that have been spilled to disk. + sbe::SlotExprPairVector mergingExprs; for (const auto& accStmt : accStmts) { - auto [aggSlots, tempEvalStage] = generateAccumulator(_state, - accStmt, - std::move(accProjEvalStage), - childOutputs, - nodeId, - &_slotIdGenerator, - accSlotToExprMap); - aggSlotsVec.emplace_back(std::move(aggSlots)); + auto [curAggSlots, tempEvalStage] = generateAccumulator(_state, + accStmt, + std::move(accProjEvalStage), + childOutputs, + nodeId, + &_slotIdGenerator, + accSlotExprPairs); + + sbe::SlotExprPairVector curMergingExprs = + generateMergingExpressions(_state, accStmt, curAggSlots.size()); + + aggSlotsVec.emplace_back(std::move(curAggSlots)); + mergingExprs.insert(mergingExprs.end(), + std::make_move_iterator(curMergingExprs.begin()), + std::make_move_iterator(curMergingExprs.end())); accProjEvalStage = std::move(tempEvalStage); } @@ -2525,9 +2604,10 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder // Builds a group stage with accumulator expressions and group-by slot(s). auto groupEvalStage = makeHashAgg(std::move(accProjEvalStage), dedupedGroupBySlots, - std::move(accSlotToExprMap), + std::move(accSlotExprPairs), _state.data->env->getSlotIfExists("collator"_sd), _cq.getExpCtx()->allowDiskUse, + std::move(mergingExprs), nodeId); tassert( diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h index 7abd0e2fa46..d1acc2c765f 100644 --- a/src/mongo/db/query/sbe_stage_builder.h +++ b/src/mongo/db/query/sbe_stage_builder.h @@ -329,9 +329,6 @@ struct PlanStageData { // This holds the output slots produced by SBE plan (resultSlot, recordIdSlot, etc). PlanStageSlots outputs; - // Map from index name to IAM. - StringMap<const IndexAccessMethod*> iamMap; - // The CompileCtx object owns the RuntimeEnvironment. The RuntimeEnvironment owns various // SlotAccessors which are accessed when the SBE plan is executed. sbe::RuntimeEnvironment* env{nullptr}; @@ -379,7 +376,6 @@ private: // RuntimeEnvironment and CompileCtx. void copyFrom(const PlanStageData& other) { outputs = other.outputs; - iamMap = other.iamMap; shouldTrackLatestOplogTimestamp = other.shouldTrackLatestOplogTimestamp; shouldTrackResumeToken = other.shouldTrackResumeToken; shouldUseTailableScan = other.shouldUseTailableScan; diff --git a/src/mongo/db/query/sbe_stage_builder_accumulator.cpp b/src/mongo/db/query/sbe_stage_builder_accumulator.cpp index f01886c4213..0cf745c15a1 100644 --- a/src/mongo/db/query/sbe_stage_builder_accumulator.cpp +++ b/src/mongo/db/query/sbe_stage_builder_accumulator.cpp @@ -42,9 +42,9 @@ namespace mongo::stage_builder { namespace { -std::unique_ptr<sbe::EExpression> wrapMinMaxArg(StageBuilderState& state, - std::unique_ptr<sbe::EExpression> arg) { - return makeLocalBind(state.frameIdGenerator, +std::unique_ptr<sbe::EExpression> wrapMinMaxArg(std::unique_ptr<sbe::EExpression> arg, + sbe::value::FrameIdGenerator& frameIdGenerator) { + return makeLocalBind(&frameIdGenerator, [](sbe::EVariable input) { return sbe::makeE<sbe::EIf>( generateNullOrMissing(input), @@ -54,30 +54,37 @@ std::unique_ptr<sbe::EExpression> wrapMinMaxArg(StageBuilderState& state, std::move(arg)); } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorMin( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorMin( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; - auto collatorSlot = state.data->env->getSlotIfExists("collator"_sd); if (collatorSlot) { aggs.push_back(makeFunction("collMin"_sd, sbe::makeE<sbe::EVariable>(*collatorSlot), - wrapMinMaxArg(state, std::move(arg)))); + wrapMinMaxArg(std::move(arg), frameIdGenerator))); } else { - aggs.push_back(makeFunction("min"_sd, wrapMinMaxArg(state, std::move(arg)))); + aggs.push_back(makeFunction("min"_sd, wrapMinMaxArg(std::move(arg), frameIdGenerator))); } - return {std::move(aggs), std::move(inputStage)}; + return aggs; } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeMin( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsMin( const AccumulationExpression& expr, - const sbe::value::SlotVector& minSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039501, + "partial agg combiner for $min should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + return buildAccumulatorMin(expr, std::move(arg), collatorSlot, frameIdGenerator); +} + +std::unique_ptr<sbe::EExpression> buildFinalizeMin(StageBuilderState& state, + const AccumulationExpression& expr, + const sbe::value::SlotVector& minSlots) { // We can get away with not building a project stage since there's no finalize step but we // will stick the slot into an EVariable in case a $min is one of many group clauses and it // can be combined into a final project stage. @@ -85,76 +92,104 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeMin( str::stream() << "Expected one input slot for finalization of min, got: " << minSlots.size(), minSlots.size() == 1); - return {makeFillEmptyNull(makeVariable(minSlots[0])), std::move(inputStage)}; + return makeFillEmptyNull(makeVariable(minSlots[0])); } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorMax( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorMax( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; - auto collatorSlot = state.data->env->getSlotIfExists("collator"_sd); if (collatorSlot) { aggs.push_back(makeFunction("collMax"_sd, sbe::makeE<sbe::EVariable>(*collatorSlot), - wrapMinMaxArg(state, std::move(arg)))); + wrapMinMaxArg(std::move(arg), frameIdGenerator))); } else { - aggs.push_back(makeFunction("max"_sd, wrapMinMaxArg(state, std::move(arg)))); + aggs.push_back(makeFunction("max"_sd, wrapMinMaxArg(std::move(arg), frameIdGenerator))); } - return {std::move(aggs), std::move(inputStage)}; + return aggs; } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeMax( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsMax( const AccumulationExpression& expr, - const sbe::value::SlotVector& maxSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039502, + "partial agg combiner for $max should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + return buildAccumulatorMax(expr, std::move(arg), collatorSlot, frameIdGenerator); +} + +std::unique_ptr<sbe::EExpression> buildFinalizeMax(StageBuilderState& state, + const AccumulationExpression& expr, + const sbe::value::SlotVector& maxSlots) { tassert(5755100, str::stream() << "Expected one input slot for finalization of max, got: " << maxSlots.size(), maxSlots.size() == 1); - return {makeFillEmptyNull(makeVariable(maxSlots[0])), std::move(inputStage)}; + return makeFillEmptyNull(makeVariable(maxSlots[0])); } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorFirst( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorFirst( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; aggs.push_back(makeFunction("first", makeFillEmptyNull(std::move(arg)))); - return {std::move(aggs), std::move(inputStage)}; + return aggs; } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorLast( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsFirst( + const AccumulationExpression& expr, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039503, + "partial agg combiner for $first should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + return buildAccumulatorFirst(expr, std::move(arg), collatorSlot, frameIdGenerator); +} + +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorLast( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; aggs.push_back(makeFunction("last", makeFillEmptyNull(std::move(arg)))); - return {std::move(aggs), std::move(inputStage)}; + return aggs; } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorAvg( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsLast( + const AccumulationExpression& expr, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039504, + "partial agg combiner for $last should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + return buildAccumulatorLast(expr, std::move(arg), collatorSlot, frameIdGenerator); +} + +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorAvg( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; // 'aggDoubleDoubleSum' will ignore non-numeric values automatically. aggs.push_back(makeFunction("aggDoubleDoubleSum", arg->clone())); // For the counter we need to skip non-numeric values ourselves. - auto addend = makeLocalBind(state.frameIdGenerator, + auto addend = makeLocalBind(&frameIdGenerator, [](sbe::EVariable input) { return sbe::makeE<sbe::EIf>( makeBinaryOp(sbe::EPrimBinary::logicOr, @@ -167,15 +202,27 @@ std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumu auto counterExpr = makeFunction("sum", std::move(addend)); aggs.push_back(std::move(counterExpr)); - return {std::move(aggs), std::move(inputStage)}; + return aggs; } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeAvg( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsAvg( const AccumulationExpression& expr, - const sbe::value::SlotVector& aggSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039539, + "partial agg combiner for $avg should have exactly two input slots", + inputSlots.size() == 2); + + std::vector<std::unique_ptr<sbe::EExpression>> aggs; + aggs.push_back(makeFunction("aggMergeDoubleDoubleSums", makeVariable(inputSlots[0]))); + aggs.push_back(makeFunction("sum", makeVariable(inputSlots[1]))); + return aggs; +} + +std::unique_ptr<sbe::EExpression> buildFinalizeAvg(StageBuilderState& state, + const AccumulationExpression& expr, + const sbe::value::SlotVector& aggSlots) { // Slot 0 contains the accumulated sum, and slot 1 contains the count of summed items. tassert(5754703, str::stream() << "Expected two slots to finalize avg, got: " << aggSlots.size(), @@ -230,7 +277,7 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeAvg( auto partialAvgFinalize = sbe::makeE<sbe::EIf>(std::move(ifCondExpr), std::move(thenExpr), std::move(elseExpr)); - return {std::move(partialAvgFinalize), std::move(inputStage)}; + return partialAvgFinalize; } else { // If we've encountered any numeric input, the counter would contain a positive integer. // Unlike $sum, when there is no numeric input, $avg should return null. @@ -243,27 +290,37 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeAvg( makeFunction("doubleDoubleSumFinalize", makeVariable(aggSlots[0])), makeVariable(aggSlots[1]))); - return {std::move(finalizingExpression), std::move(inputStage)}; + return finalizingExpression; } } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorSum( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorSum( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; aggs.push_back(makeFunction("aggDoubleDoubleSum", std::move(arg))); - return {std::move(aggs), std::move(inputStage)}; + return aggs; } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeSum( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsSum( const AccumulationExpression& expr, - const sbe::value::SlotVector& sumSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039530, + "partial agg combiner for $sum should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + std::vector<std::unique_ptr<sbe::EExpression>> aggs; + aggs.push_back(makeFunction("aggMergeDoubleDoubleSums", std::move(arg))); + return aggs; +} + +std::unique_ptr<sbe::EExpression> buildFinalizeSum(StageBuilderState& state, + const AccumulationExpression& expr, + const sbe::value::SlotVector& sumSlots) { tassert(5755300, str::stream() << "Expected one input slot for finalization of sum, got: " << sumSlots.size(), @@ -292,8 +349,7 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeSum( auto canUseNewPartialResultFormat = fcv.isVersionInitialized() && fcv.isGreaterThanOrEqualTo(multiversion::FeatureCompatibilityVersion::kVersion_6_0); if (canUseNewPartialResultFormat) { - return {makeFunction("doubleDoublePartialSumFinalize", makeVariable(sumSlots[0])), - std::move(inputStage)}; + return makeFunction("doubleDoublePartialSumFinalize", makeVariable(sumSlots[0])); } // To support the sharding behavior, the mongos splits $group into two separate $group @@ -331,43 +387,60 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeSum( input.clone()); }, std::move(sumFinalize)); - return {std::move(partialSumFinalize), std::move(inputStage)}; + return partialSumFinalize; } else { - auto sumFinalize = makeFunction("doubleDoubleSumFinalize", makeVariable(sumSlots[0])); - return {std::move(sumFinalize), std::move(inputStage)}; + return makeFunction("doubleDoubleSumFinalize", makeVariable(sumSlots[0])); } } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorAddToSet( - StageBuilderState& state, - const AccumulationExpression& expr, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorAddToSetHelper( std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + StringData funcName, + boost::optional<sbe::value::SlotId> collatorSlot, + StringData funcNameWithCollator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; const int cap = internalQueryMaxAddToSetBytes.load(); - auto collatorSlot = state.data->env->getSlotIfExists("collator"_sd); if (collatorSlot) { aggs.push_back(makeFunction( - "collAddToSetCapped"_sd, + funcNameWithCollator, sbe::makeE<sbe::EVariable>(*collatorSlot), std::move(arg), makeConstant(sbe::value::TypeTags::NumberInt32, sbe::value::bitcastFrom<int>(cap)))); } else { aggs.push_back(makeFunction( - "addToSetCapped", + funcName, std::move(arg), makeConstant(sbe::value::TypeTags::NumberInt32, sbe::value::bitcastFrom<int>(cap)))); } - return {std::move(aggs), std::move(inputStage)}; + return aggs; +} + +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorAddToSet( + const AccumulationExpression& expr, + std::unique_ptr<sbe::EExpression> arg, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + return buildAccumulatorAddToSetHelper( + std::move(arg), "addToSetCapped"_sd, collatorSlot, "collAddToSetCapped"_sd); +} + +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsAddToSet( + const AccumulationExpression& expr, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039506, + "partial agg combiner for $addToSet should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + return buildAccumulatorAddToSetHelper( + std::move(arg), "aggSetUnionCapped"_sd, collatorSlot, "aggCollSetUnionCapped"_sd); } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeCappedAccumulator( +std::unique_ptr<sbe::EExpression> buildFinalizeCappedAccumulator( StageBuilderState& state, const AccumulationExpression& expr, - const sbe::value::SlotVector& accSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { + const sbe::value::SlotVector& accSlots) { tassert(6526500, str::stream() << "Expected one input slot for finalization of capped accumulator, got: " << accSlots.size(), @@ -382,33 +455,62 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeCappedAccum makeConstant(sbe::value::TypeTags::NumberInt32, static_cast<int>(sbe::vm::AggArrayWithSize::kValues))); - return {std::move(pushFinalize), std::move(inputStage)}; + return pushFinalize; } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorPush( - StageBuilderState& state, - const AccumulationExpression& expr, - std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorPushHelper( + std::unique_ptr<sbe::EExpression> arg, StringData aggFuncName) { const int cap = internalQueryMaxPushBytes.load(); std::vector<std::unique_ptr<sbe::EExpression>> aggs; aggs.push_back(makeFunction( - "addToArrayCapped"_sd, + aggFuncName, std::move(arg), makeConstant(sbe::value::TypeTags::NumberInt32, sbe::value::bitcastFrom<int>(cap)))); - return {std::move(aggs), std::move(inputStage)}; + return aggs; } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorStdDev( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorPush( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + return buildAccumulatorPushHelper(std::move(arg), "addToArrayCapped"_sd); +} + +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsPush( + const AccumulationExpression& expr, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039505, + "partial agg combiner for $push should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + return buildAccumulatorPushHelper(std::move(arg), "aggConcatArraysCapped"_sd); +} + +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorStdDev( + const AccumulationExpression& expr, + std::unique_ptr<sbe::EExpression> arg, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; aggs.push_back(makeFunction("aggStdDev", std::move(arg))); - return {std::move(aggs), std::move(inputStage)}; + return aggs; +} + +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsStdDev( + const AccumulationExpression& expr, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039540, + "partial agg combiner for stddev should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + std::vector<std::unique_ptr<sbe::EExpression>> aggs; + aggs.push_back(makeFunction("aggMergeStdDevs", std::move(arg))); + return aggs; } std::unique_ptr<sbe::EExpression> buildFinalizePartialStdDev(sbe::value::SlotId stdDevSlot) { @@ -439,69 +541,75 @@ std::unique_ptr<sbe::EExpression> buildFinalizePartialStdDev(sbe::value::SlotId static_cast<int>(sbe::vm::AggStdDevValueElems::kCount)))}); } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeStdDevPop( +std::unique_ptr<sbe::EExpression> buildFinalizeStdDevPop( StageBuilderState& state, const AccumulationExpression& expr, - const sbe::value::SlotVector& stdDevSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { + const sbe::value::SlotVector& stdDevSlots) { tassert(5755204, str::stream() << "Expected one input slot for finalization of stdDevPop, got: " << stdDevSlots.size(), stdDevSlots.size() == 1); if (state.needsMerge) { - return {buildFinalizePartialStdDev(stdDevSlots[0]), std::move(inputStage)}; + return buildFinalizePartialStdDev(stdDevSlots[0]); } else { auto stdDevPopFinalize = makeFunction("stdDevPopFinalize", makeVariable(stdDevSlots[0])); - return {std::move(stdDevPopFinalize), std::move(inputStage)}; + return stdDevPopFinalize; } } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalizeStdDevSamp( +std::unique_ptr<sbe::EExpression> buildFinalizeStdDevSamp( StageBuilderState& state, const AccumulationExpression& expr, - const sbe::value::SlotVector& stdDevSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { + const sbe::value::SlotVector& stdDevSlots) { tassert(5755209, str::stream() << "Expected one input slot for finalization of stdDevSamp, got: " << stdDevSlots.size(), stdDevSlots.size() == 1); if (state.needsMerge) { - return {buildFinalizePartialStdDev(stdDevSlots[0]), std::move(inputStage)}; + return buildFinalizePartialStdDev(stdDevSlots[0]); } else { - auto stdDevSampFinalize = makeFunction("stdDevSampFinalize", makeVariable(stdDevSlots[0])); - return {std::move(stdDevSampFinalize), std::move(inputStage)}; + return makeFunction("stdDevSampFinalize", makeVariable(stdDevSlots[0])); } } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulatorMergeObjects( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulatorMergeObjects( const AccumulationExpression& expr, std::unique_ptr<sbe::EExpression> arg, - EvalStage inputStage, - PlanNodeId planNodeId) { + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector<std::unique_ptr<sbe::EExpression>> aggs; - auto filterExpr = makeLocalBind( - state.frameIdGenerator, - [](sbe::EVariable input) { - return makeBinaryOp( - sbe::EPrimBinary::logicOr, - generateNullOrMissing(input), - makeBinaryOp(sbe::EPrimBinary::logicOr, - makeFunction("isObject", input.clone()), - sbe::makeE<sbe::EFail>(ErrorCodes::Error{5911200}, - "$mergeObjects only supports objects"))); - }, - arg->clone()); - - inputStage = makeFilter<false>(std::move(inputStage), std::move(filterExpr), planNodeId); + auto filterExpr = + makeLocalBind(&frameIdGenerator, + [](sbe::EVariable input) { + auto typeCheckExpr = + makeBinaryOp(sbe::EPrimBinary::logicOr, + generateNullOrMissing(input), + makeFunction("isObject", input.clone())); + return sbe::makeE<sbe::EIf>( + std::move(typeCheckExpr), + makeFunction("mergeObjects", input.clone()), + sbe::makeE<sbe::EFail>(ErrorCodes::Error{5911200}, + "$mergeObjects only supports objects")); + }, + std::move(arg)); + + aggs.push_back(std::move(filterExpr)); + return aggs; +} - aggs.push_back(makeFunction("mergeObjects", std::move(arg))); - return {std::move(aggs), std::move(inputStage)}; +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggsMergeObjects( + const AccumulationExpression& expr, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + tassert(7039507, + "partial agg combiner for $mergeObjects should have exactly one input slot", + inputSlots.size() == 1); + auto arg = makeVariable(inputSlots[0]); + return buildAccumulatorMergeObjects(expr, std::move(arg), collatorSlot, frameIdGenerator); } }; // namespace @@ -516,19 +624,16 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildArgument( return {argExpr.extractExpr(), std::move(outStage)}; } -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulator( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulator( const AccumulationStatement& acc, - EvalStage inputStage, - std::unique_ptr<sbe::EExpression> inputExpr, - PlanNodeId planNodeId) { - using BuildAccumulatorFn = - std::function<std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage>( - StageBuilderState&, - const AccumulationExpression&, - std::unique_ptr<sbe::EExpression>, - EvalStage, - PlanNodeId)>; + std::unique_ptr<sbe::EExpression> argExpr, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + using BuildAccumulatorFn = std::function<std::vector<std::unique_ptr<sbe::EExpression>>( + const AccumulationExpression&, + std::unique_ptr<sbe::EExpression>, + boost::optional<sbe::value::SlotId>, + sbe::value::FrameIdGenerator&)>; static const StringDataMap<BuildAccumulatorFn> kAccumulatorBuilders = { {AccumulatorMin::kName, &buildAccumulatorMin}, @@ -550,25 +655,51 @@ std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumu kAccumulatorBuilders.find(accExprName) != kAccumulatorBuilders.end()); return std::invoke(kAccumulatorBuilders.at(accExprName), - state, acc.expr, - std::move(inputExpr), - std::move(inputStage), - planNodeId); + std::move(argExpr), + collatorSlot, + frameIdGenerator); } -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalize( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggregates( const AccumulationStatement& acc, - const sbe::value::SlotVector& aggSlots, - EvalStage inputStage, - PlanNodeId planNodeId) { - using BuildFinalizeFn = std::function<std::pair<std::unique_ptr<sbe::EExpression>, EvalStage>( - StageBuilderState&, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator& frameIdGenerator) { + using BuildAggCombinerFn = std::function<std::vector<std::unique_ptr<sbe::EExpression>>( const AccumulationExpression&, - sbe::value::SlotVector, - EvalStage, - PlanNodeId)>; + const sbe::value::SlotVector&, + boost::optional<sbe::value::SlotId>, + sbe::value::FrameIdGenerator&)>; + + static const StringDataMap<BuildAggCombinerFn> kAggCombinerBuilders = { + {AccumulatorAddToSet::kName, &buildCombinePartialAggsAddToSet}, + {AccumulatorAvg::kName, &buildCombinePartialAggsAvg}, + {AccumulatorFirst::kName, &buildCombinePartialAggsFirst}, + {AccumulatorLast::kName, &buildCombinePartialAggsLast}, + {AccumulatorMax::kName, &buildCombinePartialAggsMax}, + {AccumulatorMergeObjects::kName, &buildCombinePartialAggsMergeObjects}, + {AccumulatorMin::kName, &buildCombinePartialAggsMin}, + {AccumulatorPush::kName, &buildCombinePartialAggsPush}, + {AccumulatorStdDevPop::kName, &buildCombinePartialAggsStdDev}, + {AccumulatorStdDevSamp::kName, &buildCombinePartialAggsStdDev}, + {AccumulatorSum::kName, &buildCombinePartialAggsSum}, + }; + + auto accExprName = acc.expr.name; + uassert(7039500, + str::stream() << "Unsupported Accumulator in SBE accumulator builder: " << accExprName, + kAggCombinerBuilders.find(accExprName) != kAggCombinerBuilders.end()); + + return std::invoke( + kAggCombinerBuilders.at(accExprName), acc.expr, inputSlots, collatorSlot, frameIdGenerator); +} + +std::unique_ptr<sbe::EExpression> buildFinalize(StageBuilderState& state, + const AccumulationStatement& acc, + const sbe::value::SlotVector& aggSlots) { + using BuildFinalizeFn = std::function<std::unique_ptr<sbe::EExpression>( + StageBuilderState&, const AccumulationExpression&, sbe::value::SlotVector)>; static const StringDataMap<BuildFinalizeFn> kAccumulatorBuilders = { {AccumulatorMin::kName, &buildFinalizeMin}, @@ -590,10 +721,10 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalize( kAccumulatorBuilders.find(accExprName) != kAccumulatorBuilders.end()); if (auto fn = kAccumulatorBuilders.at(accExprName); fn) { - return std::invoke(fn, state, acc.expr, aggSlots, std::move(inputStage), planNodeId); + return std::invoke(fn, state, acc.expr, aggSlots); } else { // nullptr for 'EExpression' signifies that no final project is necessary. - return {nullptr, std::move(inputStage)}; + return nullptr; } } } // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_accumulator.h b/src/mongo/db/query/sbe_stage_builder_accumulator.h index 508a34f07f0..7477ffb2a3a 100644 --- a/src/mongo/db/query/sbe_stage_builder_accumulator.h +++ b/src/mongo/db/query/sbe_stage_builder_accumulator.h @@ -51,24 +51,32 @@ std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildArgument( /** * Translates an input AccumulationStatement into an SBE EExpression for accumulation expressions. - * The 'stage' parameter provides the input subtree to build on top of. */ -std::pair<std::vector<std::unique_ptr<sbe::EExpression>>, EvalStage> buildAccumulator( - StageBuilderState& state, +std::vector<std::unique_ptr<sbe::EExpression>> buildAccumulator( const AccumulationStatement& acc, - EvalStage stage, std::unique_ptr<sbe::EExpression> argExpr, - PlanNodeId planNodeId); + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator&); + +/** + * When SBE hash aggregation spills to disk, it spills partial aggregates which need to be combined + * later. This function returns the expressions that can be used to combine partial aggregates for + * the given accumulator 'acc'. The aggregate-of-aggregates will be stored in a slots owned by the + * hash agg stage, while the new partial aggregates to combine can be read from the given + * 'inputSlots'. + */ +std::vector<std::unique_ptr<sbe::EExpression>> buildCombinePartialAggregates( + const AccumulationStatement& acc, + const sbe::value::SlotVector& inputSlots, + boost::optional<sbe::value::SlotId> collatorSlot, + sbe::value::FrameIdGenerator&); /** * Translates an input AccumulationStatement into an SBE EExpression that represents an * AccumulationStatement's finalization step. The 'stage' parameter provides the input subtree to * build on top of. */ -std::pair<std::unique_ptr<sbe::EExpression>, EvalStage> buildFinalize( - StageBuilderState& state, - const AccumulationStatement& acc, - const sbe::value::SlotVector& aggSlots, - EvalStage stage, - PlanNodeId planNodeId); +std::unique_ptr<sbe::EExpression> buildFinalize(StageBuilderState& state, + const AccumulationStatement& acc, + const sbe::value::SlotVector& aggSlots); } // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp b/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp index e4fd17df36d..732b04f62a8 100644 --- a/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp +++ b/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp @@ -27,15 +27,22 @@ * it in the license file. */ +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + #include "mongo/platform/basic.h" #include <fmt/printf.h> +#include "mongo/db/exec/sbe/expression_test_base.h" +#include "mongo/db/exec/sbe/values/value_printer.h" #include "mongo/db/pipeline/document_source_group.h" #include "mongo/db/pipeline/expression_context_for_test.h" #include "mongo/db/query/collation/collator_interface_mock.h" #include "mongo/db/query/query_solution.h" +#include "mongo/db/query/sbe_stage_builder_accumulator.h" #include "mongo/db/query/sbe_stage_builder_test_fixture.h" +#include "mongo/idl/server_parameter_test_util.h" +#include "mongo/logv2/log.h" #include "mongo/unittest/unittest.h" namespace mongo { @@ -1677,4 +1684,724 @@ TEST_F(SbeStageBuilderGroupTest, SbeIncompatibleExpressionInGroup) { } } +/** + * A test fixture designed to test that the expressions generated to combine partial aggregates + * that have been spilled to disk work correctly. We use 'EExpressionTestFixture' rather than + * something like 'SbeStageBuilderTestFixture' so that the expressions can be tested in isolation, + * without actually requiring a hash agg stage or without actually spilling any data to disk. + */ +class SbeStageBuilderGroupAggCombinerTest : public sbe::EExpressionTestFixture { +public: + explicit SbeStageBuilderGroupAggCombinerTest() + : _expCtx{make_intrusive<ExpressionContextForTest>()}, + _inputSlotId{bindAccessor(&_inputAccessor)}, + _collatorSlotId{bindAccessor(&_collatorAccessor)} {} + + AccumulationStatement makeAccumulationStatement(StringData accumName) { + return makeAccumulationStatement(BSON("unused" << BSON(accumName << "unused"))); + } + + AccumulationStatement makeAccumulationStatement(BSONObj accumulationStmt) { + _accumulationStmtBson = std::move(accumulationStmt); + VariablesParseState vps = _expCtx->variablesParseState; + return AccumulationStatement::parseAccumulationStatement( + _expCtx.get(), _accumulationStmtBson.firstElement(), vps); + } + + /** + * Convenience method for producing bytecode which combines partial aggregates for the given + * 'AccumulationStatement'. + * + * Requires that accumulation statement results in a single aggregate with one input and one + * output. Furthermore, cannot be used when the test case involves a non-simple collation. + */ + std::unique_ptr<sbe::vm::CodeFragment> compileSingleInputNoCollator( + const AccumulationStatement& accStatement) { + auto exprs = stage_builder::buildCombinePartialAggregates( + accStatement, {_inputSlotId}, boost::none, _frameIdGenerator); + ASSERT_EQ(exprs.size(), 1u); + _expr = std::move(exprs[0]); + + return compileAggExpression(*_expr, &_aggAccessor); + } + + /** + * Verifies that executing the bytecode ('code') for combining partial aggregates for $group + * spilling produces the 'expected' outputs given 'inputs'. + * + * The inputs and expected outputs are expressed as BSON arrays as a convenience to the caller, + * and should have the same length. The bytecode is executed over each element of 'inputs' + * one-by-one, with the result stored into a slot holding the aggregate value. At each step, + * this function asserts that the current aggregate value is equal to the matching element in + * 'expected'. + * + * The string "MISSING" can be used as a sentinel in either 'inputs' or 'outputs' in order to + * represent the Nothing value (since nothingness cannot literally be stored in a BSON array). + */ + void aggregateAndAssertResults(BSONArray inputs, + BSONArray expected, + const sbe::vm::CodeFragment* code) { + auto [inputTag, inputVal] = makeArray(inputs); + auto [expectedTag, expectedVal] = makeArray(expected); + return aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, code); + } + + /** + * Verifies that executing the bytecode ('code') for combining partial aggregates for $group + * spilling produces the 'expectedVal' outputs given 'inputsVal'. Assumes ownership of both + * 'expectedVal' and 'inputsVal'. + * + * Identical to the overload above, except the inputs and expected outputs are provided as SBE + * arrays rather than BSON arrays. This is useful if the caller needs to construct input and + * output ways in a special way that cannot be achieved by trivial conversion from BSON. + */ + void aggregateAndAssertResults(sbe::value::TypeTags inputTag, + sbe::value::Value inputVal, + sbe::value::TypeTags expectedTag, + sbe::value::Value expectedVal, + const sbe::vm::CodeFragment* code) { + // Make sure we are starting from a clean state. + _inputAccessor.reset(); + _aggAccessor.reset(); + + sbe::value::ValueGuard inputGuard{inputTag, inputVal}; + sbe::value::ValueGuard expectedGuard{expectedTag, expectedVal}; + + sbe::value::ArrayEnumerator inputEnumerator{inputTag, inputVal}; + sbe::value::ArrayEnumerator expectedEnumerator{expectedTag, expectedVal}; + + // Aggregate the inputs one-by-one, and at each step validate that the resulting accumulator + // state is as expected. + int index = 0; + while (!inputEnumerator.atEnd()) { + ASSERT_FALSE(expectedEnumerator.atEnd()); + auto [nextInputTag, nextInputVal] = inputEnumerator.getViewOfValue(); + + // Feed in the input value, treating "MISSING" as a special sentinel to indicate the + // Nothing value. + if (sbe::value::isString(nextInputTag) && + sbe::value::getStringView(nextInputTag, nextInputVal) == "MISSING"_sd) { + _inputAccessor.reset(); + } else { + auto [copyTag, copyVal] = sbe::value::copyValue(nextInputTag, nextInputVal); + _inputAccessor.reset(true, copyTag, copyVal); + } + + auto [outputTag, outputVal] = runCompiledExpression(code); + + // Validate that the output value equals the expected value, and then put the output + // value into the slot that holds the accumulation state. + auto [expectedOutputTag, expectedOutputValue] = expectedEnumerator.getViewOfValue(); + if (sbe::value::isString(expectedOutputTag) && + sbe::value::getStringView(expectedOutputTag, expectedOutputValue) == "MISSING"_sd) { + expectedOutputTag = sbe::value::TypeTags::Nothing; + expectedOutputValue = 0; + } + auto [compareTag, compareValue] = sbe::value::compareValue( + outputTag, outputVal, expectedOutputTag, expectedOutputValue); + if (compareTag != sbe::value::TypeTags::NumberInt32 || compareValue != 0) { + // The test failed, but dump the actual and expected values to the logs for ease of + // debugging. + str::stream actualBuilder; + auto actualPrinter = makeValuePrinter(actualBuilder); + actualPrinter.writeValueToStream(outputTag, outputVal); + + str::stream expectedBuilder; + auto expectedPrinter = makeValuePrinter(expectedBuilder); + expectedPrinter.writeValueToStream(expectedOutputTag, expectedOutputValue); + + LOGV2(7039529, + "Actual value not equal to expected value", + "actual"_attr = actualBuilder, + "expected"_attr = expectedBuilder, + "index"_attr = index); + FAIL("accumulator did not have expected value"); + } + + _aggAccessor.reset(true, outputTag, outputVal); + + inputEnumerator.advance(); + expectedEnumerator.advance(); + ++index; + } + } + + /** + * A helper for converting a sequence of accumulator states for $push or $addToSet into the + * corresponding SBE value. + */ + enum class Accumulator { kPush, kAddToSet }; + std::pair<sbe::value::TypeTags, sbe::value::Value> makeArrayAccumVal(BSONArray bsonArray, + Accumulator accumType) { + auto [resultTag, resultVal] = sbe::value::makeNewArray(); + sbe::value::ValueGuard resultGuard{resultTag, resultVal}; + auto resultArr = sbe::value::getArrayView(resultVal); + + for (auto&& elt : bsonArray) { + ASSERT(elt.type() == BSONType::Array); + + BSONObjIterator arrayIt{elt.embeddedObject()}; + ASSERT_TRUE(arrayIt.more()); + auto firstElt = arrayIt.next(); + ASSERT(firstElt.type() == BSONType::Array); + BSONArray partialBsonArr{firstElt.embeddedObject()}; + + ASSERT_TRUE(arrayIt.more()); + auto secondElt = arrayIt.next(); + ASSERT(secondElt.isNumber()); + int64_t size = secondElt.safeNumberLong(); + + ASSERT_FALSE(arrayIt.more()); + + // Each partial aggregate is a two-element array whose first element is the partial + // $push result (itself an array) and whose second element is the size. + auto [partialAggTag, partialAggVal] = sbe::value::makeNewArray(); + auto partialAggArr = sbe::value::getArrayView(partialAggVal); + + auto [pushedValsTag, pushedValsVal] = accumType == Accumulator::kPush + ? makeArray(partialBsonArr) + : makeArraySet(partialBsonArr); + partialAggArr->push_back(pushedValsTag, pushedValsVal); + + partialAggArr->push_back(sbe::value::TypeTags::NumberInt64, + sbe::value::bitcastFrom<int64_t>(size)); + + resultArr->push_back(partialAggTag, partialAggVal); + } + + resultGuard.reset(); + return {resultTag, resultVal}; + } + + /** + * Given the name of an SBE agg function ('aggFuncName') and an array of values expressed as a + * BSON array, aggregates the values inside the array and returns the resulting SBE value. + */ + std::pair<sbe::value::TypeTags, sbe::value::Value> makeOnePartialAggregate( + StringData aggFuncName, BSONArray valuesToAgg) { + // Make sure we are starting from a clean state. + _inputAccessor.reset(); + _aggAccessor.reset(); + + // Construct an expression which calls the given agg function, aggregating the values in + // '_inputSlotId'. + auto expr = + stage_builder::makeFunction(aggFuncName, stage_builder::makeVariable(_inputSlotId)); + auto code = compileAggExpression(*expr, &_aggAccessor); + + // Find the first element by skipping the length. + const char* bsonElt = valuesToAgg.objdata() + 4; + const char* bsonEnd = bsonElt + valuesToAgg.objsize(); + while (*bsonElt != 0) { + auto fieldName = sbe::bson::fieldNameView(bsonElt); + + // Convert the BSON value to an SBE value and put it inside the input slot. + auto [tag, val] = sbe::bson::convertFrom<false>(bsonElt, bsonEnd, fieldName.size()); + _inputAccessor.reset(true, tag, val); + + // Run the agg function, and put the result in the slot holding the aggregate value. + auto [outputTag, outputVal] = runCompiledExpression(code.get()); + _aggAccessor.reset(true, outputTag, outputVal); + + bsonElt = sbe::bson::advance(bsonElt, fieldName.size()); + } + + return _aggAccessor.copyOrMoveValue(); + } + + /** + * Returns an SBE array which contains a sequence of partial aggregate values. Useful for + * constructing a sequence of partial aggregates when those partial aggregates are not trivial + * to describe using BSON. The input to this function is a BSON array of BSON arrays; each of + * the inner arrays is aggregated using the given 'aggFuncName' in order to produce the output + * SBE array. + * + * As an example, suppose the agg function is a simple sum. Given the input + * + * [[8, 1, 5], [6], [2,3]] + * + * the output will be the SBE array [14, 6, 5]. + */ + std::pair<sbe::value::TypeTags, sbe::value::Value> makePartialAggArray( + StringData aggFuncName, BSONArray arrayOfArrays) { + auto [arrTag, arrVal] = sbe::value::makeNewArray(); + sbe::value::ValueGuard guard{arrTag, arrVal}; + + auto arr = sbe::value::getArrayView(arrVal); + + for (auto&& element : arrayOfArrays) { + ASSERT(element.type() == BSONType::Array); + auto [tag, val] = + makeOnePartialAggregate(aggFuncName, BSONArray{element.embeddedObject()}); + arr->push_back(tag, val); + } + + guard.reset(); + return {arrTag, arrVal}; + } + +protected: + sbe::value::FrameIdGenerator _frameIdGenerator; + boost::intrusive_ptr<ExpressionContextForTest> _expCtx; + + // Accessor and corresponding slot id that holds the input to the agg expression. Each time we + // "turn the crank" this will hold the next partial aggregate to be aggregated into + // '_aggAccessor'. + sbe::value::OwnedValueAccessor _inputAccessor; + sbe::value::SlotId _inputSlotId; + + // The accessor which holds the final output resulting from combining all partial outputs. We + // check that the intermediate value is as expected after every turn of the crank. + sbe::value::OwnedValueAccessor _aggAccessor; + + sbe::value::OwnedValueAccessor _collatorAccessor; + sbe::value::SlotId _collatorSlotId; + +private: + template <typename Stream> + sbe::value::ValuePrinter<Stream> makeValuePrinter(Stream& stream) { + return sbe::value::ValuePrinters::make(stream, + sbe::PrintOptions().useTagForAmbiguousValues(true)); + } + + BSONObj _accumulationStmtBson; + std::unique_ptr<sbe::EExpression> _expr; +}; + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsMin) { + auto accStatement = makeAccumulationStatement("$min"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto inputValues = BSON_ARRAY(8 << 7 << 9 << BSONNULL << 6); + auto expectedAggStates = BSON_ARRAY(8 << 7 << 7 << 7 << 6); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); + + // Test that Nothing values are treated as expected. + inputValues = BSON_ARRAY("MISSING" << 9 << 7 << "MISSING" << 6); + expectedAggStates = BSON_ARRAY("MISSING" << 9 << 7 << 7 << 6); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsMinWithCollation) { + auto accStatement = makeAccumulationStatement("$min"_sd); + + auto exprs = stage_builder::buildCombinePartialAggregates( + accStatement, {_inputSlotId}, {_collatorSlotId}, _frameIdGenerator); + ASSERT_EQ(exprs.size(), 1u); + auto expr = std::move(exprs[0]); + + CollatorInterfaceMock collator{CollatorInterfaceMock::MockType::kReverseString}; + _collatorAccessor.reset(false, + sbe::value::TypeTags::collator, + sbe::value::bitcastFrom<const CollatorInterface*>(&collator)); + + auto compiledExpr = compileAggExpression(*expr, &_aggAccessor); + + // The strings in reverse have the opposite ordering as compared to forwards. + auto inputValues = BSON_ARRAY("az" + << "by" + << "cx"); + auto expectedAggStates = BSON_ARRAY("az" + << "by" + << "cx"); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsMax) { + auto accStatement = makeAccumulationStatement("$max"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto inputValues = BSON_ARRAY(3 << 1 << 4 << BSONNULL << 8); + auto expectedAggStates = BSON_ARRAY(3 << 3 << 4 << 4 << 8); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); + + // Test that Nothing values are treated as expected. + inputValues = BSON_ARRAY("MISSING" << 7 << 9 << "MISSING" << 10); + expectedAggStates = BSON_ARRAY("MISSING" << 7 << 9 << 9 << 10); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsMaxWithCollation) { + auto accStatement = makeAccumulationStatement("$max"_sd); + + auto exprs = stage_builder::buildCombinePartialAggregates( + accStatement, {_inputSlotId}, {_collatorSlotId}, _frameIdGenerator); + ASSERT_EQ(exprs.size(), 1u); + auto expr = std::move(exprs[0]); + + CollatorInterfaceMock collator{CollatorInterfaceMock::MockType::kReverseString}; + _collatorAccessor.reset(false, + sbe::value::TypeTags::collator, + sbe::value::bitcastFrom<const CollatorInterface*>(&collator)); + + auto compiledExpr = compileAggExpression(*expr, &_aggAccessor); + + // The strings in reverse have the opposite ordering as compared to forwards. + auto inputValues = BSON_ARRAY("cx" + << "by" + << "az"); + auto expectedAggStates = BSON_ARRAY("cx" + << "by" + << "az"); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsFirst) { + auto accStatement = makeAccumulationStatement("$first"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto inputValues = BSON_ARRAY(3 << 1 << BSONNULL << "MISSING" << 8); + auto expectedAggStates = BSON_ARRAY(3 << 3 << 3 << 3 << 3); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); + + // When the first value is missing, the resulting value is a literal null. + inputValues = BSON_ARRAY("MISSING" << 1 << BSONNULL << "MISSING" << 8); + expectedAggStates = BSON_ARRAY(BSONNULL << BSONNULL << BSONNULL << BSONNULL << BSONNULL); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsLast) { + auto accStatement = makeAccumulationStatement("$last"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto inputValues = BSON_ARRAY(3 << 1 << BSONNULL << "MISSING" << 8); + auto expectedAggStates = BSON_ARRAY(3 << 1 << BSONNULL << BSONNULL << 8); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsPush) { + auto accStatement = makeAccumulationStatement("$push"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto [inputValuesTag, inputValuesVal] = makeArrayAccumVal( + BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(5 << 4 << 3) << 10) + << BSON_ARRAY(BSON_ARRAY(2 << 1) << 20) << BSON_ARRAY(BSONArray{} << 0)), + Accumulator::kPush); + auto [expectedTag, expectedVal] = + makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(5 << 4 << 3) << 10) + << BSON_ARRAY(BSON_ARRAY(5 << 4 << 3 << 2 << 1) << 30) + << BSON_ARRAY(BSON_ARRAY(5 << 4 << 3 << 2 << 1) << 30)), + Accumulator::kPush); + aggregateAndAssertResults( + inputValuesTag, inputValuesVal, expectedTag, expectedVal, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsPushThrowsWhenExceedingSizeLimit) { + auto accStatement = makeAccumulationStatement("$push"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + // If we inject a very large size, we expect the accumulator to throw. This cap prevents the + // accumulator from consuming too much memory. + const int64_t largeSize = 1000 * 1000 * 1000; + + auto input = makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(5 << 4) << 3) + << BSON_ARRAY(BSON_ARRAY(2 << 1) << largeSize)), + Accumulator::kPush); + auto expected = makeArrayAccumVal( + BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(5 << 4) << 3) << BSON_ARRAY(BSON_ARRAY("unused") << -1)), + Accumulator::kPush); + ASSERT_THROWS_CODE( + aggregateAndAssertResults( + input.first, input.second, expected.first, expected.second, compiledExpr.get()), + DBException, + ErrorCodes::ExceededMemoryLimit); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsAddToSet) { + auto accStatement = makeAccumulationStatement("$addToSet"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto [inputValuesTag, inputValuesVal] = + makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(3 << 4 << 5) << 10) + << BSON_ARRAY(BSON_ARRAY(1 << 3 << 5 << 8) << 20) + << BSON_ARRAY(BSONArray{} << 0)), + Accumulator::kAddToSet); + + // Each SBE value is 8 bytes and its tag is 1 byte. So we expect each unique element's size to + // be calculated as 9 bytes. The sizes from the partial aggregates end up getting ignored, and + // the total size is recalculated, since we cannot predict the size of the set union in advance. + auto [expectedTag, expectedVal] = + makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(3 << 4 << 5) << 27) + << BSON_ARRAY(BSON_ARRAY(1 << 3 << 4 << 5 << 8) << 45) + << BSON_ARRAY(BSON_ARRAY(1 << 3 << 4 << 5 << 8) << 45)), + Accumulator::kAddToSet); + aggregateAndAssertResults( + inputValuesTag, inputValuesVal, expectedTag, expectedVal, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsAddToSetWithCollation) { + auto accStatement = makeAccumulationStatement("$addToSet"_sd); + + auto exprs = stage_builder::buildCombinePartialAggregates( + accStatement, {_inputSlotId}, {_collatorSlotId}, _frameIdGenerator); + ASSERT_EQ(exprs.size(), 1u); + auto expr = std::move(exprs[0]); + + CollatorInterfaceMock collator{CollatorInterfaceMock::MockType::kToLowerString}; + _collatorAccessor.reset(false, + sbe::value::TypeTags::collator, + sbe::value::bitcastFrom<const CollatorInterface*>(&collator)); + + auto compiledExpr = compileAggExpression(*expr, &_aggAccessor); + + auto [inputValuesTag, inputValuesVal] = + makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY("foo" + << "bar") + << 10) + << BSON_ARRAY(BSON_ARRAY("FOO" + << "BAR" + << "baz") + << 20)), + Accumulator::kAddToSet); + + // These strings end up as big strings copied out of the BSON array, so the size accounts for + // the value itself, the type tag, the 4-byte size of the string, and the string itself. + auto [expectedTag, expectedVal] = + makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY("bar" + << "foo") + << 34) + << BSON_ARRAY(BSON_ARRAY("bar" + << "baz" + << "foo") + << 51)), + Accumulator::kAddToSet); + aggregateAndAssertResults( + inputValuesTag, inputValuesVal, expectedTag, expectedVal, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, + CombinePartialAggsAddToSetThrowsWhenExceedingSizeLimit) { + RAIIServerParameterControllerForTest queryKnobController("internalQueryMaxAddToSetBytes", 50); + + auto accStatement = makeAccumulationStatement("$addToSet"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto input = makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(1 << 2) << 0) + << BSON_ARRAY(BSON_ARRAY(3 << 4 << 5) << 0) + << BSON_ARRAY(BSON_ARRAY(6) << 0)), + Accumulator::kAddToSet); + + auto expected = + makeArrayAccumVal(BSON_ARRAY(BSON_ARRAY(BSON_ARRAY(1 << 2) << 18) + << BSON_ARRAY(BSON_ARRAY(1 << 2 << 3 << 4 << 5) << 45) + << BSON_ARRAY(BSON_ARRAY("unused") << -1)), + Accumulator::kAddToSet); + + ASSERT_THROWS_CODE( + aggregateAndAssertResults( + input.first, input.second, expected.first, expected.second, compiledExpr.get()), + DBException, + ErrorCodes::ExceededMemoryLimit); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsMergeObjects) { + auto accStatement = makeAccumulationStatement("$mergeObjects"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + auto inputValues = BSON_ARRAY(BSONNULL << BSONObj{} << BSON("a" << 1) << BSONNULL << "MISSING" + << BSON("a" << 2 << "b" << 3 << "c" << 4) << BSONObj{}); + auto expectedAggStates = + BSON_ARRAY(BSONObj{} << BSONObj{} << BSON("a" << 1) << BSON("a" << 1) << BSON("a" << 1) + << BSON("a" << 2 << "b" << 3 << "c" << 4) + << BSON("a" << 2 << "b" << 3 << "c" << 4)); + aggregateAndAssertResults(inputValues, expectedAggStates, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsSimpleCount) { + // $sum:1 is a simple count of the incoming documents. SERVER-65465 changed this scenario to use + // a simple summation rather than the DoubleDouble summation algorithm in more recent branches, + // but the 6.0 branch still uses DoubleDouble sum. + auto inputValues = BSON_ARRAY(5 << 8 << "MISSING" << 4); + auto [inputTag, inputVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, BSON_ARRAY(BSON_ARRAY(5) << BSON_ARRAY(8) << BSON_ARRAY(4))); + auto [expectedTag, expectedVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(5) << BSON_ARRAY(5 << 8) << BSON_ARRAY(5 << 8 << 4))); + + auto accStatement = makeAccumulationStatement(BSON("unused" << BSON("$sum" << 1))); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsDoubleDoubleSum) { + auto [inputTag, inputVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(1 << 2 << 3) << BSON_ARRAY(4 << 6) << BSON_ARRAY(1 << 1 << 1))); + auto [expectedTag, expectedVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, BSON_ARRAY(BSON_ARRAY(6) << BSON_ARRAY(16) << BSON_ARRAY(19))); + + // A field path expression is needed so that the merging expression is constructed to combine + // DoubleDouble summations rather than doing a simple sum. The actual field name "foo" is + // irrelevant because the values are fed into the merging expression by the test fixture. + auto accStatement = makeAccumulationStatement(BSON("unused" << BSON("$sum" + << "$foo"))); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsDoubleDoubleSumInfAndNan) { + auto [inputTag, inputVal] = + makePartialAggArray("aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(1 << 2 << 3) + << BSON_ARRAY(4 << std::numeric_limits<double>::infinity()) + << BSON_ARRAY(1 << 1 << 1) + << BSON_ARRAY(std::numeric_limits<double>::quiet_NaN()))); + auto [expectedTag, expectedVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(6) << BSON_ARRAY(10 << std::numeric_limits<double>::infinity()) + << BSON_ARRAY(10 << std::numeric_limits<double>::infinity()) + << BSON_ARRAY(std::numeric_limits<double>::quiet_NaN()))); + + // A field path expression is needed so that the merging expression is constructed to combine + // DoubleDouble summations rather than doing a simple sum. The actual field name "foo" is + // irrelevant because the values are fed into the merging expression by the test fixture. + auto accStatement = makeAccumulationStatement(BSON("unused" << BSON("$sum" + << "$foo"))); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsDoubleDoubleSumMixedTypes) { + auto [inputTag, inputVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(1 << 2) << BSON_ARRAY(3ll << 4ll) << BSON_ARRAY(5.5 << 6.6) + << BSON_ARRAY(Decimal128(7) << Decimal128(8)))); + auto [expectedTag, expectedVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(1 << 2) << BSON_ARRAY(1 << 2 << 3ll << 4ll) + << BSON_ARRAY(1 << 2 << 3ll << 4ll << 5.5 << 6.6) + << BSON_ARRAY(1 << 2 << 3ll << 4ll << 5.5 << 6.6 + << Decimal128(7) << Decimal128(8)))); + + // A field path expression is needed so that the merging expression is constructed to combine + // DoubleDouble summations rather than doing a simple sum. The actual field name "foo" is + // irrelevant because the values are fed into the merging expression by the test fixture. + auto accStatement = makeAccumulationStatement(BSON("unused" << BSON("$sum" + << "$foo"))); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, compiledExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsDoubleDoubleSumLargeInts) { + // Large 64-bit ints can't be represented precisely as doubles. This test demonstrates that when + // summing such large longs, the sum is returned as a long and no precision is lost. + const int64_t largeLong = std::numeric_limits<int64_t>::max() - 10; + + auto [inputTag, inputVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(largeLong << 1 << 1) << BSON_ARRAY(1ll << 1ll << 1ll))); + auto [expectedTag, expectedVal] = + makePartialAggArray("aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(largeLong + 2ll) << BSON_ARRAY(largeLong + 5ll))); + + // A field path expression is needed so that the merging expression is constructed to combine + // DoubleDouble summations rather than doing a simple sum. The actual field name "foo" is + // irrelevant because the values are fed into the merging expression by the test fixture. + auto accStatement = makeAccumulationStatement(BSON("unused" << BSON("$sum" + << "$foo"))); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + + aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, compiledExpr.get()); + + // Feed the result back into the input accessor. We finalize the resulting aggregate in order + // to make sure that the resulting sum is mathematically correct. + auto [resTag, resVal] = _aggAccessor.copyOrMoveValue(); + _inputAccessor.reset(true, resTag, resVal); + auto finalizeExpr = stage_builder::makeFunction("doubleDoubleSumFinalize", + stage_builder::makeVariable(_inputSlotId)); + auto finalizeCode = compileExpression(*finalizeExpr); + auto [finalizedTag, finalizedRes] = runCompiledExpression(finalizeCode.get()); + ASSERT_EQ(finalizedTag, sbe::value::TypeTags::NumberInt64); + ASSERT_EQ(sbe::value::bitcastTo<int64_t>(finalizedRes), largeLong + 5ll); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsAvg) { + auto accStatement = makeAccumulationStatement("$avg"_sd); + + // We expect $avg to result in two separate agg expressions: one for computing the sum and the + // other for computing the count. Both agg expressions read from the same input slot. + auto exprs = stage_builder::buildCombinePartialAggregates( + accStatement, {_inputSlotId, _inputSlotId}, boost::none, _frameIdGenerator); + ASSERT_EQ(exprs.size(), 2u); + + // Compile the first expression and make sure it can combine DoubleDouble summations as + // expected. + auto [inputTag, inputVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(1 << 2) << BSON_ARRAY(3ll << 4ll) << BSON_ARRAY(5.5 << 6.6) + << BSON_ARRAY(Decimal128(7) << Decimal128(8)))); + auto [expectedTag, expectedVal] = makePartialAggArray( + "aggDoubleDoubleSum"_sd, + BSON_ARRAY(BSON_ARRAY(1 << 2) << BSON_ARRAY(1 << 2 << 3ll << 4ll) + << BSON_ARRAY(1 << 2 << 3ll << 4ll << 5.5 << 6.6) + << BSON_ARRAY(1 << 2 << 3ll << 4ll << 5.5 << 6.6 + << Decimal128(7) << Decimal128(8)))); + auto doubleDoubleSumExpr = compileAggExpression(*exprs[0], &_aggAccessor); + aggregateAndAssertResults( + inputTag, inputVal, expectedTag, expectedVal, doubleDoubleSumExpr.get()); + + // Now compile the second expression and make sure it computes a simple sum. + auto simpleSumExpr = compileAggExpression(*exprs[1], &_aggAccessor); + + auto inputValues = BSON_ARRAY(5 << 8 << 0 << 4); + auto expectedAggStates = BSON_ARRAY(5 << 13 << 13 << 17); + aggregateAndAssertResults(inputValues, expectedAggStates, simpleSumExpr.get()); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsStdDevPop) { + auto [inputTag, inputVal] = makePartialAggArray( + "aggStdDev"_sd, + BSON_ARRAY(BSON_ARRAY(5 << 10) + << BSON_ARRAY(6 << 8) << BSON_ARRAY("MISSING") << BSON_ARRAY(1 << 9 << 10))); + auto [expectedTag, expectedVal] = makePartialAggArray( + "aggStdDev"_sd, + BSON_ARRAY(BSON_ARRAY(5 << 10) + << BSON_ARRAY(5 << 10 << 6 << 8) << BSON_ARRAY(5 << 10 << 6 << 8) + << BSON_ARRAY(5 << 10 << 6 << 8 << 1 << 9 << 10))); + + auto accStatement = makeAccumulationStatement("$stdDevPop"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, compiledExpr.get()); + + // Feed the result back into the input accessor. + auto [resTag, resVal] = _aggAccessor.copyOrMoveValue(); + _inputAccessor.reset(true, resTag, resVal); + auto finalizeExpr = + stage_builder::makeFunction("stdDevPopFinalize", stage_builder::makeVariable(_inputSlotId)); + auto finalizeCode = compileExpression(*finalizeExpr); + auto [finalizedTag, finalizedRes] = runCompiledExpression(finalizeCode.get()); + ASSERT_EQ(finalizedTag, sbe::value::TypeTags::NumberDouble); + ASSERT_APPROX_EQUAL(sbe::value::bitcastTo<double>(finalizedRes), 3.0237, 0.0001); +} + +TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsStdDevSamp) { + auto [inputTag, inputVal] = makePartialAggArray( + "aggStdDev"_sd, + BSON_ARRAY(BSON_ARRAY(5 << 10) + << BSON_ARRAY(6 << 8) << BSON_ARRAY("MISSING") << BSON_ARRAY(1 << 9 << 10))); + auto [expectedTag, expectedVal] = makePartialAggArray( + "aggStdDev"_sd, + BSON_ARRAY(BSON_ARRAY(5 << 10) + << BSON_ARRAY(5 << 10 << 6 << 8) << BSON_ARRAY(5 << 10 << 6 << 8) + << BSON_ARRAY(5 << 10 << 6 << 8 << 1 << 9 << 10))); + + auto accStatement = makeAccumulationStatement("$stdDevSamp"_sd); + auto compiledExpr = compileSingleInputNoCollator(accStatement); + aggregateAndAssertResults(inputTag, inputVal, expectedTag, expectedVal, compiledExpr.get()); + + // Feed the result back into the input accessor. + auto [resTag, resVal] = _aggAccessor.copyOrMoveValue(); + _inputAccessor.reset(true, resTag, resVal); + auto finalizeExpr = stage_builder::makeFunction("stdDevSampFinalize", + stage_builder::makeVariable(_inputSlotId)); + auto finalizeCode = compileExpression(*finalizeExpr); + auto [finalizedTag, finalizedRes] = runCompiledExpression(finalizeCode.get()); + ASSERT_EQ(finalizedTag, sbe::value::TypeTags::NumberDouble); + ASSERT_APPROX_EQUAL(sbe::value::bitcastTo<double>(finalizedRes), 3.2660, 0.0001); +} + } // namespace mongo diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp index 9fd4aba048e..c21a1cc6b78 100644 --- a/src/mongo/db/query/sbe_stage_builder_expression.cpp +++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp @@ -27,16 +27,12 @@ * it in the license file. */ -#include "mongo/platform/basic.h" - #include "mongo/db/query/sbe_stage_builder_expression.h" -#include "mongo/db/query/util/make_data_structure.h" #include "mongo/base/string_data.h" #include "mongo/db/exec/sbe/stages/branch.h" #include "mongo/db/exec/sbe/stages/co_scan.h" #include "mongo/db/exec/sbe/stages/filter.h" -#include "mongo/db/exec/sbe/stages/hash_agg.h" #include "mongo/db/exec/sbe/stages/limit_skip.h" #include "mongo/db/exec/sbe/stages/loop_join.h" #include "mongo/db/exec/sbe/stages/project.h" @@ -52,6 +48,7 @@ #include "mongo/db/query/projection_parser.h" #include "mongo/db/query/sbe_stage_builder.h" #include "mongo/db/query/sbe_stage_builder_eval_frame.h" +#include "mongo/db/query/util/make_data_structure.h" #include "mongo/util/str.h" #include <absl/container/flat_hash_map.h> @@ -735,6 +732,9 @@ struct DoubleBound { static DoubleBound plusInfinity() { return DoubleBound(std::numeric_limits<double>::infinity(), false); } + static DoubleBound plusInfinityInclusive() { + return DoubleBound(std::numeric_limits<double>::infinity(), true); + } std::string printLowerBound() const { return str::stream() << (inclusive ? "[" : "(") << bound; } @@ -1181,103 +1181,46 @@ public: return; } - sbe::EExpression::Vector nullChecks; - std::vector<EvalStage> unionBranches; - std::vector<sbe::value::SlotVector> unionInputSlots; - sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projections; - - nullChecks.reserve(numChildren); - unionBranches.reserve(numChildren); - unionInputSlots.reserve(numChildren); - for (size_t idx = 0; idx < numChildren; ++idx) { - auto outputSlot = _context->state.slotId(); - projections.emplace(outputSlot, _context->popExpr()); - unionBranches.emplace_back( - EvalStage{makeLimitCoScanTree(_context->planNodeId), sbe::makeSV()}); - unionInputSlots.emplace_back(sbe::makeSV(outputSlot)); - nullChecks.emplace_back(generateNullOrMissing(outputSlot)); + auto binds = sbe::makeEs(); + for (size_t i = 0; i < numChildren; ++i) { + binds.emplace_back(_context->popExpr()); } + std::reverse(binds.begin(), binds.end()); - // Build a project to capture our child expressions. - std::reverse(std::begin(unionInputSlots), std::end(unionInputSlots)); - auto project = makeProject( - _context->extractCurrentEvalStage(), std::move(projections), _context->planNodeId); + auto frameId = _context->state.frameId(); + auto args = sbe::makeEs(); - // Build a union stage to consolidate array input branches into a stream. - auto unionOutputSlot = _context->state.slotId(); - auto unionStage = makeUnion(std::move(unionBranches), - std::move(unionInputSlots), - sbe::makeSV(unionOutputSlot), - _context->planNodeId); + std::unique_ptr<sbe::EExpression> checkArgsForNull; + for (size_t i = 0; i < numChildren; ++i) { + sbe::EVariable argRef(frameId, i); + args.emplace_back(argRef.clone()); - auto collatorSlot = _context->state.data->env->getSlotIfExists("collator"_sd); + checkArgsForNull = checkArgsForNull ? makeBinaryOp(sbe::EPrimBinary::logicOr, + std::move(checkArgsForNull), + generateNullOrMissing(argRef)) + : generateNullOrMissing(argRef); + } - // Build a filter that will throw an 'EFail' if any element coming from the union is NOT - // an array. - auto filter = makeFilter<false, false>( - std::move(unionStage), - makeBinaryOp(sbe::EPrimBinary::logicOr, - makeFunction("isArray", makeVariable(unionOutputSlot)), - sbe::makeE<sbe::EFail>(ErrorCodes::Error{5153400}, - "$concatArrays only supports arrays")), - _context->planNodeId); + auto nullOrFailExpr = + sbe::makeE<sbe::EIf>(std::move(checkArgsForNull), + makeConstant(sbe::value::TypeTags::Null, 0), + sbe::makeE<sbe::EFail>(ErrorCodes::Error{5153400}, + "$concatArrays only supports arrays")); - // Build subtree to handle nulls. If an input is null, return null. Otherwise, unwind the - // input and concatenate it into an array using addToArray. - auto unwindEvalStage = - makeUnwind(std::move(filter), _context->state.slotIdGenerator, _context->planNodeId); - auto unwindSlot = unwindEvalStage.outSlots.front(); - - // Create a group stage to append all streamed elements into one array. This is the final - // output when the input consists entirely of arrays. - auto finalAddToArrayExpr = makeFunction("addToArray", makeVariable(unwindSlot)); - auto finalGroupSlot = _context->state.slotId(); - auto finalGroupStage = - makeHashAgg(std::move(unwindEvalStage), - sbe::makeSV(), - sbe::makeEM(finalGroupSlot, std::move(finalAddToArrayExpr)), - collatorSlot, - _context->state.allowDiskUse, - _context->planNodeId); - - // Returns true if any of our input expressions return null. - using iter_t = sbe::EExpression::Vector::iterator; - auto checkPartsForNull = std::accumulate( - std::move_iterator<iter_t>(nullChecks.begin() + 1), - std::move_iterator<iter_t>(nullChecks.end()), - std::move(nullChecks.front()), - [](auto&& acc, auto&& b) { - return makeBinaryOp(sbe::EPrimBinary::logicOr, std::move(acc), std::move(b)); - }); - - // Create a branch stage to select between the branch that produces one null if any elements - // in the original input were null or missing, or otherwise select the branch that unwinds - // and concatenates elements into the output array. - auto [nullSlot, nullStage] = [&] { - auto outputSlot = _context->state.slotId(); - auto nullEvalStage = - makeProject({makeLimitCoScanTree(_context->planNodeId), sbe::makeSV()}, - _context->planNodeId, - outputSlot, - makeConstant(sbe::value::TypeTags::Null, 0)); - return std::make_pair(outputSlot, std::move(nullEvalStage)); - }(); + auto resultExpr = makeLocalBind( + _context->state.frameIdGenerator, + [&](sbe::EVariable concatArraysRef) { + // We optimize for the case where all of the args are arrays. If concatArrays() + // returns Nothing, then we deal with checking if any of the args are null and + // either returning null or raising an error. + return sbe::makeE<sbe::EIf>(makeFunction("exists", concatArraysRef.clone()), + concatArraysRef.clone(), + std::move(nullOrFailExpr)); + }, + sbe::makeE<sbe::EFunction>("concatArrays"_sd, std::move(args))); - auto branchSlot = _context->state.slotId(); - auto branchNullEvalStage = makeBranch(std::move(nullStage), - std::move(finalGroupStage), - std::move(checkPartsForNull), - sbe::makeSV(nullSlot), - sbe::makeSV(finalGroupSlot), - sbe::makeSV(branchSlot), - _context->planNodeId); - - // Create nlj to connect outer project with inner branch that handles null input. - _context->pushExpr(branchSlot, - makeLoopJoin(std::move(project), - std::move(branchNullEvalStage), - _context->planNodeId, - _context->getLexicalEnvironment())); + _context->pushExpr( + sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(resultExpr))); } void visit(const ExpressionCond* expr) final { visitConditionalExpression(expr); @@ -2465,13 +2408,7 @@ public: exprs[--i] = makeConstant(rit->first); } - auto fieldSlot{_context->state.slotIdGenerator->generate()}; - auto stage = makeProject(_context->extractCurrentEvalStage(), - _context->planNodeId, - fieldSlot, - sbe::makeE<sbe::EFunction>("newObj"_sd, std::move(exprs))); - - _context->pushExpr(fieldSlot, std::move(stage)); + _context->pushExpr(sbe::makeE<sbe::EFunction>("newObj"_sd, std::move(exprs))); } void visit(const ExpressionOr* expr) final { visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicOr); @@ -2910,7 +2847,7 @@ public: } void visit(const ExpressionHyperbolicArcCosine* expr) final { generateTrigonometricExpressionWithBounds( - "acosh", DoubleBound(1.0, true), DoubleBound::plusInfinity()); + "acosh", DoubleBound(1.0, true), DoubleBound::plusInfinityInclusive()); } void visit(const ExpressionHyperbolicArcSine* expr) final { generateTrigonometricExpression("asinh"); @@ -3353,7 +3290,8 @@ private: */ void generateTrigonometricExpressionBinary(StringData exprName) { _context->ensureArity(2); - + auto x = _context->popExpr(); + auto y = _context->popExpr(); auto genericTrignomentricExpr = makeLocalBind( _context->state.frameIdGenerator, [&](sbe::EVariable lhs, sbe::EVariable rhs) { @@ -3372,8 +3310,8 @@ private: str::stream() << "$" << exprName << " supports only numeric types")); }, - _context->popExpr(), - _context->popExpr()); + std::move(y), + std::move(x)); _context->pushExpr(std::move(genericTrignomentricExpr)); } @@ -3414,13 +3352,17 @@ private: str::stream() << "$" << exprName.toString() << " supports only numeric types"), sbe::makeE<sbe::EIf>( - std::move(checkBounds), - makeFunction(exprName.toString(), inputRef.clone()), - sbe::makeE<sbe::EFail>(ErrorCodes::Error{4995503}, - str::stream() << "Cannot apply $" << exprName.toString() - << ", value must be in " - << lowerBound.printLowerBound() << ", " - << upperBound.printUpperBound())))); + // return NaN when NaN is the input. + generateNaNCheck(inputRef), + inputRef.clone(), + sbe::makeE<sbe::EIf>( + std::move(checkBounds), + makeFunction(exprName.toString(), inputRef.clone()), + sbe::makeE<sbe::EFail>( + ErrorCodes::Error{4995503}, + str::stream() << "Cannot apply $" << exprName.toString() + << ", value must be in " << lowerBound.printLowerBound() + << ", " << upperBound.printUpperBound()))))); _context->pushExpr(sbe::makeE<sbe::ELocalBind>( frameId, std::move(binds), std::move(genericTrignomentricExpr))); diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.cpp b/src/mongo/db/query/sbe_stage_builder_helpers.cpp index c36947a23ad..e04233e11c5 100644 --- a/src/mongo/db/query/sbe_stage_builder_helpers.cpp +++ b/src/mongo/db/query/sbe_stage_builder_helpers.cpp @@ -470,14 +470,20 @@ EvalStage makeUnion(std::vector<EvalStage> inputStages, EvalStage makeHashAgg(EvalStage stage, sbe::value::SlotVector gbs, - sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> aggs, + sbe::SlotExprPairVector aggs, boost::optional<sbe::value::SlotId> collatorSlot, bool allowDiskUse, + sbe::SlotExprPairVector mergingExprs, PlanNodeId planNodeId) { stage.outSlots = gbs; for (auto& [slot, _] : aggs) { stage.outSlots.push_back(slot); } + + // In debug builds, we artificially force frequent spilling. This makes sure that our tests + // exercise the spilling algorithm and the associated logic for merging partial aggregates which + // otherwise would require large data sizes to exercise. + const bool forceIncreasedSpilling = kDebugBuild && allowDiskUse; stage.stage = sbe::makeS<sbe::HashAggStage>(std::move(stage.stage), std::move(gbs), std::move(aggs), @@ -485,7 +491,9 @@ EvalStage makeHashAgg(EvalStage stage, true /* optimized close */, collatorSlot, allowDiskUse, - planNodeId); + std::move(mergingExprs), + planNodeId, + forceIncreasedSpilling); return stage; } @@ -868,12 +876,13 @@ void indexKeyCorruptionCheckCallback(OperationContext* opCtx, * or that the index keys are still part of the underlying index. */ bool indexKeyConsistencyCheckCallback(OperationContext* opCtx, - StringMap<const IndexAccessMethod*> iamTable, + StringMap<const IndexCatalogEntry*>& entryMap, sbe::value::SlotAccessor* snapshotIdAccessor, sbe::value::SlotAccessor* indexIdAccessor, sbe::value::SlotAccessor* indexKeyAccessor, const CollectionPtr& collection, const Record& nextRecord) { + // The index consistency check is only performed when 'snapshotIdAccessor' is set. if (snapshotIdAccessor) { auto currentSnapshotId = opCtx->recoveryUnit()->getSnapshotId(); auto [snapshotIdTag, snapshotIdVal] = snapshotIdAccessor->getViewOfValue(); @@ -904,14 +913,29 @@ bool indexKeyConsistencyCheckCallback(OperationContext* opCtx, auto indexId = sbe::value::getStringView(indexIdTag, indexIdVal); tassert(5290712, "KeyString does not exist", keyString); - auto it = iamTable.find(indexId); - tassert(5290713, - str::stream() << "IndexAccessMethod not found for index " << indexId, - it != iamTable.end()); + auto it = entryMap.find(indexId); + + // If 'entryMap' doesn't contain an entry for 'indexId', create one. + if (it == entryMap.end()) { + auto indexCatalog = collection->getIndexCatalog(); + auto indexDesc = indexCatalog->findIndexByName(opCtx, indexId); + auto entry = indexDesc ? indexDesc->getEntry() : nullptr; + + // Throw an error if we can't get the IndexDescriptor or the IndexCatalogEntry + // (or if the index is dropped). + uassert(ErrorCodes::QueryPlanKilled, + str::stream() << "query plan killed :: index dropped: " << indexId, + indexDesc && entry && !entry->isDropped()); - auto iam = it->second->asSortedData(); + auto [newIt, _] = entryMap.emplace(indexId, entry); + + it = newIt; + } + + auto entry = it->second; + auto iam = entry->accessMethod()->asSortedData(); tassert(5290709, - str::stream() << "Expected to find SortedDataIndexAccessMethod for index " + str::stream() << "Expected to find SortedDataIndexAccessMethod for index: " << indexId, iam); @@ -939,6 +963,7 @@ bool indexKeyConsistencyCheckCallback(OperationContext* opCtx, return keys->count(*keyString); } } + return true; } @@ -950,7 +975,6 @@ makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage, sbe::value::SlotId indexKeySlot, sbe::value::SlotId indexKeyPatternSlot, const CollectionPtr& collToFetch, - StringMap<const IndexAccessMethod*> iamMap, PlanNodeId planNodeId, sbe::value::SlotVector slotsToForward, sbe::value::SlotIdGenerator& slotIdGenerator) { @@ -962,10 +986,7 @@ makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage, auto resultSlot = slotIdGenerator.generate(); auto recordIdSlot = slotIdGenerator.generate(); - using namespace std::placeholders; - sbe::ScanCallbacks callbacks( - indexKeyCorruptionCheckCallback, - std::bind(indexKeyConsistencyCheckCallback, _1, std::move(iamMap), _2, _3, _4, _5, _6)); + sbe::ScanCallbacks callbacks(indexKeyCorruptionCheckCallback, indexKeyConsistencyCheckCallback); // Scan the collection in the range [seekKeySlot, Inf). auto scanStage = sbe::makeS<sbe::ScanStage>(collToFetch->uuid(), diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.h b/src/mongo/db/query/sbe_stage_builder_helpers.h index 05cf73896e0..e2e203bdd55 100644 --- a/src/mongo/db/query/sbe_stage_builder_helpers.h +++ b/src/mongo/db/query/sbe_stage_builder_helpers.h @@ -36,6 +36,7 @@ #include "mongo/db/exec/sbe/expressions/expression.h" #include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/hash_agg.h" #include "mongo/db/exec/sbe/stages/makeobj.h" #include "mongo/db/exec/sbe/stages/project.h" #include "mongo/db/pipeline/expression.h" @@ -416,9 +417,10 @@ EvalStage makeUnion(std::vector<EvalStage> inputStages, EvalStage makeHashAgg(EvalStage stage, sbe::value::SlotVector gbs, - sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> aggs, + sbe::SlotExprPairVector aggs, boost::optional<sbe::value::SlotId> collatorSlot, bool allowDiskUse, + sbe::SlotExprPairVector mergingExprs, PlanNodeId planNodeId); EvalStage makeMkBsonObj(EvalStage stage, @@ -537,7 +539,6 @@ makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage, sbe::value::SlotId indexKeySlot, sbe::value::SlotId indexKeyPatternSlot, const CollectionPtr& collToFetch, - StringMap<const IndexAccessMethod*> iamMap, PlanNodeId planNodeId, sbe::value::SlotVector slotsToForward, sbe::value::SlotIdGenerator& slotIdGenerator); diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp index 7112e8ad56b..cc63eee572a 100644 --- a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp +++ b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp @@ -283,10 +283,6 @@ generateOptimizedMultiIntervalIndexScan(StageBuilderState& state, makeFunction("getField"_sd, makeVariable(unwindSlot), makeConstant("l"_sd))); projects.emplace(highKeySlot, makeFunction("getField"_sd, makeVariable(unwindSlot), makeConstant("h"_sd))); - if (indexIdSlot) { - // Construct a copy of 'indexName' to project for use in the index consistency check. - projects.emplace(*indexIdSlot, makeConstant(indexName)); - } if (indexKeyPatternSlot) { auto [bsonObjTag, bsonObjVal] = @@ -300,20 +296,13 @@ generateOptimizedMultiIntervalIndexScan(StageBuilderState& state, auto project = sbe::makeS<sbe::ProjectStage>(std::move(unwind), std::move(projects), planNodeId); - // Whereas 'snapshotIdSlot' is used by the caller to inspect the snapshot id of the latest index - // key, 'indexSnapshotSlot' is updated by the IndexScan below during yield to obtain the latest - // snapshot id. - boost::optional<sbe::value::SlotId> indexSnapshotSlot; - if (snapshotIdSlot) { - indexSnapshotSlot = slotIdGenerator->generate(); - } - auto stage = sbe::makeS<sbe::IndexScanStage>(collection->uuid(), indexName, forward, recordSlot, recordIdSlot, - indexSnapshotSlot, + snapshotIdSlot, + indexIdSlot, indexKeysToInclude, std::move(indexKeySlots), lowKeySlot, @@ -321,19 +310,7 @@ generateOptimizedMultiIntervalIndexScan(StageBuilderState& state, yieldPolicy, planNodeId); - // Add a project on top of the index scan to remember the snapshotId of the most recent index - // key returned by the IndexScan above. Otherwise, the index key's snapshot id would be - // overwritten during yield. - if (snapshotIdSlot) { - stage = sbe::makeProjectStage( - std::move(stage), planNodeId, *snapshotIdSlot, makeVariable(*indexSnapshotSlot)); - } - auto outerSv = sbe::makeSV(); - if (indexIdSlot) { - outerSv.push_back(*indexIdSlot); - } - if (indexKeyPatternSlot) { outerSv.push_back(*indexKeyPatternSlot); } @@ -425,10 +402,6 @@ makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection, // contain a value from the stack spool. See below for details. sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projects; projects.emplace(lowKeySlot, makeVariable(seekKeySlot)); - if (indexIdSlot) { - // Construct a copy of 'indexName' to project for use in the index consistency check. - projects.emplace(*indexIdSlot, makeConstant(indexName)); - } if (indexKeyPatternSlot) { auto [bsonObjTag, bsonObjVal] = sbe::value::copyValue( @@ -449,6 +422,7 @@ makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection, resultSlot, recordIdSlot, snapshotIdSlot, + indexIdSlot, indexKeysToInclude, std::move(savedIndexKeySlots), lowKeySlot, @@ -458,10 +432,6 @@ makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection, // Get the low key from the outer side and feed it to the inner side (ixscan). sbe::value::SlotVector outerSv = sbe::makeSV(); - if (indexIdSlot) { - outerSv.push_back(*indexIdSlot); - } - if (indexKeyPatternSlot) { outerSv.push_back(*indexKeyPatternSlot); } @@ -474,19 +444,11 @@ makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection, planNodeId); sbe::value::SlotVector correlatedSv = sbe::makeSV(seekKeySlot); - if (indexIdSlot) { - correlatedSv.push_back(*indexIdSlot); - } - if (indexKeyPatternSlot) { correlatedSv.push_back(*indexKeyPatternSlot); } auto spoolValsSV = sbe::makeSV(seekKeySlot); - if (indexIdSlot) { - spoolValsSV.push_back(*indexIdSlot); - } - if (indexKeyPatternSlot) { spoolValsSV.push_back(*indexKeyPatternSlot); } @@ -844,11 +806,6 @@ generateSingleIntervalIndexScan(StageBuilderState& state, auto lowKeySlot = makeKeySlot(std::move(lowKey)); auto highKeySlot = makeKeySlot(std::move(highKey)); - if (indexIdSlot) { - // Construct a copy of 'indexName' to project for use in the index consistency check. - projects.emplace(*indexIdSlot, makeConstant(indexName)); - } - if (indexKeyPatternSlot) { auto [bsonObjTag, bsonObjVal] = sbe::value::copyValue(sbe::value::TypeTags::bsonObject, @@ -883,14 +840,6 @@ generateSingleIntervalIndexScan(StageBuilderState& state, planNodeId); }(); - // Whereas 'snapshotIdSlot' is used by the caller to inspect the snapshot id of the latest index - // key, 'indexSnapshotSlot' is updated by the IndexScan below during yield to obtain the latest - // snapshot id. - boost::optional<sbe::value::SlotId> indexSnapshotSlot; - if (snapshotIdSlot) { - indexSnapshotSlot = slotIdGenerator->generate(); - } - // Scan the index in the range {'lowKeySlot', 'highKeySlot'} (subject to inclusive or // exclusive boundaries), and produce a single field recordIdSlot that can be used to // position into the collection. @@ -899,7 +848,8 @@ generateSingleIntervalIndexScan(StageBuilderState& state, forward, recordSlot, recordIdSlot, - indexSnapshotSlot, + snapshotIdSlot, + indexIdSlot, indexKeysToInclude, std::move(indexKeySlots), lowKeySlot, @@ -907,19 +857,7 @@ generateSingleIntervalIndexScan(StageBuilderState& state, yieldPolicy, planNodeId); - // Add a project on top of the index scan to remember the snapshotId of the most recent index - // key returned by the IndexScan above. Otherwise, the index key's snapshot id would be - // overwritten during yield. - if (snapshotIdSlot) { - stage = sbe::makeProjectStage( - std::move(stage), planNodeId, *snapshotIdSlot, makeVariable(*indexSnapshotSlot)); - } - auto outerSv = sbe::makeSV(); - if (indexIdSlot) { - outerSv.push_back(*indexIdSlot); - } - if (indexKeyPatternSlot) { outerSv.push_back(*indexKeyPatternSlot); } @@ -942,7 +880,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( const IndexScanNode* ixn, const sbe::IndexKeysInclusionSet& originalIndexKeyBitset, PlanYieldPolicy* yieldPolicy, - StringMap<const IndexAccessMethod*>* iamMap, + bool doIndexConsistencyCheck, bool needsCorruptionCheck) { auto indexName = ixn->index.identifier.catalogName; @@ -980,15 +918,12 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( auto indexKeySlots = state.slotIdGenerator->generateMultiple(indexKeyBitset.count()); sbe::value::SlotVector relevantSlots; - // Generate the relevant slots and add the access method corresponding to 'indexName' to - // 'iamMap' if a parent stage needs to execute a consistency check. + // Generate the relevant slots. boost::optional<sbe::value::SlotId> snapshotIdSlot; boost::optional<sbe::value::SlotId> indexIdSlot; boost::optional<sbe::value::SlotId> indexKeySlot; - if (iamMap) { - iamMap->insert({indexName, accessMethod}); - + if (doIndexConsistencyCheck) { snapshotIdSlot = state.slotId(); outputs.set(PlanStageSlots::kSnapshotId, *snapshotIdSlot); relevantSlots.push_back(*snapshotIdSlot); @@ -1188,7 +1123,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScanWith const IndexScanNode* ixn, const sbe::IndexKeysInclusionSet& originalIndexKeyBitset, PlanYieldPolicy* yieldPolicy, - StringMap<const IndexAccessMethod*>* iamMap, + bool doIndexConsistencyCheck, bool needsCorruptionCheck) { const bool forward = ixn->direction == 1; auto indexName = ixn->index.identifier.catalogName; @@ -1202,11 +1137,6 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScanWith // Find the IndexAccessMethod which corresponds to the 'indexName'. auto accessMethod = descriptor->getEntry()->accessMethod()->asSortedData(); - // Add the access method corresponding to 'indexName' to the 'iamMap' if a parent stage needs to - // execute a consistency check. - if (iamMap) { - iamMap->insert({indexName, accessMethod}); - } PlanStageSlots outputs; sbe::value::SlotVector relevantSlots; std::unique_ptr<sbe::PlanStage> stage; @@ -1251,9 +1181,9 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScanWith nullptr, indexKeyBitset, outputIndexKeySlots, - makeSlot(iamMap, PlanStageSlots::kSnapshotId), - makeSlot(iamMap, PlanStageSlots::kIndexId), - makeSlot(iamMap, PlanStageSlots::kIndexKey), + makeSlot(doIndexConsistencyCheck, PlanStageSlots::kSnapshotId), + makeSlot(doIndexConsistencyCheck, PlanStageSlots::kIndexId), + makeSlot(doIndexConsistencyCheck, PlanStageSlots::kIndexKey), makeSlot(needsCorruptionCheck, PlanStageSlots::kIndexKeyPattern), yieldPolicy, ixn->nodeId()); @@ -1291,11 +1221,11 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScanWith }; auto [genericIndexScanSnapshotIdSlot, optimizedIndexScanSnapshotIdSlot] = - makeSlotsForThenElseBranches(iamMap, PlanStageSlots::kSnapshotId); + makeSlotsForThenElseBranches(doIndexConsistencyCheck, PlanStageSlots::kSnapshotId); auto [genericIndexScanIndexIdSlot, optimizedIndexScanIndexIdSlot] = - makeSlotsForThenElseBranches(iamMap, PlanStageSlots::kIndexId); + makeSlotsForThenElseBranches(doIndexConsistencyCheck, PlanStageSlots::kIndexId); auto [genericIndexScanIndexKeySlot, optimizedIndexScanIndexKeySlot] = - makeSlotsForThenElseBranches(iamMap, PlanStageSlots::kIndexKey); + makeSlotsForThenElseBranches(doIndexConsistencyCheck, PlanStageSlots::kIndexKey); // Generate a slot for an index key pattern if a parent stage needs to execute a // corruption check. diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.h b/src/mongo/db/query/sbe_stage_builder_index_scan.h index 340a03051eb..08e239c018a 100644 --- a/src/mongo/db/query/sbe_stage_builder_index_scan.h +++ b/src/mongo/db/query/sbe_stage_builder_index_scan.h @@ -63,7 +63,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( const IndexScanNode* ixn, const sbe::IndexKeysInclusionSet& indexKeyBitset, PlanYieldPolicy* yieldPolicy, - StringMap<const IndexAccessMethod*>* iamMap, + bool doIndexConsistencyCheck, bool needsCorruptionCheck); /** @@ -164,6 +164,6 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScanWith const IndexScanNode* ixn, const sbe::IndexKeysInclusionSet& indexKeyBitset, PlanYieldPolicy* yieldPolicy, - StringMap<const IndexAccessMethod*>* iamMap, + bool doIndexConsistencyCheck, bool needsCorruptionCheck); } // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_lookup.cpp b/src/mongo/db/query/sbe_stage_builder_lookup.cpp index 4b62228edde..6037fb4eef7 100644 --- a/src/mongo/db/query/sbe_stage_builder_lookup.cpp +++ b/src/mongo/db/query/sbe_stage_builder_lookup.cpp @@ -342,12 +342,15 @@ std::pair<SlotId /* keyValuesSetSlot */, std::unique_ptr<sbe::PlanStage>> buildK // Re-pack the individual key values into a set. We don't cap "addToSet" here because its size // is bounded by the size of the record. SlotId keyValuesSetSlot = slotIdGenerator.generate(); + SlotId spillSlot = slotIdGenerator.generate(); EvalStage packedKeyValuesStage = makeHashAgg( EvalStage{std::move(keyValuesStage), SlotVector{}}, makeSV(), /* groupBy slots - "none" means creating a single group */ - makeEM(keyValuesSetSlot, makeFunction("addToSet"_sd, makeVariable(keyValueSlot))), + makeSlotExprPairVec(keyValuesSetSlot, + makeFunction("addToSet"_sd, makeVariable(keyValueSlot))), boost::none /* we group _all_ key values into a single set, so collator is irrelevant */, allowDiskUse, + makeSlotExprPairVec(spillSlot, makeFunction("aggSetUnion"_sd, makeVariable(spillSlot))), nodeId); // The set in 'keyValuesSetSlot' might end up empty if the localField contained only missing and @@ -403,15 +406,20 @@ std::pair<SlotId /* resultSlot */, std::unique_ptr<sbe::PlanStage>> buildForeign // are no matches, return an empty array. const int sizeCap = internalLookupStageIntermediateDocumentMaxSizeBytes.load(); SlotId accumulatorSlot = slotIdGenerator.generate(); + SlotId spillSlot = slotIdGenerator.generate(); innerBranch = makeHashAgg( std::move(innerBranch), makeSV(), /* groupBy slots */ - makeEM(accumulatorSlot, - makeFunction("addToArrayCapped"_sd, - makeVariable(foreignRecordSlot), - makeConstant(TypeTags::NumberInt32, sizeCap))), + makeSlotExprPairVec(accumulatorSlot, + makeFunction("addToArrayCapped"_sd, + makeVariable(foreignRecordSlot), + makeConstant(TypeTags::NumberInt32, sizeCap))), {} /* collatorSlot, no collation here because we want to return all matches "as is" */, allowDiskUse, + makeSlotExprPairVec(spillSlot, + makeFunction("aggConcatArraysCapped", + makeVariable(spillSlot), + makeConstant(TypeTags::NumberInt32, sizeCap))), nodeId); // 'accumulatorSlot' is either Nothing or contains an array of size two, where the front element @@ -610,7 +618,6 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( const FieldPath& foreignFieldName, const CollectionPtr& foreignColl, const IndexEntry& index, - StringMap<const IndexAccessMethod*>& iamMap, PlanYieldPolicySBE* yieldPolicy, boost::optional<SlotId> collatorSlot, const PlanNodeId nodeId, @@ -629,7 +636,6 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( foreignColl->getIndexCatalog()->getEntry(indexDescriptor)->accessMethod()->asSortedData(); const auto indexVersion = indexAccessMethod->getSortedDataInterface()->getKeyStringVersion(); const auto indexOrdering = indexAccessMethod->getSortedDataInterface()->getOrdering(); - iamMap.insert({indexName, indexAccessMethod}); // Build the outer branch that produces the correlated local key slot. auto [localKeysSetSlot, localKeysSetStage] = buildKeySet(JoinSide::Local, @@ -741,11 +747,10 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( // Calculate the low key and high key of each individual local field. They are stored in // 'lowKeySlot' and 'highKeySlot', respectively. These two slots will be made available in - // the loop join stage to perform index seek. We also set 'indexIdSlot' and - // 'indexKeyPatternSlot' constants for the seek stage later to perform consistency check. + // the loop join stage to perform index seek. We also set the 'indexKeyPatternSlot' constant + // for the seek stage later to perform consistency check. auto lowKeySlot = slotIdGenerator.generate(); auto highKeySlot = slotIdGenerator.generate(); - auto indexIdSlot = slotIdGenerator.generate(); auto indexKeyPatternSlot = slotIdGenerator.generate(); auto [_, indexKeyPatternValue] = copyValue(TypeTags::bsonObject, bitcastFrom<const char*>(index.keyPattern.objdata())); @@ -772,8 +777,6 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( makeNewKeyStringCall(KeyString::Discriminator::kExclusiveBefore), highKeySlot, makeNewKeyStringCall(KeyString::Discriminator::kExclusiveAfter), - indexIdSlot, - makeConstant(indexName), indexKeyPatternSlot, makeConstant(value::TypeTags::bsonObject, indexKeyPatternValue)); @@ -794,12 +797,14 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( auto foreignRecordIdSlot = slotIdGenerator.generate(); auto indexKeySlot = slotIdGenerator.generate(); auto snapshotIdSlot = slotIdGenerator.generate(); + auto indexIdSlot = slotIdGenerator.generate(); auto ixScanStage = makeS<IndexScanStage>(foreignCollUUID, indexName, true /* forward */, indexKeySlot, foreignRecordIdSlot, snapshotIdSlot, + indexIdSlot, IndexKeysInclusionSet{} /* indexKeysToInclude */, makeSV() /* vars */, lowKeySlot, @@ -812,7 +817,7 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( auto ixScanNljStage = makeS<LoopJoinStage>(std::move(indexBoundKeyStage), std::move(ixScanStage), - makeSV(indexIdSlot, indexKeyPatternSlot) /* outerProjects */, + makeSV(indexKeyPatternSlot) /* outerProjects */, makeSV(lowKeySlot, highKeySlot) /* outerCorrelated */, nullptr /* predicate */, nodeId); @@ -832,8 +837,7 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( // Loop join the foreign record id produced by the index seek on the outer side with seek // stage on the inner side to get matched foreign documents. The foreign documents are // stored in 'foreignRecordSlot'. We also pass in 'snapshotIdSlot', 'indexIdSlot', - // 'indexKeySlot' and 'indexKeyPatternSlot' to perform index consistency check during the - // seek. + // 'indexKeySlot' and 'indexKeyPatternSlot' to perform index consistency check during the seek. auto [foreignRecordSlot, __, scanNljStage] = makeLoopJoinForFetch(std::move(ixScanNljStage), foreignRecordIdSlot, snapshotIdSlot, @@ -841,7 +845,6 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( indexKeySlot, indexKeyPatternSlot, foreignColl, - iamMap, nodeId, makeSV() /* slotsToForward */, slotIdGenerator); @@ -1088,7 +1091,6 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder eqLookupNode->joinFieldForeign, foreignColl, *eqLookupNode->idxEntry, - _data.iamMap, _yieldPolicy, collatorSlot, eqLookupNode->nodeId(), diff --git a/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp b/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp index b54872582eb..bcc426984e8 100644 --- a/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp +++ b/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp @@ -235,7 +235,7 @@ public: expectedDocuments.reserve(expectedPairs.size()); for (auto& [localDocument, matchedDocuments] : expectedPairs) { MutableDocument expectedDocument; - expectedDocument.reset(localDocument, false /* stripMetadata */); + expectedDocument.reset(localDocument, false /* bsonHasMetadata */); std::vector<mongo::Value> matchedValues{matchedDocuments.begin(), matchedDocuments.end()}; diff --git a/src/mongo/db/query/sbe_utils.cpp b/src/mongo/db/query/sbe_utils.cpp index 695b2904562..7f327848fc7 100644 --- a/src/mongo/db/query/sbe_utils.cpp +++ b/src/mongo/db/query/sbe_utils.cpp @@ -62,9 +62,11 @@ bool isQuerySbeCompatible(const CollectionPtr* collection, const bool doesNotHaveElemMatchProject = !cq->getProj() || !cq->getProj()->containsElemMatch(); + const bool isNotInnerSideOfLookup = !(expCtx && expCtx->inLookup); + return allExpressionsSupported && isNotCount && doesNotContainMetadataRequirements && isQueryNotAgainstTimeseriesCollection && isQueryNotAgainstClusteredCollection && doesNotSortOnMetaOrPathWithNumericComponents && isNotOplog && doesNotRequireMatchDetails && - doesNotHaveElemMatchProject; + doesNotHaveElemMatchProject && isNotInnerSideOfLookup; } } // namespace mongo::sbe diff --git a/src/mongo/db/query/sort_pattern.h b/src/mongo/db/query/sort_pattern.h index 9c74208ac43..b659ed0124e 100644 --- a/src/mongo/db/query/sort_pattern.h +++ b/src/mongo/db/query/sort_pattern.h @@ -147,6 +147,6 @@ private: std::vector<SortPatternPart> _sortPattern; // The set of paths on which we're sorting. - std::set<std::string> _paths; + OrderedPathSet _paths; }; } // namespace mongo diff --git a/src/mongo/db/query/wildcard_multikey_paths.cpp b/src/mongo/db/query/wildcard_multikey_paths.cpp index fb27c8b34f5..eea11ccd4b5 100644 --- a/src/mongo/db/query/wildcard_multikey_paths.cpp +++ b/src/mongo/db/query/wildcard_multikey_paths.cpp @@ -31,7 +31,7 @@ #include "mongo/db/query/wildcard_multikey_paths.h" -#include "mongo/db/concurrency/write_conflict_exception.h" +#include "mongo/db/concurrency/exception_util.h" #include "mongo/db/index/wildcard_access_method.h" #include "mongo/db/query/index_bounds_builder.h" #include "mongo/db/record_id_helpers.h" |