SERVER-63317 Basic plan generation for columnar indexes

11 files changed, 126 insertions, 4 deletions

diff --git a/jstests/noPassthroughWithMongod/column_index_skeleton.js b/jstests/noPassthroughWithMongod/column_index_skeleton.js
new file mode 100644
index 00000000000..bb938fef8ec
--- /dev/null
+++ b/jstests/noPassthroughWithMongod/column_index_skeleton.js
@@ -0,0 +1,36 @@
+/**
+ * Testing of just the query layer's integration for columnar index.
+ * This test is intended to be temporary.
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/fail_point_util.js");
+load("jstests/libs/analyze_plan.js");
+load("jstests/libs/sbe_util.js");  // For checkSBEEnabled.
+
+const isSBEEnabled = checkSBEEnabled(db);
+
+if (!isSBEEnabled) {
+    // This test is only relevant when SBE is enabled.
+    return;
+}
+
+const testDB = db;
+const coll = db.column_index_skeleton;
+
+assert.commandWorked(coll.insert({a: 1}));
+
+// Enable the columnar fail point.
+const failPoint = configureFailPoint(testDB, "includeFakeColumnarIndex");
+try {
+    // Run an explain.
+    const expl = coll.find({}, {a: 1}).explain();
+    assert(planHasStage(db, expl, "COLUMN_IXSCAN"));
+
+    // Run a query (should return nothing).
+    assert.eq(coll.find({}, {a: 1}).itcount(), 0);
+} finally {
+    failPoint.off();
+}
+})();
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.
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 7f148d2b961..d81f329e715 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -105,6 +105,7 @@
 #include "mongo/util/str.h"
 
 namespace mongo {
+MONGO_FAIL_POINT_DEFINE(includeFakeColumnarIndex);
 
 boost::intrusive_ptr<ExpressionContext> makeExpressionContextForGetExecutor(
     OperationContext* opCtx, const BSONObj& requestCollation, const NamespaceString& nss) {
@@ -307,6 +308,10 @@ void fillOutPlannerParams(OperationContext* opCtx,
     // If it's not NULL, we may have indices. Access the catalog and fill out IndexEntry(s)
     fillOutIndexEntries(opCtx, apiStrict, canonicalQuery, collection, plannerParams->indices);
 
+    if (includeFakeColumnarIndex.shouldFail()) {
+        plannerParams->columnarIndexes.push_back(ColumnIndexEntry{"fakeColumnIndex"});
+    }
+
     // If query supports index filters, filter params.indices by indices in query settings.
     // Ignore index filters when it is possible to use the id-hack.
     applyIndexFilters(collection, *canonicalQuery, plannerParams);
diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp
index 26a605c081f..ad91bd0c323 100644
--- a/src/mongo/db/query/query_planner.cpp
+++ b/src/mongo/db/query/query_planner.cpp
@@ -1196,7 +1196,8 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan(
     // Check whether we're eligible to use the columnar index, assuming no other indexes can be
     // used.
     if (out.empty() && !params.columnarIndexes.empty() && query.getProj() &&
-        !query.getProj()->requiresDocument()) {
+        !query.getProj()->requiresDocument() && query.isSbeCompatible() &&
+        !query.getForceClassicEngine()) {
         // TODO SERVER-63123: Check if the columnar index actually provides the fields we need.
         auto columnScan = std::make_unique<ColumnIndexScanNode>(params.columnarIndexes.front());
         columnScan->fields = query.getProj()->getRequiredFields();
diff --git a/src/mongo/db/query/query_planner_columnar_test.cpp b/src/mongo/db/query/query_planner_columnar_test.cpp
index dbca638506c..16a9f9c1854 100644
--- a/src/mongo/db/query/query_planner_columnar_test.cpp
+++ b/src/mongo/db/query/query_planner_columnar_test.cpp
@@ -45,6 +45,9 @@ protected:
     void setUp() final {
         QueryPlannerTest::setUp();
 
+        // Treat all queries as SBE compatible for this test.
+        QueryPlannerTest::setMarkQueriesSbeCompatible(true);
+
         // We're interested in testing plans that use a columnar index, so don't generate collection
         // scans.
         params.options &= ~QueryPlannerParams::INCLUDE_COLLSCAN;
diff --git a/src/mongo/db/query/query_planner_test_fixture.cpp b/src/mongo/db/query/query_planner_test_fixture.cpp
index aed4d8a7dd8..867ac63cdfb 100644
--- a/src/mongo/db/query/query_planner_test_fixture.cpp
+++ b/src/mongo/db/query/query_planner_test_fixture.cpp
@@ -363,6 +363,7 @@ void QueryPlannerTest::runQueryFull(
                                      std::move(pipeline));
     ASSERT_OK(statusWithCQ.getStatus());
     cq = std::move(statusWithCQ.getValue());
+    cq->setSbeCompatible(markQueriesSbeCompatible);
 
     auto statusWithMultiPlanSolns = QueryPlanner::plan(*cq, params);
     ASSERT_OK(statusWithMultiPlanSolns.getStatus());
@@ -439,6 +440,7 @@ void QueryPlannerTest::runInvalidQueryFull(const BSONObj& query,
                                      MatchExpressionParser::kAllowAllSpecialFeatures);
     ASSERT_OK(statusWithCQ.getStatus());
     cq = std::move(statusWithCQ.getValue());
+    cq->setSbeCompatible(markQueriesSbeCompatible);
 
     auto statusWithMultiPlanSolns = QueryPlanner::plan(*cq, params);
     plannerStatus = statusWithMultiPlanSolns.getStatus();
@@ -467,6 +469,7 @@ void QueryPlannerTest::runQueryAsCommand(const BSONObj& cmdObj) {
                                      MatchExpressionParser::kAllowAllSpecialFeatures);
     ASSERT_OK(statusWithCQ.getStatus());
     cq = std::move(statusWithCQ.getValue());
+    cq->setSbeCompatible(markQueriesSbeCompatible);
 
     auto statusWithMultiPlanSolns = QueryPlanner::plan(*cq, params);
     ASSERT_OK(statusWithMultiPlanSolns.getStatus());
@@ -494,6 +497,7 @@ void QueryPlannerTest::runInvalidQueryAsCommand(const BSONObj& cmdObj) {
                                      MatchExpressionParser::kAllowAllSpecialFeatures);
     ASSERT_OK(statusWithCQ.getStatus());
     cq = std::move(statusWithCQ.getValue());
+    cq->setSbeCompatible(markQueriesSbeCompatible);
 
     auto statusWithMultiPlanSolns = QueryPlanner::plan(*cq, params);
     plannerStatus = statusWithMultiPlanSolns.getStatus();
diff --git a/src/mongo/db/query/query_planner_test_fixture.h b/src/mongo/db/query/query_planner_test_fixture.h
index a34589548a3..02232623eef 100644
--- a/src/mongo/db/query/query_planner_test_fixture.h
+++ b/src/mongo/db/query/query_planner_test_fixture.h
@@ -242,6 +242,10 @@ protected:
     std::unique_ptr<MatchExpression> parseMatchExpression(
         const BSONObj& obj, const boost::intrusive_ptr<ExpressionContext>& expCtx = nullptr);
 
+    void setMarkQueriesSbeCompatible(bool sbeCompatible) {
+        markQueriesSbeCompatible = sbeCompatible;
+    }
+
     //
     // Data members.
     //
@@ -259,6 +263,7 @@ protected:
     std::vector<std::unique_ptr<QuerySolution>> solns;
 
     bool relaxBoundsCheck = false;
+    bool markQueriesSbeCompatible = false;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index cc64d5cea97..c82baaa1235 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -37,6 +37,7 @@
 
 #include "mongo/db/catalog/collection.h"
 #include "mongo/db/exec/sbe/stages/co_scan.h"
+#include "mongo/db/exec/sbe/stages/column_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/hash_join.h"
@@ -869,6 +870,37 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     return {std::move(stage), std::move(outputs)};
 }
 
+std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder::buildColumnScan(
+    const QuerySolutionNode* root, const PlanStageReqs& reqs) {
+    invariant(!reqs.getIndexKeyBitset());
+
+    auto csn = static_cast<const ColumnIndexScanNode*>(root);
+
+    PlanStageSlots outputs;
+
+    auto recordSlot = _slotIdGenerator.generate();
+    outputs.set(kResult, recordSlot);
+
+    boost::optional<sbe::value::SlotId> ridSlot;
+
+    if (reqs.has(kRecordId)) {
+        ridSlot = _slotIdGenerator.generate();
+        outputs.set(kRecordId, *ridSlot);
+    }
+
+    auto fieldSlotIds = _slotIdGenerator.generateMultiple(csn->fields.size());
+    auto stage = std::make_unique<sbe::ColumnScanStage>(_collection->uuid(),
+                                                        csn->indexEntry.catalogName,
+                                                        fieldSlotIds,
+                                                        csn->fields,
+                                                        recordSlot,
+                                                        ridSlot,
+                                                        _yieldPolicy,
+                                                        csn->nodeId());
+
+    return {std::move(stage), std::move(outputs)};
+}
+
 std::tuple<sbe::value::SlotId, sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>>
 SlotBasedStageBuilder::makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage,
                                             sbe::value::SlotId seekKeySlot,
@@ -2941,6 +2973,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
             {STAGE_COLLSCAN, &SlotBasedStageBuilder::buildCollScan},
             {STAGE_VIRTUAL_SCAN, &SlotBasedStageBuilder::buildVirtualScan},
             {STAGE_IXSCAN, &SlotBasedStageBuilder::buildIndexScan},
+            {STAGE_COLUMN_IXSCAN, &SlotBasedStageBuilder::buildColumnScan},
             {STAGE_FETCH, &SlotBasedStageBuilder::buildFetch},
             {STAGE_LIMIT, &SlotBasedStageBuilder::buildLimit},
             {STAGE_SKIP, &SlotBasedStageBuilder::buildSkip},
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index 6ff9bf615fb..9efd2703504 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -349,6 +349,9 @@ private:
     std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> buildIndexScan(
         const QuerySolutionNode* root, const PlanStageReqs& reqs);
 
+    std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> buildColumnScan(
+        const QuerySolutionNode* root, const PlanStageReqs& reqs);
+
     std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> buildFetch(
         const QuerySolutionNode* root, const PlanStageReqs& reqs);