SERVER-75541 Pretty printer for SargableNode

3 files changed, 266 insertions, 84 deletions

diff --git a/buildscripts/gdb/optimizer_printers.py b/buildscripts/gdb/optimizer_printers.py
index 8c6c399e694..756b7580e35 100644
--- a/buildscripts/gdb/optimizer_printers.py
+++ b/buildscripts/gdb/optimizer_printers.py
@@ -51,6 +51,14 @@ class IntervalPrinter(OptimizerTypePrinter):
         super().__init__(val, "ExplainGenerator::explainInterval")
 
 
+class CandidateIndexEntryPrinter(OptimizerTypePrinter):
+    """Pretty-printer for mongo::optimizer::CandidateIndexEntry."""
+
+    def __init__(self, val):
+        """Initialize CandidateIndexEntryPrinter."""
+        super().__init__(val, "ExplainGenerator::explainCandidateIndex")
+
+
 class IntervalExprPrinter(OptimizerTypePrinter):
     """Pretty-printer for mongo::optimizer::IntervalRequirement::Node."""
 
@@ -99,6 +107,7 @@ class FixedArityNodePrinter(object):
         self.val = val
         self.arity = arity
         self.name = name
+        self.custom_children = []
 
     @staticmethod
     def display_hint():
@@ -109,7 +118,16 @@ class FixedArityNodePrinter(object):
         """children."""
 
         prior_indent = ABTPrinter.indent_level
-        current_indent = ABTPrinter.indent_level + self.arity - 1
+        current_indent = ABTPrinter.indent_level + self.arity + len(self.custom_children) - 1
+        for child in self.custom_children:
+            lhs = "\n"
+            for _ in range(current_indent):
+                lhs += "|   "
+
+            ABTPrinter.indent_level = current_indent
+            yield lhs, child
+            current_indent -= 1
+
         for i in range(self.arity):
             lhs = "\n"
             for _ in range(current_indent):
@@ -121,6 +139,10 @@ class FixedArityNodePrinter(object):
             current_indent -= 1
         ABTPrinter.indent_level = prior_indent
 
+    # Adds a custom child node which is not directly contained in the "_nodes" member variable.
+    def add_child(self, child):
+        self.custom_children.append(child)
+
     def to_string(self):
         # Default for nodes which just print their type.
         return self.name
@@ -550,9 +572,9 @@ class FieldProjectionMapPrinter(object):
         if get_boost_optional(root_proj) is not None:
             res += "<root>: " + str(root_proj) + ", "
         # Rely on default printer for std::set, but remove the extra metadata at the start.
-        # TODO SERVER-75541 pretty print field projections map.
-        # field_projections = self.val["_fieldProjections"]
-        # res += str(field_projections).split("elems  =")[-1]
+        field_projections = self.val["_fieldProjections"]
+        res += "<empty>" if field_projections["size_"] == 0 else str(field_projections).split(
+            "elems  =")[-1]
         res += "}"
         return res
 
@@ -599,7 +621,7 @@ class IndexScanNodePrinter(FixedArityNodePrinter):
     def to_string(self):
         return "IndexScan[{{{}}}, scanDef={}, indexDef={}, interval={}]".format(
             self.val["_fieldProjectionMap"], self.val["_scanDefName"], self.val["_indexDefName"],
-            self.val["_indexInterval"])
+            self.val["_indexInterval"]).replace("\n", "")
 
 
 class SeekNodePrinter(FixedArityNodePrinter):
@@ -610,7 +632,7 @@ class SeekNodePrinter(FixedArityNodePrinter):
         super().__init__(val, 2, "Seek")
 
     def to_string(self):
-        return "Seek[rid_projection: {}, {}, scanDef: {}]".format(self.val["_rid_projectionName"],
+        return "Seek[rid_projection: {}, {}, scanDef: {}]".format(self.val["_ridProjectionName"],
                                                                   self.val["_fieldProjectionMap"],
                                                                   self.val["_scanDefName"])
 
@@ -638,12 +660,70 @@ class MemoPhysicalDelegatorNodePrinter(FixedArityNodePrinter):
                                                                     self.val["_nodeId"]["_index"])
 
 
+class ResidualRequirementPrinter(object):
+    """Pretty-printer for ResidualRequirement."""
+
+    def __init__(self, val):
+        """Initialize ResidualRequirementPrinter."""
+        self.val = val
+
+    def to_string(self):
+        key = self.val["_key"]
+        req = self.val["_req"]
+        res = "<"
+        if get_boost_optional(key["_projectionName"]) is not None:
+            res += "refProj: " + str(get_boost_optional(key["_projectionName"])) + ", "
+
+        res += "path: '" + str(key["_path"]).replace("|   ", "").replace("\n", " -> ") + "'"
+
+        if get_boost_optional(req["_boundProjectionName"]) is not None:
+            res += "boundProj: " + str(get_boost_optional(req["_boundProjectionName"])) + ", "
+
+        res += ">"
+        return res
+
+
 class SargableNodePrinter(FixedArityNodePrinter):
     """Pretty-printer for SargableNode."""
 
     def __init__(self, val):
         """Initialize SargableNodePrinter."""
-        super().__init__(val, 3, "Sargable")
+        # Although Sargable technically has 3 children, avoid printing the refs (child1) and bind block (child2).
+        super().__init__(val, 1, "Sargable")
+
+        # Add children for requirements, candidateIndex, and scan_params.
+        self.add_child(str(self.val["_reqMap"]).replace("\n", ""))
+        self.add_child(self.print_candidate_indexes())
+
+        self.scan_params = get_boost_optional(self.val["_scanParams"])
+        if self.scan_params is not None:
+            self.add_child(self.print_scan_params())
+
+    def print_scan_params(self):
+        res = "scan_params: (proj: " + str(self.scan_params["_fieldProjectionMap"]) + ", "
+        residual_reqs = get_boost_optional(self.scan_params["_residualRequirements"])
+        if residual_reqs is not None:
+            res += "residual: " + str(residual_reqs)
+        res += ")"
+        return res
+
+    def print_candidate_indexes(self):
+        res = "candidateIndexes: ["
+        indexes = Vector(self.val["_candidateIndexes"])
+        for i in range(indexes.count()):
+            if i > 0:
+                res += ", "
+            res += "<id: " + str(i) + ", " + str(indexes.get(i)).replace("\n", "") + ">"
+        res += "]"
+        return res
+
+    @staticmethod
+    def index_req_to_string(index_req):
+        req_map = ["Index", "Seek", "Complete"]
+        return req_map[index_req]
+
+    def to_string(self):
+        return "Sargable [" + self.index_req_to_string(self.val["_target"]) + "]"
 
 
 class RIDIntersectNodePrinter(FixedArityNodePrinter):
@@ -908,6 +988,69 @@ class ABTPrinter(PolyValuePrinter):
         super().__init__(ABTPrinter.abt_type_set, ABTPrinter.abt_namespace, val)
 
 
+class AtomPrinter(object):
+    """Pretty-printer for Atom."""
+
+    def __init__(self, val):
+        """Initialize AtomPrinter."""
+        self.val = val
+
+    def to_string(self):
+        return self.val["_expr"]
+
+
+class ConjunctionPrinter(object):
+    """Pretty-printer for Conjunction."""
+
+    def __init__(self, val, separator=" ^ "):
+        """Initialize ConjunctionPrinter."""
+        self.val = val
+        self.dynamic_nodes = Vector(self.val["_dyNodes"])
+        self.dynamic_count = self.dynamic_nodes.count()
+        self.separator = separator
+
+    def to_string(self):
+        if self.dynamic_count == 0:
+            return "<empty>"
+
+        res = ""
+        first = True
+        for child in self.dynamic_nodes:
+            if first:
+                first = False
+            else:
+                res += self.separator
+
+            res += str(child)
+        return res
+
+
+class DisjunctionPrinter(ConjunctionPrinter):
+    """Pretty-printer for Disjunction."""
+
+    def __init__(self, val):
+        super().__init__(val, " U ")
+
+
+class BoolExprPrinter(PolyValuePrinter):
+    """Pretty-printer for BoolExpr."""
+
+    type_set = ["Atom", "Conjunction", "Disjunction"]
+
+    def __init__(self, val, template_type):
+        """Initialize BoolExprPrinter."""
+        namespace = "mongo::optimizer::BoolExpr<" + template_type + ">::"
+        super().__init__(BoolExprPrinter.type_set, namespace, val)
+
+
+class ResidualReqExprPrinter(BoolExprPrinter):
+    """Pretty-printer for BoolExpr<ResidualRequirement>."""
+
+    def __init__(self, val):
+        """Initialize ResidualReqExprPrinter."""
+        super().__init__(val, "mongo::optimizer::ResidualRequirement")
+
+
 def register_abt_printers(pp):
     """Registers a number of pretty printers related to the CQF optimizer."""
 
@@ -934,6 +1077,28 @@ def register_abt_printers(pp):
     pp.add("PartialSchemaRequirements", "mongo::optimizer::PartialSchemaRequirements", False,
            PartialSchemaReqMapPrinter)
 
+    # ResidualRequirement printer.
+    pp.add("ResidualRequirement", "mongo::optimizer::ResidualRequirement", False,
+           ResidualRequirementPrinter)
+
+    # CandidateIndexEntry printer.
+    pp.add("CandidateIndexEntry", "mongo::optimizer::CandidateIndexEntry", False,
+           CandidateIndexEntryPrinter)
+
+    pp.add(
+        "ResidualRequirementExpr",
+        ("mongo::optimizer::algebra::PolyValue<" +
+         "mongo::optimizer::BoolExpr<mongo::optimizer::ResidualRequirement>::Atom, " +
+         "mongo::optimizer::BoolExpr<mongo::optimizer::ResidualRequirement>::Conjunction, " +
+         "mongo::optimizer::BoolExpr<mongo::optimizer::ResidualRequirement>::Disjunction>"),
+        False,
+        ResidualReqExprPrinter,
+    )
+    for bool_type in BoolExprPrinter.type_set:
+        pp.add(bool_type,
+               "mongo::optimizer::BoolExpr<mongo::optimizer::ResidualRequirement>::" + bool_type,
+               False, getattr(sys.modules[__name__], bool_type + "Printer"))
+
     # Utility types within the optimizer.
     pp.add("StrongStringAlias", "mongo::optimizer::StrongStringAlias", True,
            StrongStringAliasPrinter)
diff --git a/src/mongo/db/query/optimizer/explain.cpp b/src/mongo/db/query/optimizer/explain.cpp
index 4cb4c18abf7..f5c349ac664 100644
--- a/src/mongo/db/query/optimizer/explain.cpp
+++ b/src/mongo/db/query/optimizer/explain.cpp
@@ -1045,6 +1045,91 @@ public:
         return printer.str();
     }
 
+    void printCandidateIndexEntry(ExplainPrinter& local,
+                                  const CandidateIndexEntry& candidateIndexEntry) {
+        local.fieldName("indexDefName", ExplainVersion::V3)
+            .print(candidateIndexEntry._indexDefName)
+            .separator(", ");
+
+        local.separator("{");
+        printFieldProjectionMap(local, candidateIndexEntry._fieldProjectionMap);
+        local.separator("}, {");
+
+        {
+            if constexpr (version < ExplainVersion::V3) {
+                bool first = true;
+                for (const auto type : candidateIndexEntry._predTypes) {
+                    if (first) {
+                        first = false;
+                    } else {
+                        local.print(", ");
+                    }
+                    local.print(IndexFieldPredTypeEnum::toString[static_cast<int>(type)]);
+                }
+            } else if constexpr (version == ExplainVersion::V3) {
+                std::vector<ExplainPrinter> printers;
+                for (const auto type : candidateIndexEntry._predTypes) {
+                    ExplainPrinter local1;
+                    local1.print(IndexFieldPredTypeEnum::toString[static_cast<int>(type)]);
+                    printers.push_back(std::move(local1));
+                }
+                local.fieldName("predType").print(printers);
+            } else {
+                MONGO_UNREACHABLE;
+            }
+        }
+
+        local.separator("}, ");
+        {
+            if (candidateIndexEntry._eqPrefixes.size() == 1) {
+                local.fieldName("intervals", ExplainVersion::V3);
+
+                ExplainPrinter intervals = printIntervalExpr<CompoundIntervalRequirement>(
+                    candidateIndexEntry._eqPrefixes.front()._interval);
+                local.printSingleLevel(intervals, "" /*singleLevelSpacer*/);
+            } else {
+                std::vector<ExplainPrinter> eqPrefixPrinters;
+                for (const auto& entry : candidateIndexEntry._eqPrefixes) {
+                    ExplainPrinter eqPrefixPrinter;
+                    eqPrefixPrinter.fieldName("startPos", ExplainVersion::V3)
+                        .print(entry._startPos)
+                        .separator(", ");
+
+                    ExplainPrinter intervals =
+                        printIntervalExpr<CompoundIntervalRequirement>(entry._interval);
+                    eqPrefixPrinter.separator("[")
+                        .fieldName("interval", ExplainVersion::V3)
+                        .printSingleLevel(intervals, "" /*singleLevelSpacer*/)
+                        .separator("]");
+
+                    eqPrefixPrinters.push_back(std::move(eqPrefixPrinter));
+                }
+
+                local.print(eqPrefixPrinters);
+            }
+        }
+
+        if (const auto& residualReqs = candidateIndexEntry._residualRequirements) {
+            local.separator("}, ");
+            if constexpr (version < ExplainVersion::V3) {
+                ExplainPrinter residualReqMapPrinter;
+                printResidualRequirements(residualReqMapPrinter, *residualReqs);
+                local.print(residualReqMapPrinter);
+            } else if (version == ExplainVersion::V3) {
+                printResidualRequirements(local, *residualReqs);
+            } else {
+                MONGO_UNREACHABLE;
+            }
+        }
+    }
+
+
+    std::string printCandidateIndexEntry(const CandidateIndexEntry& indexEntry) {
+        ExplainPrinter printer;
+        printCandidateIndexEntry(printer, indexEntry);
+        return printer.str();
+    }
+
     void printPartialSchemaEntry(ExplainPrinter& printer, const PartialSchemaEntry& entry) {
         const auto& [key, req] = entry;
 
@@ -1429,83 +1514,8 @@ public:
                 const CandidateIndexEntry& candidateIndexEntry = candidateIndexes.at(index);
 
                 ExplainPrinter local;
-                local.fieldName("candidateId")
-                    .print(index + 1)
-                    .separator(", ")
-                    .fieldName("indexDefName", ExplainVersion::V3)
-                    .print(candidateIndexEntry._indexDefName)
-                    .separator(", ");
-
-                local.separator("{");
-                printFieldProjectionMap(local, candidateIndexEntry._fieldProjectionMap);
-                local.separator("}, {");
-
-                {
-                    if constexpr (version < ExplainVersion::V3) {
-                        bool first = true;
-                        for (const auto type : candidateIndexEntry._predTypes) {
-                            if (first) {
-                                first = false;
-                            } else {
-                                local.print(", ");
-                            }
-                            local.print(IndexFieldPredTypeEnum::toString[static_cast<int>(type)]);
-                        }
-                    } else if constexpr (version == ExplainVersion::V3) {
-                        std::vector<ExplainPrinter> printers;
-                        for (const auto type : candidateIndexEntry._predTypes) {
-                            ExplainPrinter local1;
-                            local1.print(IndexFieldPredTypeEnum::toString[static_cast<int>(type)]);
-                            printers.push_back(std::move(local1));
-                        }
-                        local.fieldName("predType").print(printers);
-                    } else {
-                        MONGO_UNREACHABLE;
-                    }
-                }
-
-                local.separator("}, ");
-                {
-                    if (candidateIndexEntry._eqPrefixes.size() == 1) {
-                        local.fieldName("intervals", ExplainVersion::V3);
-
-                        ExplainPrinter intervals = printIntervalExpr<CompoundIntervalRequirement>(
-                            candidateIndexEntry._eqPrefixes.front()._interval);
-                        local.printSingleLevel(intervals, "" /*singleLevelSpacer*/);
-                    } else {
-                        std::vector<ExplainPrinter> eqPrefixPrinters;
-                        for (const auto& entry : candidateIndexEntry._eqPrefixes) {
-                            ExplainPrinter eqPrefixPrinter;
-                            eqPrefixPrinter.fieldName("startPos", ExplainVersion::V3)
-                                .print(entry._startPos)
-                                .separator(", ");
-
-                            ExplainPrinter intervals =
-                                printIntervalExpr<CompoundIntervalRequirement>(entry._interval);
-                            eqPrefixPrinter.separator("[")
-                                .fieldName("interval", ExplainVersion::V3)
-                                .printSingleLevel(intervals, "" /*singleLevelSpacer*/)
-                                .separator("]");
-
-                            eqPrefixPrinters.push_back(std::move(eqPrefixPrinter));
-                        }
-
-                        local.print(eqPrefixPrinters);
-                    }
-                }
-
-                if (const auto& residualReqs = candidateIndexEntry._residualRequirements) {
-                    if constexpr (version < ExplainVersion::V3) {
-                        ExplainPrinter residualReqMapPrinter;
-                        printResidualRequirements(residualReqMapPrinter, *residualReqs);
-                        local.print(residualReqMapPrinter);
-                    } else if (version == ExplainVersion::V3) {
-                        printResidualRequirements(local, *residualReqs);
-                    } else {
-                        MONGO_UNREACHABLE;
-                    }
-                }
-
+                local.fieldName("candidateId").print(index + 1).separator(", ");
+                printCandidateIndexEntry(local, candidateIndexEntry);
                 candidateIndexesPrinters.push_back(std::move(local));
             }
             ExplainPrinter candidateIndexesPrinter;
@@ -3057,4 +3067,9 @@ std::string ExplainGenerator::explainIntervalExpr(
     ExplainGeneratorV2 gen;
     return gen.printIntervalExpr<CompoundIntervalRequirement>(intervalExpr).str();
 }
+
+std::string ExplainGenerator::explainCandidateIndex(const CandidateIndexEntry& indexEntry) {
+    ExplainGeneratorV2 gen;
+    return gen.printCandidateIndexEntry(indexEntry);
+}
 }  // namespace mongo::optimizer
diff --git a/src/mongo/db/query/optimizer/explain.h b/src/mongo/db/query/optimizer/explain.h
index 20137ac9157..f93caf2722f 100644
--- a/src/mongo/db/query/optimizer/explain.h
+++ b/src/mongo/db/query/optimizer/explain.h
@@ -122,6 +122,8 @@ public:
     static std::string explainIntervalExpr(const IntervalReqExpr::Node& intervalExpr);
 
     static std::string explainIntervalExpr(const CompoundIntervalReqExpr::Node& intervalExpr);
+
+    static std::string explainCandidateIndex(const CandidateIndexEntry& indexEntry);
 };
 
 }  // namespace mongo::optimizer