SERVER-67638 [CQF] Error when projection has >10k fields

5 files changed, 62 insertions, 44 deletions

diff --git a/src/mongo/db/pipeline/abt/abt_pipeline_test.cpp b/src/mongo/db/pipeline/abt/abt_pipeline_test.cpp
index 12ae02dd9ea..93fd6c57703 100644
--- a/src/mongo/db/pipeline/abt/abt_pipeline_test.cpp
+++ b/src/mongo/db/pipeline/abt/abt_pipeline_test.cpp
@@ -214,17 +214,17 @@ TEST(ABTTranslate, MatchWithOrConvertedToIn) {
         "Sargable [Complete]\n"
         "|   |   |   |   |   requirementsMap: \n"
         "|   |   |   |   |       refProjection: scan_0, path: 'PathGet [a] PathTraverse [1] "
-        "PathIdentity []', intervals: {{{[Const [1], Const [1]]}} U {{[Const [2], Const [2]]}} U "
-        "{{[Const [3], Const [3]]}}}\n"
+        "PathIdentity []', intervals: {{{[Const [3], Const [3]]}} U {{[Const [1], Const [1]]}} U "
+        "{{[Const [2], Const [2]]}}}\n"
         "|   |   |   |   candidateIndexes: \n"
-        "|   |   |   |       candidateId: 1, index1, {}, {0}, {{{[Const [1], Const [1]]}} U "
-        "{{[Const [2], Const [2]]}} U {{[Const [3], Const [3]]}}}\n"
+        "|   |   |   |       candidateId: 1, index1, {}, {0}, {{{[Const [3], Const [3]]}} U "
+        "{{[Const [1], Const [1]]}} U {{[Const [2], Const [2]]}}}\n"
         "|   |   |   scanParams: \n"
         "|   |   |       {'a': evalTemp_0}\n"
         "|   |   |           residualReqs: \n"
         "|   |   |               refProjection: evalTemp_0, path: 'PathTraverse [1] PathIdentity "
-        "[]', intervals: {{{[Const [1], Const [1]]}} U {{[Const [2], Const [2]]}} U {{[Const [3], "
-        "Const [3]]}}}, entryIndex: 0\n"
+        "[]', intervals: {{{[Const [3], Const [3]]}} U {{[Const [1], Const [1]]}} U {{[Const [2], "
+        "Const [2]]}}}, entryIndex: 0\n"
         "|   |   BindBlock:\n"
         "|   RefBlock: \n"
         "|       Variable [scan_0]\n"
@@ -618,17 +618,17 @@ TEST(ABTTranslate, MatchBasic) {
         "Filter []\n"
         "|   EvalFilter []\n"
         "|   |   Variable [scan_0]\n"
-        "|   PathGet [b]\n"
+        "|   PathGet [a]\n"
         "|   PathTraverse [1]\n"
         "|   PathCompare [Eq]\n"
-        "|   Const [2]\n"
+        "|   Const [1]\n"
         "Filter []\n"
         "|   EvalFilter []\n"
         "|   |   Variable [scan_0]\n"
-        "|   PathGet [a]\n"
+        "|   PathGet [b]\n"
         "|   PathTraverse [1]\n"
         "|   PathCompare [Eq]\n"
-        "|   Const [1]\n"
+        "|   Const [2]\n"
         "Scan [collection]\n"
         "    BindBlock:\n"
         "        [scan_0]\n"
@@ -794,17 +794,17 @@ TEST(ABTTranslate, MatchProject) {
         "|   EvalFilter []\n"
         "|   |   Variable [combinedProjection_0]\n"
         "|   PathComposeA []\n"
-        "|   |   PathGet [s]\n"
+        "|   |   PathGet [c]\n"
         "|   |   PathTraverse [1]\n"
-        "|   |   PathComposeM []\n"
-        "|   |   |   PathCompare [Lt]\n"
-        "|   |   |   Const [\"\"]\n"
-        "|   |   PathCompare [Gte]\n"
-        "|   |   Const [10]\n"
-        "|   PathGet [c]\n"
+        "|   |   PathCompare [Eq]\n"
+        "|   |   Const [2]\n"
+        "|   PathGet [s]\n"
         "|   PathTraverse [1]\n"
-        "|   PathCompare [Eq]\n"
-        "|   Const [2]\n"
+        "|   PathComposeM []\n"
+        "|   |   PathCompare [Lt]\n"
+        "|   |   Const [\"\"]\n"
+        "|   PathCompare [Gte]\n"
+        "|   Const [10]\n"
         "Evaluation []\n"
         "|   BindBlock:\n"
         "|       [combinedProjection_0]\n"
@@ -1812,20 +1812,20 @@ TEST(ABTTranslate, RangeIndex) {
         "|   PathGet [a]\n"
         "|   PathTraverse [1]\n"
         "|   PathComposeM []\n"
-        "|   |   PathCompare [Gte]\n"
-        "|   |   Const [nan]\n"
-        "|   PathCompare [Lt]\n"
-        "|   Const [90]\n"
+        "|   |   PathCompare [Lt]\n"
+        "|   |   Const [\"\"]\n"
+        "|   PathCompare [Gt]\n"
+        "|   Const [70]\n"
         "Filter []\n"
         "|   EvalFilter []\n"
         "|   |   Variable [scan_0]\n"
         "|   PathGet [a]\n"
         "|   PathTraverse [1]\n"
         "|   PathComposeM []\n"
-        "|   |   PathCompare [Lt]\n"
-        "|   |   Const [\"\"]\n"
-        "|   PathCompare [Gt]\n"
-        "|   Const [70]\n"
+        "|   |   PathCompare [Gte]\n"
+        "|   |   Const [nan]\n"
+        "|   PathCompare [Lt]\n"
+        "|   Const [90]\n"
         "Scan [collection]\n"
         "    BindBlock:\n"
         "        [scan_0]\n"
@@ -1886,7 +1886,7 @@ TEST(ABTTranslate, RangeIndex) {
         "|   |       [sideId_0]\n"
         "|   |           Const [1]\n"
         "|   IndexScan [{'<rid>': rid_0}, scanDefName: collection, indexDefName: index1, interval: "
-        "{[Const [nan], Const [90])}]\n"
+        "{(Const [70], Const [\"\"])}]\n"
         "|       BindBlock:\n"
         "|           [rid_0]\n"
         "|               Source []\n"
@@ -1895,7 +1895,7 @@ TEST(ABTTranslate, RangeIndex) {
         "|       [sideId_0]\n"
         "|           Const [0]\n"
         "IndexScan [{'<rid>': rid_0}, scanDefName: collection, indexDefName: index1, interval: "
-        "{(Const [70], Const [\"\"])}]\n"
+        "{[Const [nan], Const [90])}]\n"
         "    BindBlock:\n"
         "        [rid_0]\n"
         "            Source []\n",
@@ -1926,14 +1926,14 @@ TEST(ABTTranslate, Index1) {
             "Filter []\n"
             "|   EvalFilter []\n"
             "|   |   Variable [scan_0]\n"
-            "|   PathGet [b]\n"
+            "|   PathGet [a]\n"
             "|   PathTraverse [1]\n"
             "|   PathCompare [Eq]\n"
             "|   Const [2]\n"
             "Filter []\n"
             "|   EvalFilter []\n"
             "|   |   Variable [scan_0]\n"
-            "|   PathGet [a]\n"
+            "|   PathGet [b]\n"
             "|   PathTraverse [1]\n"
             "|   PathCompare [Eq]\n"
             "|   Const [2]\n"
@@ -1998,14 +1998,14 @@ TEST(ABTTranslate, Index1) {
             "Filter []\n"
             "|   EvalFilter []\n"
             "|   |   Variable [scan_0]\n"
-            "|   PathGet [b]\n"
+            "|   PathGet [a]\n"
             "|   PathTraverse [1]\n"
             "|   PathCompare [Eq]\n"
             "|   Const [2]\n"
             "Filter []\n"
             "|   EvalFilter []\n"
             "|   |   Variable [scan_0]\n"
-            "|   PathGet [a]\n"
+            "|   PathGet [b]\n"
             "|   PathTraverse [1]\n"
             "|   PathCompare [Eq]\n"
             "|   Const [2]\n"
diff --git a/src/mongo/db/pipeline/abt/match_expression_visitor.cpp b/src/mongo/db/pipeline/abt/match_expression_visitor.cpp
index 4e736889730..2c6642bfe5a 100644
--- a/src/mongo/db/pipeline/abt/match_expression_visitor.cpp
+++ b/src/mongo/db/pipeline/abt/match_expression_visitor.cpp
@@ -715,11 +715,14 @@ private:
             return;
         }
 
-        ABT node = _ctx.pop();
-        for (size_t i = 0; i < childCount - 1; i++) {
-            node = make<Composition>(_ctx.pop(), std::move(node));
+        ABTVector nodes;
+        for (size_t i = 0; i < childCount; i++) {
+            nodes.push_back(_ctx.pop());
         }
-        _ctx.push(std::move(node));
+
+        // Construct a balanced composition tree.
+        maybeComposePaths<Composition>(nodes);
+        _ctx.push(std::move(nodes.front()));
     }
 
     void unsupportedExpression(const MatchExpression* expr) const {
diff --git a/src/mongo/db/query/ce/ce_heuristic_test.cpp b/src/mongo/db/query/ce/ce_heuristic_test.cpp
index 37d80beed0b..e833484547f 100644
--- a/src/mongo/db/query/ce/ce_heuristic_test.cpp
+++ b/src/mongo/db/query/ce/ce_heuristic_test.cpp
@@ -159,9 +159,9 @@ TEST(CEHeuristicTest, CEWithoutOptimizationNestedConjAndDisj3) {
         "]}"
         "]}";
     ASSERT_MATCH_CE_CARD_NO_OPT(ht, query, 0.0, 0.0);
-    ASSERT_MATCH_CE_CARD_NO_OPT(ht, query, 3.48545, 9.0);
-    ASSERT_MATCH_CE_CARD_NO_OPT(ht, query, 9.96732, 99.0);
-    ASSERT_MATCH_CE_CARD_NO_OPT(ht, query, 25.3708, 1000.0);
+    ASSERT_MATCH_CE_CARD_NO_OPT(ht, query, 3.01561, 9.0);
+    ASSERT_MATCH_CE_CARD_NO_OPT(ht, query, 9.37173, 99.0);
+    ASSERT_MATCH_CE_CARD_NO_OPT(ht, query, 19.3064, 1000.0);
 }
 
 TEST(CEHeuristicTest, CEWithoutOptimizationNestedConjAndDisj4) {
diff --git a/src/mongo/db/query/optimizer/interval_intersection_test.cpp b/src/mongo/db/query/optimizer/interval_intersection_test.cpp
index c83443ce010..dba4aa9ebf1 100644
--- a/src/mongo/db/query/optimizer/interval_intersection_test.cpp
+++ b/src/mongo/db/query/optimizer/interval_intersection_test.cpp
@@ -147,13 +147,13 @@ TEST(IntervalIntersection, SingleFieldIntersection) {
          "|   |   BindBlock:\n"
          "|   |       [rid_0]\n"
          "|   |           Source []\n"
-         "|   IndexScan [{'<rid>': rid_0}, scanDefName: coll, indexDefName: index1, "
-         "interval: {(Const [40], Const [44])}]\n"
+         "|   IndexScan [{'<rid>': rid_0}, scanDefName: coll, indexDefName: index1, interval: "
+         "{(Const [9], Const [12])}]\n"
          "|       BindBlock:\n"
          "|           [rid_0]\n"
          "|               Source []\n"
-         "IndexScan [{'<rid>': rid_0}, scanDefName: coll, indexDefName: index1, interval: "
-         "{(Const [9], Const [12])}]\n"
+         "IndexScan [{'<rid>': rid_0}, scanDefName: coll, indexDefName: index1, interval: {(Const "
+         "[40], Const [44])}]\n"
          "    BindBlock:\n"
          "        [rid_0]\n"
          "            Source []\n"},
diff --git a/src/mongo/db/query/optimizer/utils/utils.h b/src/mongo/db/query/optimizer/utils/utils.h
index 28690c50262..b5fe8c652ca 100644
--- a/src/mongo/db/query/optimizer/utils/utils.h
+++ b/src/mongo/db/query/optimizer/utils/utils.h
@@ -81,6 +81,21 @@ inline void maybeComposePath(ABT& composition, ABT child) {
 }
 
 /**
+ * Creates a balanced tree of composition elements over the input vector which it modifies in place.
+ * In the end at most one element remains in the vector.
+ */
+template <class Element = PathComposeM>
+inline void maybeComposePaths(ABTVector& paths) {
+    while (paths.size() > 1) {
+        const size_t half = paths.size() / 2;
+        for (size_t i = 0; i < half; i++) {
+            maybeComposePath<Element>(paths.at(i), std::move(paths.at(paths.size() - i - 1)));
+        }
+        paths.resize(paths.size() - half, make<Blackhole>());
+    }
+}
+
+/**
  * Used to vend out fresh ids for projection names.
  */
 class PrefixId {