SERVER-63317 Basic plan generation for columnar indexes

3 files changed, 35 insertions, 3 deletions

diff --git a/src/mongo/db/exec/sbe/parser/parser.cpp b/src/mongo/db/exec/sbe/parser/parser.cpp
index e86549ba938..d6e89450d05 100644
--- a/src/mongo/db/exec/sbe/parser/parser.cpp
+++ b/src/mongo/db/exec/sbe/parser/parser.cpp
@@ -130,6 +130,7 @@ static constexpr auto kSyntax = R"(
                 COLUMNSCAN <- 'columnscan' STRING_LIST # paths
                                            IDENT_LIST # output variables
                                            IDENT # output record
+                                           IDENT # output RID
                                            IDENT # collection name
                                            IDENT # index name to scan
 
@@ -956,14 +957,16 @@ void Parser::walkColumnScan(AstQuery& ast) {
     auto paths = ast.nodes[0]->identifiers;
     auto outputs = lookupSlots(ast.nodes[1]->identifiers);
     auto record = lookupSlot(ast.nodes[2]->identifier);
-    auto collName = ast.nodes[3]->identifier;
-    auto indexName = ast.nodes[4]->identifier;
+    auto recordId = lookupSlot(ast.nodes[3]->identifier);
+    auto collName = ast.nodes[4]->identifier;
+    auto indexName = ast.nodes[5]->identifier;
 
     ast.stage = makeS<ColumnScanStage>(getCollectionUuid(collName),
                                        indexName,
                                        std::move(outputs),
                                        std::move(paths),
                                        record,
+                                       recordId,
                                        nullptr,
                                        getCurrentPlanNodeId());
 }
diff --git a/src/mongo/db/exec/sbe/stages/column_scan.cpp b/src/mongo/db/exec/sbe/stages/column_scan.cpp
index f3553f1c144..90928bce6d4 100644
--- a/src/mongo/db/exec/sbe/stages/column_scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/column_scan.cpp
@@ -40,6 +40,7 @@ ColumnScanStage::ColumnScanStage(UUID collectionUuid,
                                  value::SlotVector fieldSlots,
                                  std::vector<std::string> paths,
                                  boost::optional<value::SlotId> recordSlot,
+                                 boost::optional<value::SlotId> recordIdSlot,
                                  PlanYieldPolicy* yieldPolicy,
                                  PlanNodeId nodeId)
     : PlanStage("columnscan"_sd, yieldPolicy, nodeId),
@@ -47,7 +48,8 @@ ColumnScanStage::ColumnScanStage(UUID collectionUuid,
       _columnIndexName(columnIndexName),
       _fieldSlots(std::move(fieldSlots)),
       _paths(std::move(paths)),
-      _recordSlot(recordSlot) {
+      _recordSlot(recordSlot),
+      _recordIdSlot(recordIdSlot) {
     invariant(_fieldSlots.size() == _paths.size());
 }
 
@@ -57,6 +59,7 @@ std::unique_ptr<PlanStage> ColumnScanStage::clone() const {
                                              _fieldSlots,
                                              _paths,
                                              _recordSlot,
+                                             _recordIdSlot,
                                              _yieldPolicy,
                                              _commonStats.nodeId);
 }
@@ -72,12 +75,23 @@ void ColumnScanStage::prepare(CompileCtx& ctx) {
     if (_recordSlot) {
         _recordAccessor = std::make_unique<value::OwnedValueAccessor>();
     }
+    if (_recordIdSlot) {
+        _recordIdAccessor = std::make_unique<value::OwnedValueAccessor>();
+    }
 
     tassert(6298602, "'_coll' should not be initialized prior to 'acquireCollection()'", !_coll);
     std::tie(_coll, _collName, _catalogEpoch) = acquireCollection(_opCtx, _collUuid);
 }
 
 value::SlotAccessor* ColumnScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot) {
+    if (_recordSlot && slot == *_recordSlot) {
+        return _recordAccessor.get();
+    }
+
+    if (_recordIdSlot && slot == *_recordIdSlot) {
+        return _recordIdAccessor.get();
+    }
+
     if (auto it = _outputFieldsMap.find(slot); it != _outputFieldsMap.end()) {
         return it->second;
     }
@@ -230,6 +244,18 @@ std::vector<DebugPrinter::Block> ColumnScanStage::debugPrint() const {
     }
     ret.emplace_back(DebugPrinter::Block("`]"));
 
+    if (_recordSlot) {
+        DebugPrinter::addIdentifier(ret, _recordSlot.get());
+    } else {
+        DebugPrinter::addIdentifier(ret, DebugPrinter::kNoneKeyword);
+    }
+
+    if (_recordIdSlot) {
+        DebugPrinter::addIdentifier(ret, _recordIdSlot.get());
+    } else {
+        DebugPrinter::addIdentifier(ret, DebugPrinter::kNoneKeyword);
+    }
+
     // Print out paths.
     ret.emplace_back(DebugPrinter::Block("[`"));
     for (size_t idx = 0; idx < _paths.size(); ++idx) {
diff --git a/src/mongo/db/exec/sbe/stages/column_scan.h b/src/mongo/db/exec/sbe/stages/column_scan.h
index dcdcb233c59..da62f9290b6 100644
--- a/src/mongo/db/exec/sbe/stages/column_scan.h
+++ b/src/mongo/db/exec/sbe/stages/column_scan.h
@@ -46,6 +46,7 @@ public:
                     value::SlotVector fieldSlots,
                     std::vector<std::string> paths,
                     boost::optional<value::SlotId> recordSlot,
+                    boost::optional<value::SlotId> recordIdSlot,
                     PlanYieldPolicy* yieldPolicy,
                     PlanNodeId nodeId);
 
@@ -77,10 +78,12 @@ private:
     const value::SlotVector _fieldSlots;
     const std::vector<std::string> _paths;
     const boost::optional<value::SlotId> _recordSlot;
+    const boost::optional<value::SlotId> _recordIdSlot;
 
     std::vector<value::OwnedValueAccessor> _outputFields;
     value::SlotAccessorMap _outputFieldsMap;
     std::unique_ptr<value::OwnedValueAccessor> _recordAccessor;
+    std::unique_ptr<value::OwnedValueAccessor> _recordIdAccessor;
 
     // These members are default constructed to boost::none and are initialized when 'prepare()'
     // is called. Once they are set, they are never modified again.