SERVER-54405 Fall back to classic engine when match contains numeric path components

6 files changed, 219 insertions, 25 deletions

diff --git a/jstests/core/match_numeric_components.js b/jstests/core/match_numeric_components.js
new file mode 100644
index 00000000000..8db464fdb69
--- /dev/null
+++ b/jstests/core/match_numeric_components.js
@@ -0,0 +1,162 @@
+/**
+ * Tests behavior of the match language when using numeric path components.
+ */
+(function() {
+const coll = db.match_numeric_components;
+coll.drop();
+
+const kDocs = [
+    {_id: 0, "a": 42},
+    {_id: 1, "a": [42]},
+    {_id: 2, "a": {"0": 42}},
+    {_id: 3, "a": [[42]]},
+    {_id: 4, "a": [{"0": 42}]},
+    {_id: 5, "a": {"0": [42]}},
+    {_id: 6, "a": {"0": {"0": 42}}},
+    {_id: 7, "a": [[[42]]]},
+    {_id: 8, "a": [[{"0": 42}]]},
+    {_id: 9, "a": [{"0": [42]}]},
+    {_id: 10, "a": [{"0": {"0": 42}}]},
+    {_id: 11, "a": {"0": [[42]]}},
+    {_id: 12, "a": {"0": [{"0": 42}]}},
+    {_id: 13, "a": {"0": {"0": [42]}}},
+    {_id: 14, "a": {"0": {"0": {"0": 42}}}},
+    {_id: 15, "a": [[[[42]]]]},
+    {_id: 16, "a": [[[{"0": 42}]]]},
+    {_id: 17, "a": [[{"0": [42]}]]},
+    {_id: 18, "a": [[{"0": {"0": 42}}]]},
+    {_id: 19, "a": [{"0": [[42]]}]},
+    {_id: 20, "a": [{"0": [{"0": 42}]}]},
+    {_id: 21, "a": [{"0": {"0": [42]}}]},
+    {_id: 22, "a": [{"0": {"0": {"0": 42}}}]},
+    {_id: 23, "a": {"0": [[[42]]]}},
+    {_id: 24, "a": {"0": [[{"0": 42}]]}},
+    {_id: 25, "a": {"0": [{"0": [42]}]}},
+    {_id: 26, "a": {"0": [{"0": {"0": 42}}]}},
+    {_id: 27, "a": {"0": {"0": [[42]]}}},
+    {_id: 28, "a": {"0": {"0": [{"0": 42}]}}},
+    {_id: 29, "a": {"0": {"0": {"0": [42]}}}},
+    {_id: 30, "a": {"0": {"0": {"0": {"0": 42}}}}},
+];
+
+for (let doc of kDocs) {
+    coll.insert(doc);
+}
+
+{
+    const res = coll.find({"a.0": 42}).toArray();
+    const expected = [
+        {_id: 1, "a": [42]},
+        {_id: 2, "a": {"0": 42}},
+        {_id: 4, "a": [{"0": 42}]},
+        {_id: 5, "a": {"0": [42]}},
+        {_id: 9, "a": [{"0": [42]}]}
+    ];
+
+    assert.sameMembers(res, expected);
+}
+
+{
+    const res = coll.find({"a.0.0": 42}).toArray();
+    const expected = [
+        {_id: 3, "a": [[42]]},
+        {_id: 4, "a": [{"0": 42}]},
+        {_id: 5, "a": {"0": [42]}},
+        {_id: 6, "a": {"0": {"0": 42}}},
+        {_id: 8, "a": [[{"0": 42}]]},
+        {_id: 9, "a": [{"0": [42]}]},
+        {_id: 10, "a": [{"0": {"0": 42}}]},
+        {_id: 12, "a": {"0": [{"0": 42}]}},
+        {_id: 13, "a": {"0": {"0": [42]}}},
+        {_id: 17, "a": [[{"0": [42]}]]},
+        {_id: 20, "a": [{"0": [{"0": 42}]}]},
+        {_id: 21, "a": [{"0": {"0": [42]}}]},
+        {_id: 25, "a": {"0": [{"0": [42]}]}}
+    ];
+
+    assert.sameMembers(res, expected);
+}
+
+// Using a comparison.
+{
+    const res = coll.find({"a.0": {$gt: 41}}).toArray();
+    const expected = [
+        {_id: 1, "a": [42]},
+        {_id: 2, "a": {"0": 42}},
+        {_id: 4, "a": [{"0": 42}]},
+        {_id: 5, "a": {"0": [42]}},
+        {_id: 9, "a": [{"0": [42]}]}
+    ];
+
+    assert.sameMembers(res, expected);
+}
+
+// Using $in.
+{
+    const res = coll.find({"a.0": {$in: [41, 42, 43]}}).toArray();
+    const expected = [
+        {_id: 1, "a": [42]},
+        {_id: 2, "a": {"0": 42}},
+        {_id: 4, "a": [{"0": 42}]},
+        {_id: 5, "a": {"0": [42]}},
+        {_id: 9, "a": [{"0": [42]}]}
+    ];
+
+    assert.sameMembers(res, expected);
+}
+
+// Using an expression.
+{
+    const res = coll.find({"a.0": {$type: "number"}}).toArray();
+    const expected = [
+        {_id: 1, "a": [42]},
+        {_id: 2, "a": {"0": 42}},
+        {_id: 4, "a": [{"0": 42}]},
+        {_id: 5, "a": {"0": [42]}},
+        {_id: 9, "a": [{"0": [42]}]}
+    ];
+
+    assert.sameMembers(res, expected);
+}
+
+{
+    const res = coll.find({"a.0": {$mod: [42, 0]}}).toArray();
+    const expected = [
+        {_id: 1, "a": [42]},
+        {_id: 2, "a": {"0": 42}},
+        {_id: 4, "a": [{"0": 42}]},
+        {_id: 5, "a": {"0": [42]}},
+        {_id: 9, "a": [{"0": [42]}]}
+    ];
+
+    assert.sameMembers(res, expected);
+}
+
+// In a $or.
+{
+    const res = coll.find({$or: [{"a.0": 42}, {"notARealField": 123}]}).toArray();
+    const expected = [
+        {_id: 1, "a": [42]},
+        {_id: 2, "a": {"0": 42}},
+        {_id: 4, "a": [{"0": 42}]},
+        {_id: 5, "a": {"0": [42]}},
+        {_id: 9, "a": [{"0": [42]}]}
+    ];
+
+    assert.sameMembers(res, expected);
+}
+
+const coll2 = db.match_numeric_components2;
+coll2.drop();
+assert.commandWorked(coll2.insert({_id: 1, "b": "hello"}));
+assert.commandWorked(coll2.insert({_id: 2, "b": {"0": "hello"}}));
+assert.commandWorked(coll2.insert({_id: 3, "b": ["hello", "abc", "abc"]}));
+
+// Regexes are often something of a special case.
+{
+    const res = coll2.find({"b.0": {$regex: "hello"}}).toArray();
+    const expected = [{_id: 2, "b": {"0": "hello"}}, {_id: 3, "b": ["hello", "abc", "abc"]}];
+
+    assert.sameMembers(res, expected);
+}
+})();
diff --git a/jstests/core/multikey_geonear.js b/jstests/core/multikey_geonear.js
index 7e77293fe49..c389576bb1a 100644
--- a/jstests/core/multikey_geonear.js
+++ b/jstests/core/multikey_geonear.js
@@ -70,3 +70,13 @@ t.insert({_id: 2, a: [{b: [2, 2]}, {c: 2}]});
 
 cursor = t.find({"a.b": {$near: [2, 2]}, "a.c": {$gte: 0}});
 checkResults(cursor);
+
+// Check the case where we create a geo index on a specific array index.
+t.drop();
+t.createIndex({"a.0": "2dsphere"});
+t.insert({_id: 0, a: [{type: "Point", coordinates: [0, 0]}]});
+t.insert({_id: 1, a: [{type: "Point", coordinates: [1, 1]}]});
+t.insert({_id: 2, a: {0: {type: "Point", coordinates: [2, 2]}}});
+
+cursor = t.find({"a.0": {$near: {$geometry: {type: "Point", coordinates: [2, 2]}}}});
+checkResults(cursor);
diff --git a/src/mongo/db/matcher/expression.h b/src/mongo/db/matcher/expression.h
index 3c3077ea5df..2a9216e85e6 100644
--- a/src/mongo/db/matcher/expression.h
+++ b/src/mongo/db/matcher/expression.h
@@ -336,6 +336,12 @@ public:
     virtual const StringData path() const {
         return StringData();
     }
+    /**
+     * Similar to path(), but returns a FieldRef. Returns nullptr if there is no path.
+     */
+    virtual const FieldRef* fieldRef() const {
+        return nullptr;
+    }
 
     enum class MatchCategory {
         // Expressions that are leaves on the AST, these do not have any children.
diff --git a/src/mongo/db/matcher/expression_parser.cpp b/src/mongo/db/matcher/expression_parser.cpp
index 434955fcf95..cb2e31c67cb 100644
--- a/src/mongo/db/matcher/expression_parser.cpp
+++ b/src/mongo/db/matcher/expression_parser.cpp
@@ -90,6 +90,16 @@ bool hasNode(const MatchExpression* root, MatchExpression::MatchType type) {
     return false;
 }
 
+void addExpressionToRoot(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                         AndMatchExpression* root,
+                         std::unique_ptr<MatchExpression> newNode) {
+    if (newNode->fieldRef() && newNode->fieldRef()->hasNumericPathComponents()) {
+        // TODO SERVER-49852: Currently SBE cannot handle match expressions with numeric path
+        // components due to some of the complexity around how arrays are handled.
+        expCtx->sbeCompatible = false;
+    }
+    root->add(std::move(newNode));
+}
 }  // namespace
 
 namespace mongo {
@@ -306,7 +316,8 @@ StatusWithMatchExpression parse(const BSONObj& obj,
             auto result = parseRegexElement(e.fieldNameStringData(), e, expCtx);
             if (!result.isOK())
                 return result;
-            root->add(result.getValue().release());
+
+            addExpressionToRoot(expCtx, root.get(), std::move(result.getValue()));
             continue;
         }
 
@@ -322,7 +333,7 @@ StatusWithMatchExpression parse(const BSONObj& obj,
         if (!eq.isOK())
             return eq;
 
-        root->add(eq.getValue().release());
+        addExpressionToRoot(expCtx, root.get(), std::move(eq.getValue()));
     }
 
     if (root->numChildren() == 1) {
@@ -1954,7 +1965,7 @@ Status parseSub(StringData name,
                 auto s =
                     parseGeo(name, PathAcceptingKeyword::GEO_NEAR, sub, expCtx, allowedFeatures);
                 if (s.isOK()) {
-                    root->add(s.getValue().release());
+                    addExpressionToRoot(expCtx, root, std::move(s.getValue()));
                 }
 
                 // Propagate geo parsing result to caller.
@@ -1969,8 +1980,9 @@ Status parseSub(StringData name,
         if (!s.isOK())
             return s.getStatus();
 
-        if (s.getValue())
-            root->add(s.getValue().release());
+        if (s.getValue()) {
+            addExpressionToRoot(expCtx, root, std::move(s.getValue()));
+        }
     }
 
     return Status::OK();
diff --git a/src/mongo/db/matcher/expression_path.h b/src/mongo/db/matcher/expression_path.h
index 191e0356121..4604e57e81b 100644
--- a/src/mongo/db/matcher/expression_path.h
+++ b/src/mongo/db/matcher/expression_path.h
@@ -68,10 +68,14 @@ public:
         return false;
     }
 
-    const StringData path() const final {
+    const StringData path() const override final {
         return _elementPath.fieldRef().dottedField();
     }
 
+    const FieldRef* fieldRef() const override final {
+        return &(_elementPath.fieldRef());
+    }
+
     /**
      * Resets the path for this expression. Note that this method will make a copy of 'path' such
      * that there's no lifetime requirements for the string which 'path' points into.
diff --git a/src/mongo/db/query/sbe_stage_builder_filter.cpp b/src/mongo/db/query/sbe_stage_builder_filter.cpp
index b63a43296fd..8f5d9b36b55 100644
--- a/src/mongo/db/query/sbe_stage_builder_filter.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_filter.cpp
@@ -453,7 +453,7 @@ EvalExprStagePair generatePathTraversal(EvalStage inputStage,
  * evaluating the predicate on a single value.
  */
 void generatePredicate(MatchExpressionVisitorContext* context,
-                       StringData path,
+                       const FieldRef* path,
                        MakePredicateFn makePredicate,
                        LeafTraversalMode mode = LeafTraversalMode::kArrayAndItsElements,
                        bool useCombinator = true) {
@@ -461,10 +461,10 @@ void generatePredicate(MatchExpressionVisitorContext* context,
 
     auto&& [expr, stage] = [&]() {
         if (frame.data().inputSlot) {
-            if (!path.empty()) {
+            if (path && !path->empty()) {
                 return generatePathTraversal(frame.extractStage(),
                                              *frame.data().inputSlot,
-                                             FieldRef{path},
+                                             *path,
                                              0,
                                              context->planNodeId,
                                              context->slotIdGenerator,
@@ -474,7 +474,7 @@ void generatePredicate(MatchExpressionVisitorContext* context,
                                              context->stateHelper);
             } else {
                 // If matchExpr's parent is a ElemMatchValueMatchExpression, then
-                // matchExpr()->path() will be empty. In this case, 'inputSlot' will be a
+                // matchExpr()->fieldRef() will be nullptr. In this case, 'inputSlot' will be a
                 // "correlated slot" that holds the value of the ElemMatchValueMatchExpression's
                 // field path, and we should apply the predicate directly on 'inputSlot' without
                 // array traversal.
@@ -491,11 +491,11 @@ void generatePredicate(MatchExpressionVisitorContext* context,
             // current field path - the index scan will extract the value for this field path and
             // will store it in a corresponding slot in the 'indexKeySlots' map.
 
-            tassert(5273402, "Field path cannot be empty for an index filter", !path.empty());
+            tassert(5273402, "Field path cannot be empty for an index filter", path);
 
-            auto it = context->indexKeySlots.find(path.toString());
+            auto it = context->indexKeySlots.find(path->dottedField());
             tassert(5273403,
-                    str::stream() << "Unknown field path in index filter: " << path,
+                    str::stream() << "Unknown field path in index filter: " << path->dottedField(),
                     it != context->indexKeySlots.end());
 
             auto result = makePredicate(it->second, frame.extractStage());
@@ -576,7 +576,7 @@ void generateArraySize(MatchExpressionVisitorContext* context,
     };
 
     generatePredicate(context,
-                      matchExpr->path(),
+                      matchExpr->fieldRef(),
                       std::move(makePredicate),
                       LeafTraversalMode::kDoNotTraverseLeaf);
 }
@@ -694,7 +694,7 @@ void generateComparison(MatchExpressionVisitorContext* context,
                 std::move(inputStage)};
     };
 
-    generatePredicate(context, expr->path(), std::move(makePredicate));
+    generatePredicate(context, expr->fieldRef(), std::move(makePredicate));
 }
 
 /**
@@ -782,7 +782,7 @@ void generateBitTest(MatchExpressionVisitorContext* context,
                 std::move(inputStage)};
     };
 
-    generatePredicate(context, expr->path(), std::move(makePredicate));
+    generatePredicate(context, expr->fieldRef(), std::move(makePredicate));
 }
 
 // Each logical expression child is evaluated in a separate EvalFrame. Set up a new EvalFrame with a
@@ -1170,7 +1170,7 @@ public:
         // 'makePredicate' defined above returns a state instead of plain boolean value, so there is
         // no need to use combinator for it.
         generatePredicate(_context,
-                          matchExpr->path(),
+                          matchExpr->fieldRef(),
                           std::move(makePredicate),
                           LeafTraversalMode::kDoNotTraverseLeaf,
                           false /* useCombinator */);
@@ -1222,7 +1222,7 @@ public:
         // 'makePredicate' defined above returns a state instead of plain boolean value, so there is
         // no need to use combinator for it.
         generatePredicate(_context,
-                          matchExpr->path(),
+                          matchExpr->fieldRef(),
                           std::move(makePredicate),
                           LeafTraversalMode::kDoNotTraverseLeaf,
                           false /* useCombinator */);
@@ -1240,7 +1240,7 @@ public:
                     std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->path(), std::move(makePredicate));
+        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
     }
 
     void visit(const ExprMatchExpression* matchExpr) final {
@@ -1330,7 +1330,7 @@ public:
                         std::move(inputStage)};
             };
 
-            generatePredicate(_context, expr->path(), std::move(makePredicate));
+            generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
             return;
         } else {
             // If the InMatchExpression contains regex patterns, then we need to handle a regex-only
@@ -1427,7 +1427,7 @@ public:
                 return {regexOutputSlot, std::move(regexStage)};
             };
             generatePredicate(_context,
-                              expr->path(),
+                              expr->fieldRef(),
                               std::move(makePredicate),
                               LeafTraversalMode::kArrayAndItsElements);
         }
@@ -1506,7 +1506,7 @@ public:
                     std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->path(), std::move(makePredicate));
+        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
     }
 
     void visit(const NorMatchExpression* expr) final {
@@ -1553,7 +1553,7 @@ public:
                     std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->path(), std::move(makePredicate));
+        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
     }
 
     void visit(const SizeMatchExpression* expr) final {
@@ -1573,7 +1573,7 @@ public:
                     std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->path(), std::move(makePredicate));
+        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
     }
 
     void visit(const WhereMatchExpression* expr) final {
@@ -1589,7 +1589,7 @@ public:
             return {std::move(whereExpr), std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->path(), std::move(makePredicate));
+        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
     }
 
     void visit(const WhereNoOpMatchExpression* expr) final {}