SERVER-63317 Basic plan generation for columnar indexes

7 files changed, 55 insertions, 1 deletions

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);