diff options
Diffstat (limited to 'src/mongo')
-rw-r--r-- | src/mongo/db/exec/sort_key_generator.cpp | 9 | ||||
-rw-r--r-- | src/mongo/db/exec/sort_key_generator.h | 6 | ||||
-rw-r--r-- | src/mongo/db/exec/sort_test.cpp | 13 | ||||
-rw-r--r-- | src/mongo/db/index/SConscript | 2 | ||||
-rw-r--r-- | src/mongo/db/index/sort_key_generator.cpp | 68 | ||||
-rw-r--r-- | src/mongo/db/index/sort_key_generator.h | 19 | ||||
-rw-r--r-- | src/mongo/db/index/sort_key_generator_test.cpp | 105 | ||||
-rw-r--r-- | src/mongo/db/pipeline/document_source_sort.cpp | 2 | ||||
-rw-r--r-- | src/mongo/db/pipeline/expression.h | 18 | ||||
-rw-r--r-- | src/mongo/db/query/canonical_query.cpp | 5 | ||||
-rw-r--r-- | src/mongo/db/query/canonical_query.h | 4 | ||||
-rw-r--r-- | src/mongo/db/query/get_executor.cpp | 5 | ||||
-rw-r--r-- | src/mongo/db/query/sort_pattern.cpp | 17 | ||||
-rw-r--r-- | src/mongo/db/query/sort_pattern.h | 2 | ||||
-rw-r--r-- | src/mongo/db/query/stage_builder.cpp | 3 | ||||
-rw-r--r-- | src/mongo/dbtests/query_stage_ensure_sorted.cpp | 14 | ||||
-rw-r--r-- | src/mongo/dbtests/query_stage_sort.cpp | 7 | ||||
-rw-r--r-- | src/mongo/dbtests/query_stage_sort_key_generator.cpp | 6 |
18 files changed, 166 insertions, 139 deletions
diff --git a/src/mongo/db/exec/sort_key_generator.cpp b/src/mongo/db/exec/sort_key_generator.cpp index 7b846f18bad..90fe3b48fea 100644 --- a/src/mongo/db/exec/sort_key_generator.cpp +++ b/src/mongo/db/exec/sort_key_generator.cpp @@ -49,12 +49,13 @@ namespace mongo { const char* SortKeyGeneratorStage::kStageType = "SORT_KEY_GENERATOR"; -SortKeyGeneratorStage::SortKeyGeneratorStage(OperationContext* opCtx, +SortKeyGeneratorStage::SortKeyGeneratorStage(const boost::intrusive_ptr<ExpressionContext>& pExpCtx, PlanStage* child, WorkingSet* ws, - const BSONObj& sortSpecObj, - const CollatorInterface* collator) - : PlanStage(kStageType, opCtx), _ws(ws), _sortKeyGen(sortSpecObj, collator) { + const BSONObj& sortSpecObj) + : PlanStage(kStageType, pExpCtx->opCtx), + _ws(ws), + _sortKeyGen({{sortSpecObj, pExpCtx}, pExpCtx->getCollator()}) { _children.emplace_back(child); } diff --git a/src/mongo/db/exec/sort_key_generator.h b/src/mongo/db/exec/sort_key_generator.h index 79cb5a6fe49..c7a9e2dfd4b 100644 --- a/src/mongo/db/exec/sort_key_generator.h +++ b/src/mongo/db/exec/sort_key_generator.h @@ -34,6 +34,7 @@ #include "mongo/bson/bsonobj.h" #include "mongo/db/exec/plan_stage.h" #include "mongo/db/index/sort_key_generator.h" +#include "mongo/db/pipeline/expression_context.h" #include "mongo/db/query/index_bounds.h" #include "mongo/db/query/stage_types.h" @@ -49,11 +50,10 @@ class WorkingSetMember; */ class SortKeyGeneratorStage final : public PlanStage { public: - SortKeyGeneratorStage(OperationContext* opCtx, + SortKeyGeneratorStage(const boost::intrusive_ptr<ExpressionContext>& pExpCtx, PlanStage* child, WorkingSet* ws, - const BSONObj& sortSpecObj, - const CollatorInterface* collator); + const BSONObj& sortSpecObj); bool isEOF() final; diff --git a/src/mongo/db/exec/sort_test.cpp b/src/mongo/db/exec/sort_test.cpp index 66d6f3c1d86..7c8a1696399 100644 --- a/src/mongo/db/exec/sort_test.cpp +++ b/src/mongo/db/exec/sort_test.cpp @@ -103,8 +103,12 @@ public: params.pattern = fromjson(patternStr); params.limit = limit; + // Create an ExpressionContext for the SortKeyGeneratorStage. + boost::intrusive_ptr<ExpressionContext> pExpCtx( + new ExpressionContext(getOpCtx(), collator)); + auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>( - getOpCtx(), queuedDataStage.release(), &ws, params.pattern, collator); + pExpCtx, queuedDataStage.release(), &ws, params.pattern); SortStage sort(getOpCtx(), params, &ws, sortKeyGen.release()); @@ -158,10 +162,13 @@ private: TEST_F(SortStageTest, SortEmptyWorkingSet) { WorkingSet ws; + // Create an ExpressionContext for the SortKeyGeneratorStage. + boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(getOpCtx(), nullptr)); + // QueuedDataStage will be owned by SortStage. auto queuedDataStage = std::make_unique<QueuedDataStage>(getOpCtx(), &ws); - auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>( - getOpCtx(), queuedDataStage.release(), &ws, BSONObj(), nullptr); + auto sortKeyGen = + std::make_unique<SortKeyGeneratorStage>(pExpCtx, queuedDataStage.release(), &ws, BSONObj()); SortStageParams params; SortStage sort(getOpCtx(), params, &ws, sortKeyGen.release()); diff --git a/src/mongo/db/index/SConscript b/src/mongo/db/index/SConscript index 07784425da7..3d7f912bd6f 100644 --- a/src/mongo/db/index/SConscript +++ b/src/mongo/db/index/SConscript @@ -53,6 +53,7 @@ env.Library( '$BUILD_DIR/mongo/db/mongohasher', '$BUILD_DIR/mongo/db/projection_exec_agg', '$BUILD_DIR/mongo/db/query/collation/collator_interface', + '$BUILD_DIR/mongo/db/query/sort_pattern', '$BUILD_DIR/third_party/s2/s2', 'expression_params', 'index_descriptor', @@ -163,5 +164,6 @@ env.CppUnitTest( "$BUILD_DIR/mongo/db/matcher/expressions", '$BUILD_DIR/mongo/db/mongohasher', '$BUILD_DIR/mongo/db/query/collation/collator_interface_mock', + '$BUILD_DIR/mongo/db/query/query_test_service_context', ], ) diff --git a/src/mongo/db/index/sort_key_generator.cpp b/src/mongo/db/index/sort_key_generator.cpp index d07f4d9bc62..82122a65448 100644 --- a/src/mongo/db/index/sort_key_generator.cpp +++ b/src/mongo/db/index/sort_key_generator.cpp @@ -36,40 +36,23 @@ namespace mongo { -SortKeyGenerator::SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterface* collator) - : _collator(collator) { +SortKeyGenerator::SortKeyGenerator(SortPattern sortPattern, const CollatorInterface* collator) + : _collator(collator), _sortPattern(std::move(sortPattern)) { BSONObjBuilder btreeBob; + size_t nFields = 0; - for (auto&& elt : sortSpec) { - if (elt.isNumber()) { - btreeBob.append(elt); - _patternPartTypes.push_back(SortPatternPartType::kFieldPath); - } else { - // If this field of the sort pattern is non-numeric, we expect it to be a text-score - // meta sort. - invariant(elt.type() == BSONType::Object); - invariant(elt.embeddedObject().nFields() == 1); - auto metaElem = elt.embeddedObject().firstElement(); - invariant(metaElem.fieldNameStringData() == "$meta"_sd); - if (metaElem.valueStringData() == "textScore"_sd) { - _patternPartTypes.push_back(SortPatternPartType::kMetaTextScore); - } else if (metaElem.valueStringData() == "randVal"_sd) { - _patternPartTypes.push_back(SortPatternPartType::kMetaRandVal); - } else if (metaElem.valueStringData() == "searchScore"_sd) { - uasserted(31218, "$meta sort by 'searchScore' metadata is not supported"); - } else if (metaElem.valueStringData() == "searchHighlights"_sd) { - uasserted(31219, "$meta sort by 'searchHighlights' metadata is not supported"); - } else { - uasserted(31138, "Illegal $meta sort: " + metaElem.valueStringData()); - } - _sortHasMeta = true; + for (auto&& part : _sortPattern) { + if (part.fieldPath) { + btreeBob.append(part.fieldPath->fullPath(), part.isAscending ? 1 : -1); + ++nFields; } } // The fake index key pattern used to generate Btree keys. _sortSpecWithoutMeta = btreeBob.obj(); + _sortHasMeta = nFields < _sortPattern.size(); - // If we're just sorting by meta, don't bother with all the key stuff. + // If we're just sorting by meta, don't bother creating an index key generator. if (_sortSpecWithoutMeta.isEmpty()) { return; } @@ -77,11 +60,11 @@ SortKeyGenerator::SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterf // We'll need to treat arrays as if we were to create an index over them. that is, we may need // to unnest the first level and consider each array element to decide the sort order. In order // to do this, we make a BtreeKeyGenerator. + std::vector<BSONElement> fixed(nFields); std::vector<const char*> fieldNames; - std::vector<BSONElement> fixed; - for (auto&& patternElt : _sortSpecWithoutMeta) { - fieldNames.push_back(patternElt.fieldName()); - fixed.push_back(BSONElement()); + fieldNames.reserve(fixed.size()); + for (auto&& elem : _sortSpecWithoutMeta) { + fieldNames.push_back(elem.fieldName()); } constexpr bool isSparse = false; @@ -105,10 +88,10 @@ StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromIndexKey(const WorkingSetMem invariant(!_sortHasMeta); BSONObjBuilder objBuilder; - for (BSONElement specElt : _sortSpecWithoutMeta) { - invariant(specElt.isNumber()); + for (auto&& elem : _sortSpecWithoutMeta) { BSONElement sortKeyElt; - invariant(member.getFieldDotted(specElt.fieldName(), &sortKeyElt)); + invariant(elem.isNumber()); + invariant(member.getFieldDotted(elem.fieldName(), &sortKeyElt)); // If we were to call 'collationAwareIndexKeyAppend' with a non-simple collation and a // 'sortKeyElt' representing a collated index key we would incorrectly encode for the // collation twice. This is not currently possible as the query planner will ensure that @@ -139,18 +122,19 @@ StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromDocument(const BSONObj& obj, // Merge metadata into the key. BSONObjIterator sortKeyIt(sortKeyNoMetadata.getValue()); - for (auto type : _patternPartTypes) { - switch (type) { - case SortPatternPartType::kFieldPath: { - invariant(sortKeyIt.more()); - mergedKeyBob.append(sortKeyIt.next()); - continue; - } - case SortPatternPartType::kMetaTextScore: { + for (auto& part : _sortPattern) { + if (part.fieldPath) { + invariant(sortKeyIt.more()); + mergedKeyBob.append(sortKeyIt.next()); + continue; + } + invariant(part.expression); + switch (part.expression->getMetaType()) { + case ExpressionMeta::MetaType::TEXT_SCORE: { mergedKeyBob.append("", metadata->textScore); continue; } - case SortPatternPartType::kMetaRandVal: { + case ExpressionMeta::MetaType::RAND_VAL: { mergedKeyBob.append("", metadata->randVal); continue; } diff --git a/src/mongo/db/index/sort_key_generator.h b/src/mongo/db/index/sort_key_generator.h index 64afec01cb8..2ddc186f340 100644 --- a/src/mongo/db/index/sort_key_generator.h +++ b/src/mongo/db/index/sort_key_generator.h @@ -34,6 +34,7 @@ #include "mongo/db/index/btree_key_generator.h" #include "mongo/db/operation_context.h" #include "mongo/db/query/collation/collator_interface.h" +#include "mongo/db/query/sort_pattern.h" namespace mongo { @@ -50,11 +51,11 @@ public: }; /** - * Constructs a sort key generator which will generate keys for sort pattern 'sortSpec'. The + * Constructs a sort key generator which will generate keys for sort pattern 'sortPattern'. The * keys will incorporate the collation given by 'collator', and thus when actually compared to * one another should use the simple collation. */ - SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterface* collator); + SortKeyGenerator(SortPattern sortPattern, const CollatorInterface* collator); /** * Returns the key which should be used to sort the WorkingSetMember, or a non-OK status if no @@ -77,14 +78,6 @@ public: StatusWith<BSONObj> getSortKeyFromDocument(const BSONObj& obj, const Metadata*) const; private: - // Describes whether a component of the sort pattern is a field path (e.g. sort by "a.b"), or - // else describes the type of $meta sort. - enum class SortPatternPartType { - kFieldPath, - kMetaTextScore, - kMetaRandVal, - }; - // Extracts the sort key from a WorkingSetMember which represents an index key. It is illegal to // call this if the working set member is not in RID_AND_IDX state. It is also illegal to call // this if the sort pattern has any $meta components. @@ -97,14 +90,12 @@ private: const CollatorInterface* _collator = nullptr; + SortPattern _sortPattern; + // The sort pattern with any $meta sort components stripped out, since the underlying index key // generator does not understand $meta sort. BSONObj _sortSpecWithoutMeta; - // For each element of the raw sort spec, describes whether the element is sorting by a field - // path or by a particular meta-sort. - std::vector<SortPatternPartType> _patternPartTypes; - // If we're not sorting with a $meta value we can short-cut some work. bool _sortHasMeta = false; diff --git a/src/mongo/db/index/sort_key_generator_test.cpp b/src/mongo/db/index/sort_key_generator_test.cpp index 36d910349c0..701c8975fd4 100644 --- a/src/mongo/db/index/sort_key_generator_test.cpp +++ b/src/mongo/db/index/sort_key_generator_test.cpp @@ -33,6 +33,7 @@ #include "mongo/bson/json.h" #include "mongo/db/index/sort_key_generator.h" +#include "mongo/db/pipeline/expression_context_for_test.h" #include "mongo/db/query/collation/collator_interface_mock.h" #include "mongo/unittest/death_test.h" #include "mongo/unittest/unittest.h" @@ -40,15 +41,24 @@ namespace mongo { namespace { +// A method to create mock ExpressionContexts with a specified collation +std::unique_ptr<SortKeyGenerator> makeSortKeyGen(const BSONObj& sortSpec, + const CollatorInterface* collator) { + boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContextForTest()); + pExpCtx->setCollator(collator); + SortPattern sortPattern{sortSpec, pExpCtx}; + return std::make_unique<SortKeyGenerator>(std::move(sortPattern), collator); +} + TEST(SortKeyGeneratorTest, ExtractNumberKeyForNonCompoundSortNonNested) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, a: 5}"), nullptr); ASSERT_OK(sortKey.getStatus()); ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 5)); } TEST(SortKeyGeneratorTest, ExtractNumberKeyFromDocWithSeveralFields) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -56,7 +66,7 @@ TEST(SortKeyGeneratorTest, ExtractNumberKeyFromDocWithSeveralFields) { } TEST(SortKeyGeneratorTest, ExtractStringKeyNonCompoundNonNested) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument( fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -66,7 +76,7 @@ TEST(SortKeyGeneratorTest, ExtractStringKeyNonCompoundNonNested) { } TEST(SortKeyGeneratorTest, CompoundSortPattern) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1 << "b" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1 << "b" << 1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument( fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -74,7 +84,7 @@ TEST(SortKeyGeneratorTest, CompoundSortPattern) { } TEST(SortKeyGeneratorTest, CompoundSortPatternWithDottedPath) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("c.a" << 1 << "b" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("c.a" << 1 << "b" << 1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument( fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -82,7 +92,7 @@ TEST(SortKeyGeneratorTest, CompoundSortPatternWithDottedPath) { } TEST(SortKeyGeneratorTest, CompoundPatternLeadingFieldIsArray) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("c" << 1 << "b" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("c" << 1 << "b" << 1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument( fromjson("{_id: 0, z: 'thing1', a: 99, c: [2, 4, 1], b: 16}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -91,7 +101,7 @@ TEST(SortKeyGeneratorTest, CompoundPatternLeadingFieldIsArray) { TEST(SortKeyGeneratorTest, ExtractStringSortKeyWithCollatorUsesComparisonKey) { CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString); - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator); auto sortKey = sortKeyGen->getSortKeyFromDocument( fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -102,7 +112,7 @@ TEST(SortKeyGeneratorTest, ExtractStringSortKeyWithCollatorUsesComparisonKey) { TEST(SortKeyGeneratorTest, CollatorHasNoEffectWhenExtractingNonStringSortKey) { CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString); - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator); auto sortKey = sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -110,7 +120,7 @@ TEST(SortKeyGeneratorTest, CollatorHasNoEffectWhenExtractingNonStringSortKey) { } TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysChoosesCorrectKey) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << -1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << -1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, a: [1, 2, 3, 4]}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -119,7 +129,7 @@ TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysChoosesCorrectKey) { TEST(SortKeyGeneratorTest, EnsureSortKeyGenerationForArraysRespectsCollation) { CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString); - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator); auto sortKey = sortKeyGen->getSortKeyFromDocument( fromjson("{_id: 0, a: ['aaz', 'zza', 'yya', 'zzb']}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -129,7 +139,7 @@ TEST(SortKeyGeneratorTest, EnsureSortKeyGenerationForArraysRespectsCollation) { } TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysRespectsCompoundOrdering) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a.b" << 1 << "a.c" << -1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a.b" << 1 << "a.c" << -1), nullptr); auto sortKey = sortKeyGen->getSortKeyFromDocument( fromjson("{_id: 0, a: [{b: 1, c: 0}, {b: 0, c: 3}, {b: 0, c: 1}]}"), nullptr); ASSERT_OK(sortKey.getStatus()); @@ -138,62 +148,59 @@ TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysRespectsCompoundOrdering) { DEATH_TEST(SortKeyGeneratorTest, SortPatternComponentWithStringIsFatal, - "Invariant failure elt.type() == BSONType::Object") { - MONGO_COMPILER_VARIABLE_UNUSED auto ignored = std::make_unique<SortKeyGenerator>(BSON("a" - << "foo"), - nullptr); + "Illegal key in $sort specification: a: \"foo\"") { + MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a" + << "foo"), + nullptr); } DEATH_TEST(SortKeyGeneratorTest, SortPatternComponentWhoseObjectHasMultipleKeysIsFatal, - "Invariant failure elt.embeddedObject().nFields() == 1") { - MONGO_COMPILER_VARIABLE_UNUSED auto ignored = - std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta" - << "textScore" - << "extra" << 1)), - nullptr); + "Cannot have additional keys in a $meta sort specification") { + MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a" << BSON("$meta" + << "textScore" + << "extra" << 1)), + nullptr); } DEATH_TEST(SortKeyGeneratorTest, SortPatternComponentWithNonMetaObjectSortIsFatal, - "Invariant failure metaElem.fieldNameStringData() == \"$meta\"_sd") { - MONGO_COMPILER_VARIABLE_UNUSED auto ignored = - std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$unknown" - << "textScore")), - nullptr); + "$meta is the only expression supported by $sort right now") { + MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a" << BSON("$unknown" + << "textScore")), + nullptr); } DEATH_TEST(SortKeyGeneratorTest, SortPatternComponentWithUnknownMetaKeywordIsFatal, "Illegal $meta sort: unknown") { - MONGO_COMPILER_VARIABLE_UNUSED auto ignored = - std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta" - << "unknown")), - nullptr); + MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a" << BSON("$meta" + << "unknown")), + nullptr); } DEATH_TEST(SortKeyGeneratorTest, NoMetadataWhenPatternHasMetaTextScoreIsFatal, "Invariant failure metadata") { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta" - << "textScore")), - nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta" + << "textScore")), + nullptr); uassertStatusOK(sortKeyGen->getSortKeyFromDocument(BSONObj{}, nullptr).getStatus()); } DEATH_TEST(SortKeyGeneratorTest, NoMetadataWhenPatternHasMetaRandValIsFatal, "Invariant failure metadata") { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta" - << "randVal")), - nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta" + << "randVal")), + nullptr); uassertStatusOK(sortKeyGen->getSortKeyFromDocument(BSONObj{}, nullptr).getStatus()); } TEST(SortKeyGeneratorTest, CanGenerateKeysForTextScoreMetaSort) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta" - << "textScore")), - nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta" + << "textScore")), + nullptr); SortKeyGenerator::Metadata metadata; metadata.textScore = 1.5; auto sortKey = sortKeyGen->getSortKeyFromDocument(BSONObj{}, &metadata); @@ -202,9 +209,9 @@ TEST(SortKeyGeneratorTest, CanGenerateKeysForTextScoreMetaSort) { } TEST(SortKeyGeneratorTest, CanGenerateKeysForRandValMetaSort) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta" - << "randVal")), - nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta" + << "randVal")), + nullptr); SortKeyGenerator::Metadata metadata; metadata.randVal = 0.3; auto sortKey = sortKeyGen->getSortKeyFromDocument(BSONObj{}, &metadata); @@ -215,7 +222,7 @@ TEST(SortKeyGeneratorTest, CanGenerateKeysForRandValMetaSort) { TEST(SortKeyGeneratorTest, CanGenerateKeysForCompoundMetaSort) { BSONObj pattern = fromjson( "{a: 1, b: {$meta: 'randVal'}, c: {$meta: 'textScore'}, d: -1, e: {$meta: 'textScore'}}"); - auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr); + auto sortKeyGen = makeSortKeyGen(pattern, nullptr); SortKeyGenerator::Metadata metadata; metadata.randVal = 0.3; metadata.textScore = 1.5; @@ -258,7 +265,7 @@ private: }; TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithObj) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr); setRecordIdAndObj(BSON("x" << 1 << "a" << 2 << "y" << 3)); auto sortKey = sortKeyGen->getSortKey(member()); ASSERT_OK(sortKey); @@ -266,7 +273,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithObj) } TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithOwnedObj) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr); setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3)); auto sortKey = sortKeyGen->getSortKey(member()); ASSERT_OK(sortKey); @@ -275,7 +282,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithOwne TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateKeyFromWSMForTextScoreMetaSort) { BSONObj pattern = fromjson("{a: 1, b: {$meta: 'textScore'}, c: -1}}"); - auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr); + auto sortKeyGen = makeSortKeyGen(pattern, nullptr); setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3 << "c" << BSON_ARRAY(4 << 5 << 6))); member().metadata().setTextScore(9.9); auto sortKey = sortKeyGen->getSortKey(member()); @@ -284,7 +291,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateKeyFromWSMForTextScoreMetaSort } TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyState) { - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr); setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << 2 << "" << 3)); auto sortKey = sortKeyGen->getSortKey(member()); ASSERT_OK(sortKey); @@ -293,7 +300,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyState) TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyStateWithCollator) { CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString); - auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator); + auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator); setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << "string1" @@ -310,7 +317,7 @@ DEATH_TEST_F(SortKeyGeneratorWorkingSetTest, DeathOnAttemptToGetSortKeyFromIndexKeyWithMetadata, "Invariant failure !_sortHasMeta") { BSONObj pattern = fromjson("{z: {$meta: 'textScore'}}"); - auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr); + auto sortKeyGen = makeSortKeyGen(pattern, nullptr); setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << 2 << "" << 3)); member().metadata().setTextScore(9.9); MONGO_COMPILER_VARIABLE_UNUSED auto ignored = sortKeyGen->getSortKey(member()); diff --git a/src/mongo/db/pipeline/document_source_sort.cpp b/src/mongo/db/pipeline/document_source_sort.cpp index fe32168c170..09afc8fc2e6 100644 --- a/src/mongo/db/pipeline/document_source_sort.cpp +++ b/src/mongo/db/pipeline/document_source_sort.cpp @@ -112,7 +112,7 @@ DocumentSourceSort::DocumentSourceSort(const boost::intrusive_ptr<ExpressionCont pExpCtx->allowDiskUse}), // The SortKeyGenerator expects the expressions to be serialized in order to detect a sort // by a metadata field. - _sortKeyGen({sortOrder, pExpCtx->getCollator()}) { + _sortKeyGen({{sortOrder, pExpCtx}, pExpCtx->getCollator()}) { uassert(15976, "$sort stage must have at least one sort key", !_sortExecutor->sortPattern().empty()); diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h index ab70a4d4319..641d35d5aa4 100644 --- a/src/mongo/db/pipeline/expression.h +++ b/src/mongo/db/pipeline/expression.h @@ -1642,6 +1642,13 @@ private: class ExpressionMeta final : public Expression { public: + enum MetaType { + TEXT_SCORE, + RAND_VAL, + SEARCH_SCORE, + SEARCH_HIGHLIGHTS, + }; + Value serialize(bool explain) const final; Value evaluate(const Document& root, Variables* variables) const final; @@ -1654,17 +1661,14 @@ public: return visitor->visit(this); } + MetaType getMetaType() { + return _metaType; + } + protected: void _doAddDependencies(DepsTracker* deps) const final; private: - enum MetaType { - TEXT_SCORE, - RAND_VAL, - SEARCH_SCORE, - SEARCH_HIGHLIGHTS, - }; - ExpressionMeta(const boost::intrusive_ptr<ExpressionContext>& expCtx, MetaType metaType); MetaType _metaType; diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp index 03424a838d6..e0a164bd2a2 100644 --- a/src/mongo/db/query/canonical_query.cpp +++ b/src/mongo/db/query/canonical_query.cpp @@ -205,7 +205,7 @@ StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::canonicalize( // Make the CQ we'll hopefully return. std::unique_ptr<CanonicalQuery> cq(new CanonicalQuery()); Status initStatus = cq->init(opCtx, - nullptr, // no expression context + baseQuery.getExpCtx(), std::move(qr), baseQuery.canHaveNoopMatchNodes(), root->shallowClone(), @@ -261,6 +261,9 @@ void CanonicalQuery::setCollator(std::unique_ptr<CollatorInterface> collator) { // the object owned by '_collator'. We must associate the match expression tree with the new // value of '_collator'. _root->setCollator(_collator.get()); + + // In a similar vein, we must give the ExpressionContext the same collator. + _expCtx->setCollator(_collator.get()); } // static diff --git a/src/mongo/db/query/canonical_query.h b/src/mongo/db/query/canonical_query.h index 366265add3c..bbd7103707d 100644 --- a/src/mongo/db/query/canonical_query.h +++ b/src/mongo/db/query/canonical_query.h @@ -181,6 +181,10 @@ public: return _canHaveNoopMatchNodes; } + const boost::intrusive_ptr<ExpressionContext>& getExpCtx() const { + return _expCtx; + } + private: // You must go through canonicalize to create a CanonicalQuery. CanonicalQuery() {} diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index 178821fb880..7e8c0d1e9b4 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -395,11 +395,10 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx, // Add a SortKeyGeneratorStage if there is a $meta sortKey projection. if (canonicalQuery->getProj()->wantSortKey()) { root = std::make_unique<SortKeyGeneratorStage>( - opCtx, + canonicalQuery->getExpCtx(), root.release(), ws, - canonicalQuery->getQueryRequest().getSort(), - canonicalQuery->getCollator()); + canonicalQuery->getQueryRequest().getSort()); } // Stuff the right data into the params depending on what proj impl we use. diff --git a/src/mongo/db/query/sort_pattern.cpp b/src/mongo/db/query/sort_pattern.cpp index 17df5adcca7..1964164cf5e 100644 --- a/src/mongo/db/query/sort_pattern.cpp +++ b/src/mongo/db/query/sort_pattern.cpp @@ -51,7 +51,22 @@ SortPattern::SortPattern(const BSONObj& obj, metaDoc.nFields() == 1); VariablesParseState vps = pExpCtx->variablesParseState; - patternPart.expression = ExpressionMeta::parse(pExpCtx, metaDoc.firstElement(), vps); + BSONElement metaElem = metaDoc.firstElement(); + + if (metaElem.valueStringData() == "textScore"_sd) { + // Valid meta sort. Just fall through. + } else if (metaElem.valueStringData() == "randVal"_sd) { + // Valid meta sort. Just fall through. + } else if (metaElem.valueStringData() == "searchScore"_sd) { + uasserted(31218, "$meta sort by 'searchScore' metadata is not supported"); + } else if (metaElem.valueStringData() == "searchHighlights"_sd) { + uasserted(31219, "$meta sort by 'searchHighlights' metadata is not supported"); + } else { + uasserted(31138, + str::stream() << "Illegal $meta sort: " << metaElem.valueStringData()); + } + patternPart.expression = + static_cast<ExpressionMeta*>(ExpressionMeta::parse(pExpCtx, metaElem, vps).get()); // If sorting by textScore, sort highest scores first. If sorting by randVal, order // doesn't matter, so just always use descending. diff --git a/src/mongo/db/query/sort_pattern.h b/src/mongo/db/query/sort_pattern.h index 5dbb00c7ca2..d818bd82f38 100644 --- a/src/mongo/db/query/sort_pattern.h +++ b/src/mongo/db/query/sort_pattern.h @@ -49,7 +49,7 @@ public: struct SortPatternPart { bool isAscending = true; boost::optional<FieldPath> fieldPath; - boost::intrusive_ptr<Expression> expression; + boost::intrusive_ptr<ExpressionMeta> expression; }; SortPattern(const BSONObj&, const boost::intrusive_ptr<ExpressionContext>&); diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp index dcba367378b..382e3652354 100644 --- a/src/mongo/db/query/stage_builder.cpp +++ b/src/mongo/db/query/stage_builder.cpp @@ -142,8 +142,7 @@ PlanStage* buildStages(OperationContext* opCtx, if (nullptr == childStage) { return nullptr; } - return new SortKeyGeneratorStage( - opCtx, childStage, ws, keyGenNode->sortSpec, cq.getCollator()); + return new SortKeyGeneratorStage(cq.getExpCtx(), childStage, ws, keyGenNode->sortSpec); } case STAGE_PROJECTION_DEFAULT: { auto pn = static_cast<const ProjectionNodeDefault*>(root); diff --git a/src/mongo/dbtests/query_stage_ensure_sorted.cpp b/src/mongo/dbtests/query_stage_ensure_sorted.cpp index 49e9383442f..20d5750c958 100644 --- a/src/mongo/dbtests/query_stage_ensure_sorted.cpp +++ b/src/mongo/dbtests/query_stage_ensure_sorted.cpp @@ -78,10 +78,15 @@ public: queuedDataStage->pushBack(id); } + // Create a mock ExpressionContext. + boost::intrusive_ptr<ExpressionContext> pExpCtx( + new ExpressionContext(opCtx.get(), collator)); + pExpCtx->setCollator(collator); + // Initialization. BSONObj pattern = fromjson(patternStr); auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>( - opCtx.get(), queuedDataStage.release(), &ws, pattern, collator); + pExpCtx, queuedDataStage.release(), &ws, pattern); EnsureSortedStage ess(opCtx.get(), pattern, &ws, sortKeyGen.release()); WorkingSetID id = WorkingSet::INVALID_ID; PlanStage::StageState state = PlanStage::NEED_TIME; @@ -115,10 +120,13 @@ protected: TEST_F(QueryStageEnsureSortedTest, EnsureSortedEmptyWorkingSet) { auto opCtx = _serviceContext.makeOperationContext(); + // Create a mock ExpressionContext. + boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), nullptr)); + WorkingSet ws; auto queuedDataStage = std::make_unique<QueuedDataStage>(opCtx.get(), &ws); - auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>( - opCtx.get(), queuedDataStage.release(), &ws, BSONObj(), nullptr); + auto sortKeyGen = + std::make_unique<SortKeyGeneratorStage>(pExpCtx, queuedDataStage.release(), &ws, BSONObj()); EnsureSortedStage ess(opCtx.get(), BSONObj(), &ws, sortKeyGen.release()); WorkingSetID id = WorkingSet::INVALID_ID; diff --git a/src/mongo/dbtests/query_stage_sort.cpp b/src/mongo/dbtests/query_stage_sort.cpp index 5b855933793..0919fdea1ad 100644 --- a/src/mongo/dbtests/query_stage_sort.cpp +++ b/src/mongo/dbtests/query_stage_sort.cpp @@ -120,7 +120,7 @@ public: params.limit = limit(); auto keyGenStage = std::make_unique<SortKeyGeneratorStage>( - &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr); + _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern); auto ss = std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release()); @@ -158,7 +158,7 @@ public: params.limit = limit(); auto keyGenStage = std::make_unique<SortKeyGeneratorStage>( - &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr); + _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern); auto sortStage = std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release()); @@ -227,6 +227,7 @@ public: protected: const ServiceContext::UniqueOperationContext _txnPtr = cc().makeOperationContext(); OperationContext& _opCtx = *_txnPtr; + boost::intrusive_ptr<ExpressionContext> _pExpCtx = new ExpressionContext(&_opCtx, nullptr); DBDirectClient _client; }; @@ -558,7 +559,7 @@ public: params.pattern = BSON("b" << -1 << "c" << 1 << "a" << 1); auto keyGenStage = std::make_unique<SortKeyGeneratorStage>( - &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr); + _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern); auto sortStage = std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release()); diff --git a/src/mongo/dbtests/query_stage_sort_key_generator.cpp b/src/mongo/dbtests/query_stage_sort_key_generator.cpp index daa667313e4..8675a3e7a51 100644 --- a/src/mongo/dbtests/query_stage_sort_key_generator.cpp +++ b/src/mongo/dbtests/query_stage_sort_key_generator.cpp @@ -65,6 +65,7 @@ BSONObj extractKeyFromKeyGenStage(SortKeyGeneratorStage* sortKeyGen, WorkingSet* BSONObj extractSortKey(const char* sortSpec, const char* doc, const CollatorInterface* collator) { QueryTestServiceContext serviceContext; auto opCtx = serviceContext.makeOperationContext(); + boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), collator)); WorkingSet workingSet; @@ -77,7 +78,7 @@ BSONObj extractSortKey(const char* sortSpec, const char* doc, const CollatorInte BSONObj sortPattern = fromjson(sortSpec); SortKeyGeneratorStage sortKeyGen{ - opCtx.get(), mockStage.release(), &workingSet, std::move(sortPattern), collator}; + pExpCtx, mockStage.release(), &workingSet, std::move(sortPattern)}; return extractKeyFromKeyGenStage(&sortKeyGen, &workingSet); } @@ -93,6 +94,7 @@ BSONObj extractSortKeyCovered(const char* sortSpec, const CollatorInterface* collator) { QueryTestServiceContext serviceContext; auto opCtx = serviceContext.makeOperationContext(); + boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), collator)); WorkingSet workingSet; @@ -105,7 +107,7 @@ BSONObj extractSortKeyCovered(const char* sortSpec, BSONObj sortPattern = fromjson(sortSpec); SortKeyGeneratorStage sortKeyGen{ - opCtx.get(), mockStage.release(), &workingSet, std::move(sortPattern), collator}; + pExpCtx, mockStage.release(), &workingSet, std::move(sortPattern)}; return extractKeyFromKeyGenStage(&sortKeyGen, &workingSet); } |