SERVER-73724 [CQF] Support for optional v1 and v2 explain for js tests

11 files changed, 205 insertions, 26 deletions

diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript
index a9b3dcc6d2b..e87d9e4b72f 100644
--- a/src/mongo/db/query/SConscript
+++ b/src/mongo/db/query/SConscript
@@ -256,6 +256,7 @@ env.Library(
     source=[
         'ce_mode_parameter.cpp',
         'cost_model/cost_model_on_update.cpp',
+        'explain_version_validator.cpp',
         'framework_control.cpp',
         'query_feature_flags.idl',
         'query_knobs.idl',
diff --git a/src/mongo/db/query/ce/test_utils.cpp b/src/mongo/db/query/ce/test_utils.cpp
index 32ae7f2c06c..cd1e9939a54 100644
--- a/src/mongo/db/query/ce/test_utils.cpp
+++ b/src/mongo/db/query/ce/test_utils.cpp
@@ -86,6 +86,7 @@ CEType CETester::getCE(ABT& abt, std::function<bool(const ABT&)> nodePredicate)
                                  makeCostEstimator(),
                                  defaultConvertPathToInterval,
                                  ConstEval::constFold,
+                                 true /*supportExplain*/,
                                  DebugInfo::kDefaultForTests,
                                  _hints};
     optimize(phaseManager, abt);
diff --git a/src/mongo/db/query/cqf_get_executor.cpp b/src/mongo/db/query/cqf_get_executor.cpp
index ab0c85e0aa1..16bc95963de 100644
--- a/src/mongo/db/query/cqf_get_executor.cpp
+++ b/src/mongo/db/query/cqf_get_executor.cpp
@@ -44,6 +44,7 @@
 #include "mongo/db/query/cost_model/cost_model_manager.h"
 #include "mongo/db/query/cost_model/on_coefficients_change_updater_impl.h"
 #include "mongo/db/query/cqf_command_utils.h"
+#include "mongo/db/query/explain_version_validator.h"
 #include "mongo/db/query/optimizer/explain.h"
 #include "mongo/db/query/optimizer/metadata_factory.h"
 #include "mongo/db/query/optimizer/node.h"
@@ -269,7 +270,8 @@ static ExecParams createExecutor(OptPhaseManager phaseManager,
                                  const NamespaceString& nss,
                                  const CollectionPtr& collection,
                                  const bool requireRID,
-                                 const ScanOrder scanOrder) {
+                                 const ScanOrder scanOrder,
+                                 const bool needsExplain) {
     auto env = VariableEnvironment::build(abt);
     SlotVarMap slotMap;
     auto runtimeEnvironment = std::make_unique<sbe::RuntimeEnvironment>();  // TODO use factory
@@ -299,7 +301,7 @@ static ExecParams createExecutor(OptPhaseManager phaseManager,
     }
 
     sbePlan->attachToOperationContext(opCtx);
-    if (expCtx->explain || expCtx->mayDbProfile) {
+    if (needsExplain || expCtx->mayDbProfile) {
         sbePlan->markShouldCollectTimingInfo();
     }
 
@@ -311,13 +313,41 @@ static ExecParams createExecutor(OptPhaseManager phaseManager,
                                              nullptr,
                                              std::make_unique<YieldPolicyCallbacksImpl>(nss));
 
+    std::unique_ptr<ABTPrinter> abtPrinter;
+    if (needsExplain) {
+        // By default, we print the optimized ABT. For test-only versions we output the post-memo
+        // phan instead.
+        ABT toExplain = std::move(abt);
+
+        ExplainVersion explainVersion = ExplainVersion::Vmax;
+        const auto& explainVersionStr = internalCascadesOptimizerExplainVersion.get();
+        if (explainVersionStr == "v1"_sd) {
+            explainVersion = ExplainVersion::V1;
+            toExplain = *phaseManager.getPostMemoPlan();
+        } else if (explainVersionStr == "v2"_sd) {
+            explainVersion = ExplainVersion::V2;
+            toExplain = *phaseManager.getPostMemoPlan();
+        } else if (explainVersionStr == "v2compact"_sd) {
+            explainVersion = ExplainVersion::V2Compact;
+            toExplain = *phaseManager.getPostMemoPlan();
+        } else if (explainVersionStr == "bson"_sd) {
+            explainVersion = ExplainVersion::V3;
+        } else {
+            // Should have been validated.
+            MONGO_UNREACHABLE;
+        }
+
+        abtPrinter = std::make_unique<ABTPrinter>(
+            std::move(toExplain), phaseManager.getNodeToGroupPropsMap(), explainVersion);
+    }
+
     sbePlan->prepare(data.ctx);
     CurOp::get(opCtx)->stopQueryPlanningTimer();
 
     return {opCtx,
             nullptr /*solution*/,
             {std::move(sbePlan), std::move(data)},
-            std::make_unique<ABTPrinter>(std::move(abt), phaseManager.getNodeToGroupPropsMap()),
+            std::move(abtPrinter),
             QueryPlannerParams::Options::DEFAULT,
             nss,
             std::move(yieldPolicy),
@@ -530,6 +560,7 @@ static OptPhaseManager createPhaseManager(const CEMode mode,
                                           const bool requireRID,
                                           Metadata metadata,
                                           const ConstFoldFn& constFold,
+                                          const bool supportExplain,
                                           QueryHints hints) {
     switch (mode) {
         case CEMode::kSampling: {
@@ -549,6 +580,7 @@ static OptPhaseManager createPhaseManager(const CEMode mode,
                                                     std::make_unique<CostEstimatorImpl>(costModel),
                                                     defaultConvertPathToInterval,
                                                     constFold,
+                                                    supportExplain,
                                                     DebugInfo::kDefaultForProd,
                                                     {} /*hints*/};
             return {OptPhaseManager::getAllRewritesSet(),
@@ -563,6 +595,7 @@ static OptPhaseManager createPhaseManager(const CEMode mode,
                     std::make_unique<CostEstimatorImpl>(costModel),
                     defaultConvertPathToInterval,
                     constFold,
+                    supportExplain,
                     DebugInfo::kDefaultForProd,
                     std::move(hints)};
         }
@@ -579,6 +612,7 @@ static OptPhaseManager createPhaseManager(const CEMode mode,
                     std::make_unique<CostEstimatorImpl>(costModel),
                     defaultConvertPathToInterval,
                     constFold,
+                    supportExplain,
                     DebugInfo::kDefaultForProd,
                     std::move(hints)};
 
@@ -592,6 +626,7 @@ static OptPhaseManager createPhaseManager(const CEMode mode,
                     std::make_unique<CostEstimatorImpl>(costModel),
                     defaultConvertPathToInterval,
                     constFold,
+                    supportExplain,
                     DebugInfo::kDefaultForProd,
                     std::move(hints)};
 
@@ -688,6 +723,7 @@ boost::optional<ExecParams> getSBEExecutorViaCascadesOptimizer(
     }
 
     auto costModel = cost_model::costModelManager(opCtx->getServiceContext()).getCoefficients();
+    const bool needsExplain = expCtx->explain.has_value();
 
     OptPhaseManager phaseManager = createPhaseManager(mode,
                                                       costModel,
@@ -698,6 +734,7 @@ boost::optional<ExecParams> getSBEExecutorViaCascadesOptimizer(
                                                       requireRID,
                                                       std::move(metadata),
                                                       constFold,
+                                                      needsExplain,
                                                       std::move(queryHints));
     if (!phaseManager.optimizeNoAssert(abt)) {
         return boost::none;
@@ -737,7 +774,8 @@ boost::optional<ExecParams> getSBEExecutorViaCascadesOptimizer(
                           nss,
                           collection,
                           requireRID,
-                          scanOrder);
+                          scanOrder,
+                          needsExplain);
 }
 
 boost::optional<ExecParams> getSBEExecutorViaCascadesOptimizer(const CollectionPtr& collection,
diff --git a/src/mongo/db/query/explain_version_validator.cpp b/src/mongo/db/query/explain_version_validator.cpp
new file mode 100644
index 00000000000..f23c04c7b7d
--- /dev/null
+++ b/src/mongo/db/query/explain_version_validator.cpp
@@ -0,0 +1,41 @@
+/**
+ *    Copyright (C) 2023-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/query/explain_version_validator.h"
+
+namespace mongo::optimizer {
+
+Status validateOptimizerExplainVersion(const std::string& value, const boost::optional<TenantId>&) {
+    if (value == "bson"_sd || value == "v1"_sd || value == "v2"_sd || value == "v2compact"_sd) {
+        return Status::OK();
+    }
+    return Status(ErrorCodes::Error{6624103}, "Invalid optimizer explain version.");
+}
+
+}  // namespace mongo::optimizer
diff --git a/src/mongo/db/query/explain_version_validator.h b/src/mongo/db/query/explain_version_validator.h
new file mode 100644
index 00000000000..fb56c0d87ba
--- /dev/null
+++ b/src/mongo/db/query/explain_version_validator.h
@@ -0,0 +1,40 @@
+/**
+ *    Copyright (C) 2023-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/base/status.h"
+#include "mongo/db/tenant_id.h"
+
+
+namespace mongo::optimizer {
+
+Status validateOptimizerExplainVersion(const std::string& value, const boost::optional<TenantId>&);
+
+}  // namespace mongo::optimizer
diff --git a/src/mongo/db/query/optimizer/explain.cpp b/src/mongo/db/query/optimizer/explain.cpp
index 96f39bd367e..a64c6a742d2 100644
--- a/src/mongo/db/query/optimizer/explain.cpp
+++ b/src/mongo/db/query/optimizer/explain.cpp
@@ -39,12 +39,44 @@
 
 namespace mongo::optimizer {
 
+ABTPrinter::ABTPrinter(ABT abt,
+                       NodeToGroupPropsMap nodeToPropsMap,
+                       const ExplainVersion explainVersion)
+    : _abt(std::move(abt)),
+      _nodeToPropsMap(std::move(nodeToPropsMap)),
+      _explainVersion(explainVersion) {}
+
 BSONObj ABTPrinter::explainBSON() const {
-    return ExplainGenerator::explainBSONObj(
-        _abtTree, true /*displayProperties*/, nullptr /*memoInterface*/, _nodeToPropsMap);
-}
+    const auto explainPlanStr = [&](std::string planStr) {
+        BSONObjBuilder builder;
+        builder.append("plan", std::move(planStr));
+        return builder.done().getOwned();
+    };
+
+    switch (_explainVersion) {
+        case ExplainVersion::V1:
+            return explainPlanStr(ExplainGenerator::explain(
+                _abt, false /*displayProperties*/, nullptr /*memoInterface*/, _nodeToPropsMap));
+
+        case ExplainVersion::V2:
+            return explainPlanStr(ExplainGenerator::explainV2(
+                _abt, false /*displayProperties*/, nullptr /*memoInterface*/, _nodeToPropsMap));
+
+        case ExplainVersion::V2Compact:
+            return explainPlanStr(ExplainGenerator::explainV2Compact(
+                _abt, false /*displayProperties*/, nullptr /*memoInterface*/, _nodeToPropsMap));
 
-enum class ExplainVersion { V1, V2, V2Compact, V3, Vmax };
+        case ExplainVersion::V3:
+            return ExplainGenerator::explainBSONObj(
+                _abt, true /*displayProperties*/, nullptr /*memoInterface*/, _nodeToPropsMap);
+
+        case ExplainVersion::Vmax:
+            // Should not be seeing this value here.
+            break;
+    }
+
+    MONGO_UNREACHABLE;
+}
 
 bool constexpr operator<(const ExplainVersion v1, const ExplainVersion v2) {
     return static_cast<int>(v1) < static_cast<int>(v2);
diff --git a/src/mongo/db/query/optimizer/explain.h b/src/mongo/db/query/optimizer/explain.h
index 2334270f888..ddec2f171dc 100644
--- a/src/mongo/db/query/optimizer/explain.h
+++ b/src/mongo/db/query/optimizer/explain.h
@@ -40,21 +40,22 @@
 
 namespace mongo::optimizer {
 
+enum class ExplainVersion { V1, V2, V2Compact, V3, Vmax };
+
 /**
- * This structure holds any data that is required by the BSON version of explain. It is
- * self-sufficient and separate because it must outlive the other optimizer state as it is used by
- * the runtime plan executor.
+ * This structure holds any data that is required by the explain. It is self-sufficient and separate
+ * because it must outlive the other optimizer state as it is used by the runtime plan executor.
  */
 class ABTPrinter : public AbstractABTPrinter {
 public:
-    ABTPrinter(ABT abtTree, NodeToGroupPropsMap nodeToPropsMap)
-        : _abtTree(std::move(abtTree)), _nodeToPropsMap(std::move(nodeToPropsMap)) {}
+    ABTPrinter(ABT abt, NodeToGroupPropsMap nodeToPropsMap, ExplainVersion explainVersion);
 
     BSONObj explainBSON() const override final;
 
 private:
-    ABT _abtTree;
+    ABT _abt;
     NodeToGroupPropsMap _nodeToPropsMap;
+    ExplainVersion _explainVersion;
 };
 
 class ExplainGenerator {
diff --git a/src/mongo/db/query/optimizer/opt_phase_manager.cpp b/src/mongo/db/query/optimizer/opt_phase_manager.cpp
index 9e2ef84ecf8..f51e94b5935 100644
--- a/src/mongo/db/query/optimizer/opt_phase_manager.cpp
+++ b/src/mongo/db/query/optimizer/opt_phase_manager.cpp
@@ -54,9 +54,11 @@ OptPhaseManager::OptPhaseManager(OptPhaseManager::PhaseSet phaseSet,
                                  std::unique_ptr<CostEstimator> costEstimator,
                                  PathToIntervalFn pathToInterval,
                                  ConstFoldFn constFold,
+                                 const bool supportExplain,
                                  DebugInfo debugInfo,
                                  QueryHints queryHints)
     : _phaseSet(std::move(phaseSet)),
+      _supportExplain(supportExplain),
       _debugInfo(std::move(debugInfo)),
       _hints(std::move(queryHints)),
       _metadata(std::move(metadata)),
@@ -68,6 +70,7 @@ OptPhaseManager::OptPhaseManager(OptPhaseManager::PhaseSet phaseSet,
       _pathToInterval(std::move(pathToInterval)),
       _constFold(std::move(constFold)),
       _physicalNodeId(),
+      _postMemoPlan(),
       _requireRID(requireRID),
       _ridProjections(),
       _prefixId(prefixId) {
@@ -238,6 +241,9 @@ bool OptPhaseManager::runMemoPhysicalRewrite(const OptPhase phase,
     _physicalNodeId = {rootGroupId, optGroupResult._index};
     std::tie(input, _nodeToGroupPropsMap) =
         extractPhysicalPlan(_physicalNodeId, _metadata, _ridProjections, _memo);
+    if (_supportExplain) {
+        _postMemoPlan = input;
+    }
 
     env.rebuild(input);
     if (env.hasFreeVariables()) {
@@ -336,6 +342,10 @@ MemoPhysicalNodeId OptPhaseManager::getPhysicalNodeId() const {
     return _physicalNodeId;
 }
 
+const boost::optional<ABT>& OptPhaseManager::getPostMemoPlan() const {
+    return _postMemoPlan;
+}
+
 const QueryHints& OptPhaseManager::getHints() const {
     return _hints;
 }
@@ -356,10 +366,6 @@ const Metadata& OptPhaseManager::getMetadata() const {
     return _metadata;
 }
 
-PrefixId& OptPhaseManager::getPrefixId() const {
-    return _prefixId;
-}
-
 const NodeToGroupPropsMap& OptPhaseManager::getNodeToGroupPropsMap() const {
     return _nodeToGroupPropsMap;
 }
@@ -368,8 +374,4 @@ NodeToGroupPropsMap& OptPhaseManager::getNodeToGroupPropsMap() {
     return _nodeToGroupPropsMap;
 }
 
-const RIDProjectionsMap& OptPhaseManager::getRIDProjections() const {
-    return _ridProjections;
-}
-
 }  // namespace mongo::optimizer
diff --git a/src/mongo/db/query/optimizer/opt_phase_manager.h b/src/mongo/db/query/optimizer/opt_phase_manager.h
index a7e719fe7ca..71e4ba989d9 100644
--- a/src/mongo/db/query/optimizer/opt_phase_manager.h
+++ b/src/mongo/db/query/optimizer/opt_phase_manager.h
@@ -84,6 +84,7 @@ public:
                     std::unique_ptr<CostEstimator> costEstimator,
                     PathToIntervalFn pathToInterval,
                     ConstFoldFn constFold,
+                    bool supportExplain,
                     DebugInfo debugInfo,
                     QueryHints queryHints = {});
 
@@ -108,6 +109,7 @@ public:
     static const PhaseSet& getAllRewritesSet();
 
     MemoPhysicalNodeId getPhysicalNodeId() const;
+    const boost::optional<ABT>& getPostMemoPlan() const;
 
     const QueryHints& getHints() const;
     QueryHints& getHints();
@@ -118,13 +120,9 @@ public:
 
     const Metadata& getMetadata() const;
 
-    PrefixId& getPrefixId() const;
-
     const NodeToGroupPropsMap& getNodeToGroupPropsMap() const;
     NodeToGroupPropsMap& getNodeToGroupPropsMap();
 
-    const RIDProjectionsMap& getRIDProjections() const;
-
 private:
     bool hasPhase(OptPhase phase) const;
 
@@ -159,6 +157,11 @@ private:
 
     const PhaseSet _phaseSet;
 
+    /**
+     * True if we should maintain extra internal state in support of explain.
+     */
+    const bool _supportExplain;
+
     const DebugInfo _debugInfo;
 
     QueryHints _hints;
@@ -210,6 +213,12 @@ private:
     MemoPhysicalNodeId _physicalNodeId;
 
     /**
+     * Post memo exploration phase plan (set if '_supportExplain' is set and if we have performed
+     * memo rewrites).
+     */
+    boost::optional<ABT> _postMemoPlan;
+
+    /**
      * Map from node to logical and physical properties.
      */
     NodeToGroupPropsMap _nodeToGroupPropsMap;
diff --git a/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp b/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
index f08fa2dccf0..9d8519a021d 100644
--- a/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
+++ b/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
@@ -151,6 +151,7 @@ OptPhaseManager makePhaseManager(
                            makeCostEstimator(costModel ? *costModel : getTestCostModel()),
                            defaultConvertPathToInterval,
                            ConstEval::constFold,
+                           true /*supportExplain*/,
                            std::move(debugInfo),
                            std::move(queryHints)};
 }
@@ -172,6 +173,7 @@ OptPhaseManager makePhaseManager(
                            makeCostEstimator(costModel ? *costModel : getTestCostModel()),
                            defaultConvertPathToInterval,
                            ConstEval::constFold,
+                           true /*supportExplain*/,
                            std::move(debugInfo),
                            std::move(queryHints)};
 }
@@ -191,6 +193,7 @@ OptPhaseManager makePhaseManagerRequireRID(OptPhaseManager::PhaseSet phaseSet,
                            makeCostEstimator(),
                            defaultConvertPathToInterval,
                            ConstEval::constFold,
+                           true /*supportExplain*/,
                            std::move(debugInfo),
                            std::move(queryHints)};
 }
diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl
index d4ca334491f..9db6b850840 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -34,6 +34,7 @@ global:
   cpp_includes:
     - "mongo/db/query/cost_model/cost_model_on_update.h"
     - "mongo/db/query/ce_mode_parameter.h"
+    - "mongo/db/query/explain_version_validator.h"
     - "mongo/db/query/sbe_plan_cache_on_parameter_change.h"
     - "mongo/db/query/telemetry_util.h"
     - "mongo/platform/atomic_proxy.h"
@@ -848,6 +849,16 @@ server_parameters:
     default: false
     test_only: true
 
+  internalCascadesOptimizerExplainVersion:
+    description: "Selects the explain version of the plans from the Cascades optimizer. The default is 'bson', other supported versions are 'v1', 'v2', and 'v2compact'"
+    set_at: [ startup, runtime ]
+    cpp_varname: "internalCascadesOptimizerExplainVersion"
+    cpp_vartype: synchronized_value<std::string>
+    default: "bson"
+    validator:
+        callback: optimizer::validateOptimizerExplainVersion
+    test_only: true
+
   internalCascadesOptimizerUseDescriptiveVarNames:
     description: "Enables generation of descriptive variable names to aid debugging."
     set_at: [ startup, runtime ]