SERVER-69217: [CQF] Convert PathObj to interval

4 files changed, 92 insertions, 5 deletions

diff --git a/src/mongo/db/pipeline/abt/utils.cpp b/src/mongo/db/pipeline/abt/utils.cpp
index 520ef41db52..fbf98c6a0df 100644
--- a/src/mongo/db/pipeline/abt/utils.cpp
+++ b/src/mongo/db/pipeline/abt/utils.cpp
@@ -215,19 +215,26 @@ public:
 
     PathToIntervalTransport() {}
 
-    ResultType transport(const ABT& /*n*/, const PathArr& /*node*/) {
-        auto [lowBound, lowInclusive] =
-            getMinMaxBoundForType(true /*isMin*/, sbe::value::TypeTags::Array);
+    template <sbe::value::TypeTags tag>
+    ResultType getBoundsForNode() {
+        auto [lowBound, lowInclusive] = getMinMaxBoundForType(true /*isMin*/, tag);
         invariant(lowBound);
 
-        auto [highBound, highInclusive] =
-            getMinMaxBoundForType(false /*isMin*/, sbe::value::TypeTags::Array);
+        auto [highBound, highInclusive] = getMinMaxBoundForType(false /*isMin*/, tag);
         invariant(highBound);
 
         return IntervalReqExpr::makeSingularDNF(IntervalRequirement{
             {lowInclusive, std::move(*lowBound)}, {highInclusive, std::move(*highBound)}});
     }
 
+    ResultType transport(const ABT& /*n*/, const PathArr& /*node*/) {
+        return getBoundsForNode<sbe::value::TypeTags::Array>();
+    }
+
+    ResultType transport(const ABT& /*n*/, const PathObj& /*node*/) {
+        return getBoundsForNode<sbe::value::TypeTags::Object>();
+    }
+
     template <typename T, typename... Ts>
     ResultType transport(const ABT& /*n*/, const T& /*node*/, Ts&&...) {
         return {};
diff --git a/src/mongo/db/query/optimizer/logical_rewriter_optimizer_test.cpp b/src/mongo/db/query/optimizer/logical_rewriter_optimizer_test.cpp
index 777e3178aa3..bd6049947ea 100644
--- a/src/mongo/db/query/optimizer/logical_rewriter_optimizer_test.cpp
+++ b/src/mongo/db/query/optimizer/logical_rewriter_optimizer_test.cpp
@@ -27,7 +27,9 @@
  *    it in the license file.
  */
 
+#include "mongo/db/pipeline/abt/utils.h"
 #include "mongo/db/query/optimizer/cascades/ce_heuristic.h"
+#include "mongo/db/query/optimizer/cascades/cost_derivation.h"
 #include "mongo/db/query/optimizer/cascades/logical_props_derivation.h"
 #include "mongo/db/query/optimizer/cascades/rewriter_rules.h"
 #include "mongo/db/query/optimizer/explain.h"
diff --git a/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp b/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp
index d82c93b720b..546f2ac297e 100644
--- a/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp
+++ b/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp
@@ -3622,6 +3622,81 @@ TEST(PhysRewriter, ObjectElemMatch) {
         optimized);
 }
 
+TEST(PhysRewriter, ObjectElemMatchPathObj) {
+    using namespace properties;
+    PrefixId prefixId;
+
+    ABT scanNode = make<ScanNode>("root", "c1");
+
+    ABT filterNode = make<FilterNode>(
+        make<EvalFilter>(
+            make<PathGet>(
+                "a",
+                make<PathComposeM>(
+                    make<PathObj>(),
+                    make<PathTraverse>(
+                        make<PathComposeM>(
+                            make<PathGet>("b",
+                                          make<PathTraverse>(
+                                              make<PathCompare>(Operations::Eq, Constant::int64(1)),
+                                              PathTraverse::kSingleLevel)),
+                            make<PathGet>("c",
+                                          make<PathTraverse>(
+                                              make<PathCompare>(Operations::Eq, Constant::int64(2)),
+                                              PathTraverse::kSingleLevel))),
+                        PathTraverse::kSingleLevel))),
+            make<Variable>("root")),
+        std::move(scanNode));
+
+    ABT rootNode =
+        make<RootNode>(ProjectionRequirement{ProjectionNameVector{"root"}}, std::move(filterNode));
+
+    OptPhaseManager phaseManager(
+        {OptPhase::MemoSubstitutionPhase,
+         OptPhase::MemoExplorationPhase,
+         OptPhase::MemoImplementationPhase},
+        prefixId,
+        false /*requireRID*/,
+        {{{"c1",
+           ScanDefinition{{},
+                          {{"index1",
+                            makeCompositeIndexDefinition(
+                                {{"b", CollationOp::Ascending, true /*isMultiKey*/},
+                                 {"a", CollationOp::Ascending, true /*isMultiKey*/}})}}}}}},
+        std::make_unique<HeuristicCE>(),
+        std::make_unique<DefaultCosting>(),
+        defaultConvertPathToInterval,
+        {true /*debugMode*/, 2 /*debugLevel*/, DebugInfo::kIterationLimitForTests});
+
+    ABT optimized = rootNode;
+    phaseManager.optimize(optimized);
+    ASSERT_EQ(4, phaseManager.getMemo().getStats()._physPlanExplorationCount);
+    // We currently cannot use indexes with ObjectElemMatch.
+    ASSERT_EXPLAIN_V2Compact(
+        "Root []\n"
+        "|   |   projections: \n"
+        "|   |       root\n"
+        "|   RefBlock: \n"
+        "|       Variable [root]\n"
+        "Filter []\n"
+        "|   EvalFilter []\n"
+        "|   |   Variable [root]\n"
+        "|   PathGet [a] PathTraverse [1] PathComposeM []\n"
+        "|   |   PathGet [c] PathTraverse [1] PathCompare [Eq] Const [2]\n"
+        "|   PathGet [b] PathTraverse [1] PathCompare [Eq] Const [1]\n"
+        "Filter []\n"
+        "|   EvalFilter []\n"
+        "|   |   Variable [evalTemp_0]\n"
+        "|   PathObj []\n"
+        "PhysicalScan [{'<root>': root, 'a': evalTemp_0}, c1]\n"
+        "    BindBlock:\n"
+        "        [evalTemp_0]\n"
+        "            Source []\n"
+        "        [root]\n"
+        "            Source []\n",
+        optimized);
+}
+
 TEST(PhysRewriter, ArrayConstantIndex) {
     using namespace properties;
     PrefixId prefixId;
diff --git a/src/mongo/db/query/optimizer/utils/utils.cpp b/src/mongo/db/query/optimizer/utils/utils.cpp
index c0d09635c0b..b514de51a13 100644
--- a/src/mongo/db/query/optimizer/utils/utils.cpp
+++ b/src/mongo/db/query/optimizer/utils/utils.cpp
@@ -1351,6 +1351,9 @@ void lowerPartialSchemaRequirement(const PartialSchemaKey& key,
         if (auto conversion = pathToInterval(make<PathArr>());
             conversion && *conversion == req.getIntervals()) {
             path = make<PathArr>();
+        } else if (auto conversion = pathToInterval(make<PathObj>());
+                   conversion && *conversion == req.getIntervals()) {
+            path = make<PathObj>();
         }
     }
     if (path.is<PathIdentity>()) {