summaryrefslogtreecommitdiff
path: root/src/mongo/db/exec
diff options
context:
space:
mode:
authorJustin Seyster <justin.seyster@mongodb.com>2020-05-29 15:37:48 -0400
committerEvergreen Agent <no-reply@evergreen.mongodb.com>2020-07-16 03:01:28 +0000
commitb1ecda7953b2d9e23857ceeb09f10f9cd18c959b (patch)
tree75d5c0322e310bf4be8b0b59de6c95d7f2efda72 /src/mongo/db/exec
parentb348b91ffdfab32e5a1ec9e63c9def884adf1acb (diff)
downloadmongo-b1ecda7953b2d9e23857ceeb09f10f9cd18c959b.tar.gz
SERVER-48721 Support for $returnKey in SBE
Diffstat (limited to 'src/mongo/db/exec')
-rw-r--r--src/mongo/db/exec/sbe/parser/parser.cpp69
-rw-r--r--src/mongo/db/exec/sbe/parser/parser.h8
-rw-r--r--src/mongo/db/exec/sbe/stages/ix_scan.cpp46
-rw-r--r--src/mongo/db/exec/sbe/stages/ix_scan.h37
-rw-r--r--src/mongo/db/exec/sbe/values/id_generators.h8
5 files changed, 142 insertions, 26 deletions
diff --git a/src/mongo/db/exec/sbe/parser/parser.cpp b/src/mongo/db/exec/sbe/parser/parser.cpp
index 546d26c8c5f..1e410bbb176 100644
--- a/src/mongo/db/exec/sbe/parser/parser.cpp
+++ b/src/mongo/db/exec/sbe/parser/parser.cpp
@@ -86,7 +86,7 @@ static constexpr auto kSyntax = R"(
IXSCAN <- 'ixscan' IDENT? # optional variable name of the root object (record) delivered by the scan
IDENT? # optional variable name of the record id delivered by the scan
- IDENT_LIST_WITH_RENAMES # list of projected fields (may be empty)
+ IX_KEY_LIST_WITH_RENAMES # list of projected fields (may be empty)
IDENT # collection name
IDENT # index name to scan
FORWARD_FLAG # forward scan or not
@@ -95,7 +95,7 @@ static constexpr auto kSyntax = R"(
IDENT # variable name of the high key
IDENT? # optional variable name of the root object (record) delivered by the scan
IDENT? # optional variable name of the record id delivered by the scan
- IDENT_LIST_WITH_RENAMES # list of projected fields (may be empty)
+ IX_KEY_LIST_WITH_RENAMES # list of projected fields (may be empty)
IDENT # collection name
IDENT # index name to scan
FORWARD_FLAG # forward scan or not
@@ -207,6 +207,9 @@ static constexpr auto kSyntax = R"(
IDENT_LIST_WITH_RENAMES <- '[' (IDENT_WITH_RENAME (',' IDENT_WITH_RENAME)*)? ']'
IDENT_WITH_RENAME <- IDENT ('=' IDENT)?
+ IX_KEY_LIST_WITH_RENAMES <- '[' (IX_KEY_WITH_RENAME (',' IX_KEY_WITH_RENAME)*)? ']'
+ IX_KEY_WITH_RENAME <- IDENT '=' NUMBER
+
IDENT_LIST <- '[' (IDENT (',' IDENT)*)? ']'
IDENT <- RAW_IDENT/ESC_IDENT
@@ -223,6 +226,30 @@ static constexpr auto kSyntax = R"(
%word <- [a-z]+
)";
+std::pair<IndexKeysInclusionSet, sbe::value::SlotVector> Parser::lookupIndexKeyRenames(
+ const std::vector<std::string>& renames, const std::vector<size_t>& indexKeys) {
+ IndexKeysInclusionSet indexKeysInclusion;
+
+ // Each indexKey is associated with the parallel remap from the 'renames' vector. This
+ // map explicitly binds each indexKey with its remap and sorts them by 'indexKey' order.
+ stdx::unordered_map<int, sbe::value::SlotId> slotsByKeyIndex;
+ invariant(renames.size() == indexKeys.size());
+ for (size_t idx = 0; idx < renames.size(); idx++) {
+ uassert(4872100,
+ str::stream() << "Cannot project index key at position " << indexKeys[idx],
+ indexKeys[idx] < indexKeysInclusion.size());
+ slotsByKeyIndex[indexKeys[idx]] = lookupSlotStrict(renames[idx]);
+ }
+
+ sbe::value::SlotVector slots;
+ for (auto&& [indexKey, slot] : slotsByKeyIndex) {
+ indexKeysInclusion.set(indexKey);
+ slots.push_back(slot);
+ }
+
+ return {indexKeysInclusion, std::move(slots)};
+}
+
void Parser::walkChildren(AstQuery& ast) {
for (const auto& node : ast.nodes) {
walk(*node);
@@ -271,6 +298,24 @@ void Parser::walkIdentListWithRename(AstQuery& ast) {
}
}
+void Parser::walkIxKeyWithRename(AstQuery& ast) {
+ walkChildren(ast);
+
+ ast.rename = ast.nodes[0]->identifier;
+
+ // This token cannot be negative, because the parser does not accept a "-" prefix.
+ ast.indexKey = static_cast<size_t>(std::stoi(ast.nodes[1]->token));
+}
+
+void Parser::walkIxKeyListWithRename(AstQuery& ast) {
+ walkChildren(ast);
+
+ for (auto& node : ast.nodes) {
+ ast.indexKeys.emplace_back(std::move(node->indexKey));
+ ast.renames.emplace_back(std::move(node->rename));
+ }
+}
+
void Parser::walkProjectList(AstQuery& ast) {
walkChildren(ast);
@@ -625,13 +670,16 @@ void Parser::walkIndexScan(AstQuery& ast) {
collection ? NamespaceStringOrUUID{dbName, collection->uuid()} : nssColl;
const auto forward = (ast.nodes[forwardPos]->token == "true") ? true : false;
+ auto [indexKeysInclusion, vars] =
+ lookupIndexKeyRenames(ast.nodes[projectsPos]->renames, ast.nodes[projectsPos]->indexKeys);
+
ast.stage = makeS<IndexScanStage>(name,
indexName,
forward,
lookupSlot(recordName),
lookupSlot(recordIdName),
- ast.nodes[projectsPos]->identifiers,
- lookupSlots(ast.nodes[projectsPos]->renames),
+ indexKeysInclusion,
+ vars,
boost::none,
boost::none,
nullptr,
@@ -678,13 +726,16 @@ void Parser::walkIndexSeek(AstQuery& ast) {
collection ? NamespaceStringOrUUID{dbName, collection->uuid()} : nssColl;
const auto forward = (ast.nodes[forwardPos]->token == "true") ? true : false;
+ auto [indexKeysInclusion, vars] =
+ lookupIndexKeyRenames(ast.nodes[projectsPos]->renames, ast.nodes[projectsPos]->indexKeys);
+
ast.stage = makeS<IndexScanStage>(name,
indexName,
forward,
lookupSlot(recordName),
lookupSlot(recordIdName),
- ast.nodes[projectsPos]->identifiers,
- lookupSlots(ast.nodes[projectsPos]->renames),
+ indexKeysInclusion,
+ vars,
lookupSlot(ast.nodes[0]->identifier),
lookupSlot(ast.nodes[1]->identifier),
nullptr,
@@ -1339,6 +1390,12 @@ void Parser::walk(AstQuery& ast) {
case "IDENT_LIST_WITH_RENAMES"_:
walkIdentListWithRename(ast);
break;
+ case "IX_KEY_WITH_RENAME"_:
+ walkIxKeyWithRename(ast);
+ break;
+ case "IX_KEY_LIST_WITH_RENAMES"_:
+ walkIxKeyListWithRename(ast);
+ break;
case "PROJECT_LIST"_:
walkProjectList(ast);
break;
diff --git a/src/mongo/db/exec/sbe/parser/parser.h b/src/mongo/db/exec/sbe/parser/parser.h
index 9d17c7deb17..87e5b04be7f 100644
--- a/src/mongo/db/exec/sbe/parser/parser.h
+++ b/src/mongo/db/exec/sbe/parser/parser.h
@@ -33,6 +33,7 @@
#include <third_party/peglib/peglib.h>
#include "mongo/db/exec/sbe/expressions/expression.h"
+#include "mongo/db/exec/sbe/stages/ix_scan.h"
#include "mongo/db/exec/sbe/stages/spool.h"
#include "mongo/db/exec/sbe/stages/stages.h"
#include "mongo/db/exec/sbe/values/id_generators.h"
@@ -42,8 +43,10 @@ namespace mongo {
namespace sbe {
struct ParsedQueryTree {
std::string identifier;
+ size_t indexKey;
std::string rename;
std::vector<std::string> identifiers;
+ std::vector<size_t> indexKeys;
std::vector<std::string> renames;
std::unique_ptr<PlanStage> stage;
@@ -153,6 +156,9 @@ private:
return result;
}
+ std::pair<IndexKeysInclusionSet, sbe::value::SlotVector> lookupIndexKeyRenames(
+ const std::vector<std::string>& renames, const std::vector<size_t>& indexKeys);
+
SpoolId lookupSpoolBuffer(const std::string& name) {
if (_spoolBuffersLookupTable.find(name) == _spoolBuffersLookupTable.end()) {
_spoolBuffersLookupTable[name] = _spoolIdGenerator.generate();
@@ -165,6 +171,8 @@ private:
void walkIdentList(AstQuery& ast);
void walkIdentWithRename(AstQuery& ast);
void walkIdentListWithRename(AstQuery& ast);
+ void walkIxKeyWithRename(AstQuery& ast);
+ void walkIxKeyListWithRename(AstQuery& ast);
void walkProjectList(AstQuery& ast);
void walkAssign(AstQuery& ast);
diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
index 9e3cdb21d44..f26dd1dbadb 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
@@ -33,6 +33,7 @@
#include "mongo/db/catalog/index_catalog.h"
#include "mongo/db/exec/sbe/expressions/expression.h"
+#include "mongo/db/exec/sbe/values/bson.h"
#include "mongo/db/index/index_access_method.h"
namespace mongo::sbe {
@@ -41,7 +42,7 @@ IndexScanStage::IndexScanStage(const NamespaceStringOrUUID& name,
bool forward,
boost::optional<value::SlotId> recordSlot,
boost::optional<value::SlotId> recordIdSlot,
- std::vector<std::string> fields,
+ IndexKeysInclusionSet indexKeysToInclude,
value::SlotVector vars,
boost::optional<value::SlotId> seekKeySlotLow,
boost::optional<value::SlotId> seekKeySlotHi,
@@ -53,15 +54,16 @@ IndexScanStage::IndexScanStage(const NamespaceStringOrUUID& name,
_forward(forward),
_recordSlot(recordSlot),
_recordIdSlot(recordIdSlot),
- _fields(std::move(fields)),
+ _indexKeysToInclude(indexKeysToInclude),
_vars(std::move(vars)),
_seekKeySlotLow(seekKeySlotLow),
_seekKeySlotHi(seekKeySlotHi),
_tracker(tracker) {
- invariant(_fields.size() == _vars.size());
// The valid state is when both boundaries, or none is set, or only low key is set.
invariant((_seekKeySlotLow && _seekKeySlotHi) || (!_seekKeySlotLow && !_seekKeySlotHi) ||
(_seekKeySlotLow && !_seekKeySlotHi));
+
+ invariant(_indexKeysToInclude.count() == _vars.size());
}
std::unique_ptr<PlanStage> IndexScanStage::clone() const {
@@ -70,7 +72,7 @@ std::unique_ptr<PlanStage> IndexScanStage::clone() const {
_forward,
_recordSlot,
_recordIdSlot,
- _fields,
+ _indexKeysToInclude,
_vars,
_seekKeySlotLow,
_seekKeySlotHi,
@@ -87,12 +89,10 @@ void IndexScanStage::prepare(CompileCtx& ctx) {
_recordIdAccessor = std::make_unique<value::ViewOfValueAccessor>();
}
- for (size_t idx = 0; idx < _fields.size(); ++idx) {
- auto [it, inserted] =
- _fieldAccessors.emplace(_fields[idx], std::make_unique<value::ViewOfValueAccessor>());
- uassert(4822821, str::stream() << "duplicate field: " << _fields[idx], inserted);
- auto [itRename, insertedRename] = _varAccessors.emplace(_vars[idx], it->second.get());
- uassert(4822822, str::stream() << "duplicate field: " << _vars[idx], insertedRename);
+ _accessors.resize(_vars.size());
+ for (size_t idx = 0; idx < _accessors.size(); ++idx) {
+ auto [it, inserted] = _accessorMap.emplace(_vars[idx], &_accessors[idx]);
+ uassert(4822821, str::stream() << "duplicate slot: " << _vars[idx], inserted);
}
if (_seekKeySlotLow) {
@@ -112,7 +112,7 @@ value::SlotAccessor* IndexScanStage::getAccessor(CompileCtx& ctx, value::SlotId
return _recordIdAccessor.get();
}
- if (auto it = _varAccessors.find(slot); it != _varAccessors.end()) {
+ if (auto it = _accessorMap.find(slot); it != _accessorMap.end()) {
return it->second;
}
@@ -206,6 +206,10 @@ void IndexScanStage::open(bool reOpen) {
_seekKeyLow = &_startPoint;
_seekKeyHi = nullptr;
}
+
+ // TODO SERVER-49385: When the 'prepare()' phase takes the collection lock, it will be
+ // possible to intialize '_ordering' there instead of here.
+ _ordering = entry->ordering();
} else {
_cursor.reset();
}
@@ -256,6 +260,12 @@ PlanState IndexScanStage::getNext() {
value::bitcastFrom<int64_t>(_nextRecord->loc.repr()));
}
+ if (_accessors.size()) {
+ _valuesBuffer.reset();
+ readKeyStringValueIntoAccessors(
+ _nextRecord->keyString, *_ordering, &_valuesBuffer, &_accessors, _indexKeysToInclude);
+ }
+
if (_tracker && _tracker->trackProgress<TrialRunProgressTracker::kNumReads>(1)) {
// If we're collecting execution stats during multi-planning and reached the end of the
// trial period (trackProgress() will return 'true' in this case), then we can reset the
@@ -308,14 +318,18 @@ std::vector<DebugPrinter::Block> IndexScanStage::debugPrint() const {
}
ret.emplace_back(DebugPrinter::Block("[`"));
- for (size_t idx = 0; idx < _fields.size(); ++idx) {
- if (idx) {
+ size_t varIndex = 0;
+ for (size_t keyIndex = 0; keyIndex < _indexKeysToInclude.size(); ++keyIndex) {
+ if (!_indexKeysToInclude[keyIndex]) {
+ continue;
+ }
+ if (varIndex) {
ret.emplace_back(DebugPrinter::Block("`,"));
}
-
- DebugPrinter::addIdentifier(ret, _vars[idx]);
+ invariant(varIndex < _vars.size());
+ DebugPrinter::addIdentifier(ret, _vars[varIndex++]);
ret.emplace_back("=");
- DebugPrinter::addIdentifier(ret, _fields[idx]);
+ ret.emplace_back(std::to_string(keyIndex));
}
ret.emplace_back(DebugPrinter::Block("`]"));
diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h
index a3d23b2bb82..fcecc127ddd 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.h
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.h
@@ -29,6 +29,7 @@
#pragma once
+#include "mongo/bson/ordering.h"
#include "mongo/db/db_raii.h"
#include "mongo/db/exec/sbe/stages/stages.h"
#include "mongo/db/exec/trial_run_progress_tracker.h"
@@ -36,6 +37,27 @@
#include "mongo/db/storage/sorted_data_interface.h"
namespace mongo::sbe {
+
+/**
+ * A stage that iterates the entries of a collection index, starting from a bound specified by the
+ * value in 'seekKeySlotLow' and ending (via IS_EOF) with the 'seekKeySlotHi' bound. (An unspecified
+ * 'seekKeySlowHi' scans to the end of the index. Leaving both bounds unspecified scans the index
+ * from beginning to end.)
+ *
+ * The input 'seekKeySlotLow' and 'seekKeySlotHi' slots get read as part of the open (or re-open)
+ * call. A common use case for an IndexScanStage is to place it as the inner child of LoopJoinStage.
+ * The outer side of the LoopJoinStage determines the bounds, and the inner IndexScanStage iterates
+ * through all the entries within those bounds.
+ *
+ * The "output" slots are
+ * - 'recordSlot': the "KeyString" representing the index entry,
+ * - 'recordIdSlot': a reference that can be used to fetch the entire document, and
+ * - 'vars': one slot for each value in the index key that should be "projected" out of the entry.
+ *
+ * The 'indexKeysToInclude' bitset determines which values are included in the projection based
+ * on their order in the index pattern. The number of bits set in 'indexKeysToInclude' must be
+ * the same as the number of slots in the 'vars' SlotVector.
+ */
class IndexScanStage final : public PlanStage {
public:
IndexScanStage(const NamespaceStringOrUUID& name,
@@ -43,7 +65,7 @@ public:
bool forward,
boost::optional<value::SlotId> recordSlot,
boost::optional<value::SlotId> recordIdSlot,
- std::vector<std::string> fields,
+ IndexKeysInclusionSet indexKeysToInclude,
value::SlotVector vars,
boost::optional<value::SlotId> seekKeySlotLow,
boost::optional<value::SlotId> seekKeySlotHi,
@@ -74,7 +96,7 @@ private:
const bool _forward;
const boost::optional<value::SlotId> _recordSlot;
const boost::optional<value::SlotId> _recordIdSlot;
- const std::vector<std::string> _fields;
+ const IndexKeysInclusionSet _indexKeysToInclude;
const value::SlotVector _vars;
const boost::optional<value::SlotId> _seekKeySlotLow;
const boost::optional<value::SlotId> _seekKeySlotHi;
@@ -82,8 +104,10 @@ private:
std::unique_ptr<value::ViewOfValueAccessor> _recordAccessor;
std::unique_ptr<value::ViewOfValueAccessor> _recordIdAccessor;
- value::FieldAccessorMap _fieldAccessors;
- value::SlotAccessorMap _varAccessors;
+ // One accessor and slot for each key component that this stage will bind from an index entry's
+ // KeyString. The accessors are in the same order as the key components they bind to.
+ std::vector<value::ViewOfValueAccessor> _accessors;
+ value::SlotAccessorMap _accessorMap;
value::SlotAccessor* _seekKeyLowAccessor{nullptr};
value::SlotAccessor* _seekKeyHiAccessor{nullptr};
@@ -94,9 +118,14 @@ private:
std::unique_ptr<SortedDataInterface::Cursor> _cursor;
std::weak_ptr<const IndexCatalogEntry> _weakIndexCatalogEntry;
+ boost::optional<Ordering> _ordering{boost::none};
boost::optional<AutoGetCollectionForRead> _coll;
boost::optional<KeyStringEntry> _nextRecord;
+ // This buffer stores values that are projected out of the index entry. Values in the
+ // '_accessors' list that are pointers point to data in this buffer.
+ BufBuilder _valuesBuffer;
+
bool _open{false};
bool _firstGetNext{true};
IndexScanStats _specificStats;
diff --git a/src/mongo/db/exec/sbe/values/id_generators.h b/src/mongo/db/exec/sbe/values/id_generators.h
index fc7fa93f840..4b3fae32cca 100644
--- a/src/mongo/db/exec/sbe/values/id_generators.h
+++ b/src/mongo/db/exec/sbe/values/id_generators.h
@@ -67,6 +67,14 @@ public:
T generate() {
return this->generateByIncrementing();
}
+
+ auto generateMultiple(size_t num) {
+ std::vector<T> idVector(num);
+ for (T& id : idVector) {
+ id = generate();
+ }
+ return idVector;
+ }
};
using SlotIdGenerator = IdGenerator<value::SlotId>;