SERVER-70726 Push down AND of supported predicates on the same field in column_scan

3 files changed, 99 insertions, 13 deletions

diff --git a/src/mongo/db/matcher/expression_algo.cpp b/src/mongo/db/matcher/expression_algo.cpp
index a51bc0e838d..e69f2aa0122 100644
--- a/src/mongo/db/matcher/expression_algo.cpp
+++ b/src/mongo/db/matcher/expression_algo.cpp
@@ -477,11 +477,6 @@ std::unique_ptr<MatchExpression> tryAddExpr(StringData path,
     if (FieldRef(path).hasNumericPathComponents())
         return me->shallowClone();
 
-    // (TODO SERVER-70726) addExpr will rightfully create AND for expressions on the same path, but
-    // we cannot translate AND into EExpressions yet.
-    if (out.find(path) != out.end())
-        return me->shallowClone();
-
     addExpr(path, me->shallowClone(), out);
     return nullptr;
 }
diff --git a/src/mongo/db/matcher/expression_algo_test.cpp b/src/mongo/db/matcher/expression_algo_test.cpp
index b5d4ef87fff..65768791cca 100644
--- a/src/mongo/db/matcher/expression_algo_test.cpp
+++ b/src/mongo/db/matcher/expression_algo_test.cpp
@@ -2179,6 +2179,63 @@ TEST(SplitMatchExpressionForColumns, SupportsTypePredicatesArray) {
     ASSERT(residual == nullptr);
 }
 
+TEST(SplitMatchExpressionForColumns, SupportsCombiningPredicatesWithAnd) {
+    ParsedMatchExpression compoundFilter(
+        "{"
+        " albatross: {$gte: 100},"
+        " albatross: {$lt: 200},"
+        " albatross: {$mod: [2, 0]}"
+        "}");
+    auto&& [splitUp, residual] = expression::splitMatchExpressionForColumns(compoundFilter.get());
+    ASSERT_EQ(splitUp.size(), 1) << splitUp.size();
+    ASSERT(splitUp.contains("albatross"));
+    ASSERT(splitUp["albatross"]->matchType() == MatchExpression::AND)
+        << splitUp.at("albatross")->toString();
+
+    // Don't care about the order of terms in AND.
+    auto andExpr = splitUp.at("albatross").get();
+    ASSERT_EQ(splitUp.at("albatross")->numChildren(), 3) << splitUp.at("albatross")->toString();
+    stdx::unordered_set<MatchExpression::MatchType> terms;
+    for (int i = 0; i < 3; i++) {
+        auto mt = andExpr->getChild(i)->matchType();
+        ASSERT(terms.insert(mt).second) << mt;
+    }
+    ASSERT(terms.count(MatchExpression::GTE) == 1) << "Should have GTE term";
+    ASSERT(terms.count(MatchExpression::LT) == 1) << "Should have LT term";
+    ASSERT(terms.count(MatchExpression::MOD) == 1) << "Should have MOD term";
+
+    ASSERT(residual == nullptr);
+}
+
+TEST(SplitMatchExpressionForColumns, SupportsAndFlattensNestedAnd) {
+    ParsedMatchExpression compoundFilter(
+        "{"
+        " $and: ["
+        "   {albatross: {$gte: 100}},"
+        "   {$and: [{albatross: {$lt: 200}}, {albatross: {$mod: [2, 0]}}]}"
+        " ]"
+        "}");
+    auto&& [splitUp, residual] = expression::splitMatchExpressionForColumns(compoundFilter.get());
+    ASSERT_EQ(splitUp.size(), 1) << splitUp.size();
+    ASSERT(splitUp.contains("albatross"));
+    ASSERT(splitUp["albatross"]->matchType() == MatchExpression::AND)
+        << splitUp.at("albatross")->toString();
+
+    // Don't care about the order of terms in AND.
+    auto andExpr = splitUp.at("albatross").get();
+    ASSERT_EQ(splitUp.at("albatross")->numChildren(), 3) << splitUp.at("albatross")->toString();
+    stdx::unordered_set<MatchExpression::MatchType> terms;
+    for (int i = 0; i < 3; i++) {
+        auto mt = andExpr->getChild(i)->matchType();
+        ASSERT(terms.insert(mt).second) << mt;
+    }
+    ASSERT(terms.count(MatchExpression::GTE) == 1) << "Should have GTE term";
+    ASSERT(terms.count(MatchExpression::LT) == 1) << "Should have LT term";
+    ASSERT(terms.count(MatchExpression::MOD) == 1) << "Should have MOD term";
+
+    ASSERT(residual == nullptr);
+}
+
 TEST(SplitMatchExpressionForColumns, DoesNotSupportNotQueries) {
     {
         ParsedMatchExpression notEqFilter("{albatross: {$not: {$eq: 2}}}");
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index a196821e3dc..ca8e52b6e1c 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -740,20 +740,54 @@ std::unique_ptr<sbe::EExpression> generatePerColumnPredicate(StageBuilderState&
     MONGO_UNREACHABLE;
 }
 
-std::unique_ptr<sbe::EExpression> generatePerColumnFilterExpr(StageBuilderState& state,
-                                                              const MatchExpression* me,
-                                                              sbe::value::SlotId inputSlot) {
-    auto lambdaFrameId = state.frameIdGenerator->generate();
-    auto lambdaParam = sbe::EVariable{lambdaFrameId, 0};
+std::unique_ptr<sbe::EExpression> generateLeafExpr(StageBuilderState& state,
+                                                   const MatchExpression* me,
+                                                   sbe::FrameId lambdaFrameId,
+                                                   sbe::value::SlotId inputSlot) {
+    sbe::EVariable lambdaParam = sbe::EVariable{lambdaFrameId, 0};
 
     auto lambdaExpr = sbe::makeE<sbe::ELocalLambda>(
         lambdaFrameId, generatePerColumnPredicate(state, me, lambdaParam));
 
-    auto traverseFuncName = (me->matchType() == MatchExpression::EXISTS ||
-                             me->matchType() == MatchExpression::TYPE_OPERATOR)
+    const MatchExpression::MatchType mt = me->matchType();
+    auto traverserName = (mt == MatchExpression::EXISTS || mt == MatchExpression::TYPE_OPERATOR)
         ? "traverseCsiCellTypes"
         : "traverseCsiCellValues";
-    return makeFunction(traverseFuncName, makeVariable(inputSlot), std::move(lambdaExpr));
+    return makeFunction(traverserName, makeVariable(inputSlot), std::move(lambdaExpr));
+}
+
+std::unique_ptr<sbe::EExpression> generatePerColumnLogicalAndExpr(StageBuilderState& state,
+                                                                  const AndMatchExpression* me,
+                                                                  sbe::FrameId lambdaFrameId,
+                                                                  sbe::value::SlotId inputSlot) {
+    const auto cTerms = me->numChildren();
+    tassert(7072600, "AND should have at least one child", cTerms > 0);
+
+    auto logical = generateLeafExpr(state, me->getChild(cTerms - 1), lambdaFrameId, inputSlot);
+    if (cTerms == 1)
+        return logical;
+
+    // TODO SERVER-70110: Replace the right-handed tree with a folded AND node when it becomes
+    // available.
+    for (int i = cTerms - 2; i >= 0; i--) {
+        logical = makeBinaryOp(sbe::EPrimBinary::logicAnd,
+                               generateLeafExpr(state, me->getChild(i), lambdaFrameId, inputSlot),
+                               std::move(logical));
+    }
+    return logical;
+}
+
+std::unique_ptr<sbe::EExpression> generatePerColumnFilterExpr(StageBuilderState& state,
+                                                              const MatchExpression* me,
+                                                              sbe::value::SlotId inputSlot) {
+    auto lambdaFrameId = state.frameIdGenerator->generate();
+
+    if (me->matchType() == MatchExpression::AND) {
+        return generatePerColumnLogicalAndExpr(
+            state, checked_cast<const AndMatchExpression*>(me), lambdaFrameId, inputSlot);
+    }
+
+    return generateLeafExpr(state, me, lambdaFrameId, inputSlot);
 }
 }  // namespace