SERVER-70639 [CQF] Implement Spool physical node

3 files changed, 239 insertions, 4 deletions

diff --git a/src/mongo/db/exec/sbe/abt/abt_lower.cpp b/src/mongo/db/exec/sbe/abt/abt_lower.cpp
index 18348ab9ac8..8d47db3830d 100644
--- a/src/mongo/db/exec/sbe/abt/abt_lower.cpp
+++ b/src/mongo/db/exec/sbe/abt/abt_lower.cpp
@@ -28,7 +28,6 @@
  */
 
 #include "mongo/db/exec/sbe/abt/abt_lower.h"
-#include "mongo/db/exec/sbe/stages/bson_scan.h"
 #include "mongo/db/exec/sbe/stages/co_scan.h"
 #include "mongo/db/exec/sbe/stages/exchange.h"
 #include "mongo/db/exec/sbe/stages/filter.h"
@@ -42,11 +41,13 @@
 #include "mongo/db/exec/sbe/stages/scan.h"
 #include "mongo/db/exec/sbe/stages/sort.h"
 #include "mongo/db/exec/sbe/stages/sorted_merge.h"
+#include "mongo/db/exec/sbe/stages/spool.h"
 #include "mongo/db/exec/sbe/stages/union.h"
 #include "mongo/db/exec/sbe/stages/unique.h"
 #include "mongo/db/exec/sbe/stages/unwind.h"
 #include "mongo/db/query/optimizer/utils/utils.h"
 
+
 namespace mongo::optimizer {
 
 static sbe::EExpression::Vector toInlinedVector(
@@ -576,6 +577,61 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const UniqueNode& n,
     return sbe::makeS<sbe::UniqueStage>(std::move(input), std::move(keySlots), planNodeId);
 }
 
+std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const SpoolProducerNode& n,
+                                                      const ABT& child,
+                                                      const ABT& filter,
+                                                      const ABT& binder,
+                                                      const ABT& refs) {
+    auto input = generateInternal(child);
+
+    sbe::value::SlotVector vals;
+    for (const ProjectionName& projectionName : n.binder().names()) {
+        auto it = _slotMap.find(projectionName);
+        uassert(6624139,
+                str::stream() << "undefined variable: " << projectionName,
+                it != _slotMap.end());
+        vals.push_back(it->second);
+    }
+
+    const PlanNodeId planNodeId = _nodeToGroupPropsMap.at(&n)._planNodeId;
+    switch (n.getType()) {
+        case SpoolProducerType::Eager:
+            return sbe::makeS<sbe::SpoolEagerProducerStage>(
+                std::move(input), n.getSpoolId(), std::move(vals), planNodeId);
+
+        case SpoolProducerType::Lazy: {
+            auto expr = SBEExpressionLowering{_env, _slotMap}.optimize(filter);
+            return sbe::makeS<sbe::SpoolLazyProducerStage>(
+                std::move(input), n.getSpoolId(), std::move(vals), std::move(expr), planNodeId);
+        }
+    }
+
+    MONGO_UNREACHABLE;
+}
+
+std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const SpoolConsumerNode& n,
+                                                      const ABT& binder) {
+    sbe::value::SlotVector vals;
+    for (const ProjectionName& projectionName : n.binder().names()) {
+        auto slot = _slotIdGenerator.generate();
+        _slotMap.emplace(projectionName, slot);
+        vals.push_back(slot);
+    }
+
+    const PlanNodeId planNodeId = _nodeToGroupPropsMap.at(&n)._planNodeId;
+    switch (n.getType()) {
+        case SpoolConsumerType::Stack:
+            return sbe::makeS<sbe::SpoolConsumerStage<true /*isStack*/>>(
+                n.getSpoolId(), std::move(vals), planNodeId);
+
+        case SpoolConsumerType::Regular:
+            return sbe::makeS<sbe::SpoolConsumerStage<false /*isStack*/>>(
+                n.getSpoolId(), std::move(vals), planNodeId);
+    }
+
+    MONGO_UNREACHABLE;
+}
+
 std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const GroupByNode& n,
                                                       const ABT& child,
                                                       const ABT& aggBinds,
diff --git a/src/mongo/db/exec/sbe/abt/abt_lower.h b/src/mongo/db/exec/sbe/abt/abt_lower.h
index bdab4602d4a..ac4c7f49ccf 100644
--- a/src/mongo/db/exec/sbe/abt/abt_lower.h
+++ b/src/mongo/db/exec/sbe/abt/abt_lower.h
@@ -108,9 +108,11 @@ public:
     // The default noop transport.
     template <typename T, typename... Ts>
     std::unique_ptr<sbe::PlanStage> walk(const T&, Ts&&...) {
-        if constexpr (std::is_base_of_v<ExclusivelyLogicalNode, T>) {
-            uasserted(6624238, "A physical plan should not contain exclusively logical nodes.");
-        }
+        // We should not be seeing a physical delegator node here.
+        static_assert(!canBePhysicalNode<T>() || std::is_same_v<MemoPhysicalDelegatorNode, T>,
+                      "Physical nodes need to implement lowering");
+
+        uasserted(6624238, "Unexpected node type.");
         return nullptr;
     }
 
@@ -127,6 +129,13 @@ public:
 
     std::unique_ptr<sbe::PlanStage> walk(const UniqueNode& n, const ABT& child, const ABT& refs);
 
+    std::unique_ptr<sbe::PlanStage> walk(const SpoolProducerNode& n,
+                                         const ABT& child,
+                                         const ABT& filter,
+                                         const ABT& binder,
+                                         const ABT& refs);
+    std::unique_ptr<sbe::PlanStage> walk(const SpoolConsumerNode& n, const ABT& binder);
+
     std::unique_ptr<sbe::PlanStage> walk(const GroupByNode& n,
                                          const ABT& child,
                                          const ABT& aggBinds,
diff --git a/src/mongo/db/exec/sbe/abt/sbe_abt_test.cpp b/src/mongo/db/exec/sbe/abt/sbe_abt_test.cpp
index 25a42fefcee..8758ae83786 100644
--- a/src/mongo/db/exec/sbe/abt/sbe_abt_test.cpp
+++ b/src/mongo/db/exec/sbe/abt/sbe_abt_test.cpp
@@ -39,6 +39,7 @@
 #include "mongo/db/query/optimizer/utils/unit_test_utils.h"
 #include "mongo/unittest/unittest.h"
 
+
 namespace mongo::optimizer {
 namespace {
 
@@ -677,5 +678,174 @@ TEST_F(NodeSBE, RequireRID) {
     ASSERT_EQ(1, resultSize);
 }
 
+/**
+ * This transport is used to populate default values into the NodeToGroupProps map to get around the
+ * fact that the plan was not obtained from the memo. At this point we are interested only in the
+ * planNodeIds being distinct.
+ */
+class PropsTransport {
+public:
+    template <typename T, typename... Ts>
+    void transport(const T& node, NodeToGroupPropsMap& propMap, Ts&&...) {
+        if constexpr (std::is_base_of_v<Node, T>) {
+            propMap.emplace(&node,
+                            NodeProps{_planNodeId++,
+                                      {-1, 0} /*groupId*/,
+                                      {} /*logicalProps*/,
+                                      {} /*physicalProps*/,
+                                      boost::none /*ridProjName*/,
+                                      CostType::kZero /*cost*/,
+                                      CostType::kZero /*localCost*/,
+                                      0.0 /*adjustedCE*/});
+        }
+    }
+
+    void updatePropsMap(const ABT& n, NodeToGroupPropsMap& propMap) {
+        algebra::transport<false>(n, *this, propMap);
+    }
+
+private:
+    int32_t _planNodeId = 0;
+};
+
+TEST_F(NodeSBE, SpoolFibonacci) {
+    using namespace unit_test_abt_literals;
+
+    PrefixId prefixId;
+    Metadata metadata{{}};
+
+    // Construct a spool-based recursive plan to compute the first 10 Fibonacci numbers. The main
+    // plan (first child of the union) sets up the initial conditions (val = 1, prev = 0, and it =
+    // 1), and the recursive subplan is computing the actual Fibonacci sequence and ensures we
+    // terminate after 10 numbers.
+    auto recursion =
+        NodeBuilder{}
+            .eval("val", _binary("Add", "valIn"_var, "valIn_prev"_var))
+            .eval("val_prev", "valIn"_var)
+            .eval("it", _binary("Add", "itIn"_var, "1"_cint64))
+            .filter(_binary("Lt", "itIn"_var, "10"_cint64))
+            .finish(_spoolc("Stack", 1 /*spoolId*/, _varnames("valIn", "valIn_prev", "itIn")));
+
+    auto tree = NodeBuilder{}
+                    .root("val")
+                    .spoolp("Lazy", 1 /*spoolId*/, _varnames("val", "val_prev", "it"), _cbool(true))
+                    .un(_varnames("val", "val_prev", "it"), {NodeHolder{std::move(recursion)}})
+                    .eval("val", "1"_cint64)
+                    .eval("val_prev", "0"_cint64)
+                    .eval("it", "1"_cint64)
+                    .ls(1, 0)
+                    .finish(_coscan());
+
+    ASSERT_EXPLAIN_V2_AUTO(
+        "Root []\n"
+        "|   |   projections: \n"
+        "|   |       val\n"
+        "|   RefBlock: \n"
+        "|       Variable [val]\n"
+        "SpoolProducer [Lazy, id: 1]\n"
+        "|   |   Const [true]\n"
+        "|   BindBlock:\n"
+        "|       [it]\n"
+        "|           Source []\n"
+        "|       [val]\n"
+        "|           Source []\n"
+        "|       [val_prev]\n"
+        "|           Source []\n"
+        "Union []\n"
+        "|   |   BindBlock:\n"
+        "|   |       [it]\n"
+        "|   |           Source []\n"
+        "|   |       [val]\n"
+        "|   |           Source []\n"
+        "|   |       [val_prev]\n"
+        "|   |           Source []\n"
+        "|   Evaluation []\n"
+        "|   |   BindBlock:\n"
+        "|   |       [val]\n"
+        "|   |           BinaryOp [Add]\n"
+        "|   |           |   Variable [valIn_prev]\n"
+        "|   |           Variable [valIn]\n"
+        "|   Evaluation []\n"
+        "|   |   BindBlock:\n"
+        "|   |       [val_prev]\n"
+        "|   |           Variable [valIn]\n"
+        "|   Evaluation []\n"
+        "|   |   BindBlock:\n"
+        "|   |       [it]\n"
+        "|   |           BinaryOp [Add]\n"
+        "|   |           |   Const [1]\n"
+        "|   |           Variable [itIn]\n"
+        "|   Filter []\n"
+        "|   |   BinaryOp [Lt]\n"
+        "|   |   |   Const [10]\n"
+        "|   |   Variable [itIn]\n"
+        "|   SpoolConsumer [Stack, id: 1]\n"
+        "|       BindBlock:\n"
+        "|           [itIn]\n"
+        "|               Source []\n"
+        "|           [valIn]\n"
+        "|               Source []\n"
+        "|           [valIn_prev]\n"
+        "|               Source []\n"
+        "Evaluation []\n"
+        "|   BindBlock:\n"
+        "|       [val]\n"
+        "|           Const [1]\n"
+        "Evaluation []\n"
+        "|   BindBlock:\n"
+        "|       [val_prev]\n"
+        "|           Const [0]\n"
+        "Evaluation []\n"
+        "|   BindBlock:\n"
+        "|       [it]\n"
+        "|           Const [1]\n"
+        "LimitSkip []\n"
+        "|   limitSkip:\n"
+        "|       limit: 1\n"
+        "|       skip: 0\n"
+        "CoScan []\n",
+        tree);
+
+    NodeToGroupPropsMap props;
+    PropsTransport{}.updatePropsMap(tree, props);
+
+    auto env = VariableEnvironment::build(tree);
+    SlotVarMap map;
+    boost::optional<sbe::value::SlotId> ridSlot;
+    sbe::value::SlotIdGenerator ids;
+    SBENodeLowering g{env, map, ridSlot, ids, metadata, props, false /*randomScan*/};
+    auto sbePlan = g.optimize(tree);
+    ASSERT_EQ(1, map.size());
+
+    auto opCtx = makeOperationContext();
+    sbe::CompileCtx ctx(std::make_unique<sbe::RuntimeEnvironment>());
+    sbePlan->prepare(ctx);
+
+    std::vector<sbe::value::SlotAccessor*> accessors;
+    for (auto& [name, slot] : map) {
+        accessors.emplace_back(sbePlan->getAccessor(ctx, slot));
+    }
+
+    sbePlan->attachToOperationContext(opCtx.get());
+    sbePlan->open(false);
+
+    std::vector<int64_t> results;
+    while (sbePlan->getNext() != sbe::PlanState::IS_EOF) {
+        const auto [resultTag, resultVal] = accessors.front()->getViewOfValue();
+        ASSERT_EQ(sbe::value::TypeTags::NumberInt64, resultTag);
+        results.push_back(resultVal);
+    };
+    sbePlan->close();
+
+    // Verify we are getting 10 Fibonacci numbers.
+    ASSERT_EQ(10, results.size());
+
+    ASSERT_EQ(1, results.at(0));
+    ASSERT_EQ(1, results.at(1));
+    for (size_t i = 2; i < 10; i++) {
+        ASSERT_EQ(results.at(i), results.at(i - 1) + results.at(i - 2));
+    }
+}
+
 }  // namespace
 }  // namespace mongo::optimizer