SERVER-60830 Added more predicates for $eq $lt $lte in internal unpack bucket

2 files changed, 30 insertions, 12 deletions

diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
index 7e8f33c5a6b..af2e276de9a 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
@@ -642,9 +642,11 @@ std::unique_ptr<MatchExpression> createComparisonPredicate(
     }
 
     BSONObj minTime;
+    BSONObj maxTime;
     if (isTimeField) {
         auto timeField = matchExprData.Date();
         minTime = BSON("" << timeField - Seconds(bucketMaxSpanSeconds));
+        maxTime = BSON("" << timeField + Seconds(bucketMaxSpanSeconds));
     }
 
     auto minPath = std::string{kControlMinFieldNamePrefix} + matchExprPath;
@@ -660,15 +662,17 @@ std::unique_ptr<MatchExpression> createComparisonPredicate(
             // is adjusted by the max range for a bucket to approximate the max bucket value given
             // the min. Also include a predicate against the _id field which is converted to the
             // minimum for the range of ObjectIds corresponding to the given date. In
-            // addition, we include a {'control.min' : {$gte: 'time - bucketMaxSpanSeconds'}}
-            // predicate which will be helpful in reducing bounds for index scans on 'time' field
-            // and routing on mongos.
+            // addition, we include a {'control.min' : {$gte: 'time - bucketMaxSpanSeconds'}} and
+            // a {'control.max' : {$lte: 'time + bucketMaxSpanSeconds'}} predicate which will be
+            // helpful in reducing bounds for index scans on 'time' field and routing on mongos.
             return isTimeField
                 ? makePredicate(
                       MatchExprPredicate<InternalExprLTEMatchExpression>(minPath, matchExprData),
                       MatchExprPredicate<InternalExprGTEMatchExpression>(minPath,
                                                                          minTime.firstElement()),
                       MatchExprPredicate<InternalExprGTEMatchExpression>(maxPath, matchExprData),
+                      MatchExprPredicate<InternalExprLTEMatchExpression>(maxPath,
+                                                                         maxTime.firstElement()),
                       MatchExprPredicate<LTEMatchExpression, Value>(
                           kBucketIdFieldName,
                           constructObjectIdValue<LTEMatchExpression>(matchExprData,
@@ -732,12 +736,14 @@ std::unique_ptr<MatchExpression> createComparisonPredicate(
             // For $lt, make a $lt predicate against 'control.min'. In addition, if the comparison
             // is against the 'time' field, include a predicate against the _id field which is
             // converted to the minimum for the corresponding range of ObjectIds. In
-            // addition, we include a {'control.min' : {$lt: 'time - bucketMaxSpanSeconds'}}
+            // addition, we include a {'control.max' : {$lt: 'time + bucketMaxSpanSeconds'}}
             // predicate which will be helpful in reducing bounds for index scans on 'time' field
             // and routing on mongos.
             return isTimeField
                 ? makePredicate(
                       MatchExprPredicate<InternalExprLTMatchExpression>(minPath, matchExprData),
+                      MatchExprPredicate<InternalExprLTMatchExpression>(maxPath,
+                                                                        maxTime.firstElement()),
                       MatchExprPredicate<LTMatchExpression, Value>(
                           kBucketIdFieldName,
                           constructObjectIdValue<LTMatchExpression>(matchExprData,
@@ -750,12 +756,14 @@ std::unique_ptr<MatchExpression> createComparisonPredicate(
             // For $lte, make a $lte predicate against 'control.min'. In addition, if the comparison
             // is against the 'time' field, include a predicate against the _id field which is
             // converted to the maximum for the corresponding range of ObjectIds. In
-            // addition, we include a {'control.min' : {$lte: 'time - bucketMaxSpanSeconds'}}
+            // addition, we include a {'control.max' : {$lte: 'time + bucketMaxSpanSeconds'}}
             // predicate which will be helpful in reducing bounds for index scans on 'time' field
             // and routing on mongos.
             return isTimeField
                 ? makePredicate(
                       MatchExprPredicate<InternalExprLTEMatchExpression>(minPath, matchExprData),
+                      MatchExprPredicate<InternalExprLTEMatchExpression>(maxPath,
+                                                                         maxTime.firstElement()),
                       MatchExprPredicate<LTEMatchExpression, Value>(
                           kBucketIdFieldName,
                           constructObjectIdValue<LTEMatchExpression>(matchExprData,
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/create_predicates_on_bucket_level_field_test.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/create_predicates_on_bucket_level_field_test.cpp
index 2c307863b6c..ff48e79592d 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/create_predicates_on_bucket_level_field_test.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test/create_predicates_on_bucket_level_field_test.cpp
@@ -405,6 +405,7 @@ TEST_F(InternalUnpackBucketPredicateMappingOptimizationTest,
 TEST_F(InternalUnpackBucketPredicateMappingOptimizationTest, OptimizeMapsTimePredicatesOnId) {
     auto date = Date_t::now();
     const auto dateMinusBucketSpan = date - Seconds{3600};
+    const auto datePlusBucketSpan = date + Seconds{3600};
     {
         auto timePred = BSON("$match" << BSON("time" << BSON("$lt" << date)));
         auto pipelines = {
@@ -427,11 +428,14 @@ TEST_F(InternalUnpackBucketPredicateMappingOptimizationTest, OptimizeMapsTimePre
             auto andExpr = dynamic_cast<AndMatchExpression*>(predicate.get());
             auto children = andExpr->getChildVector();
 
-            ASSERT_EQ(children->size(), 2);
+            ASSERT_EQ(children->size(), 3);
             ASSERT_BSONOBJ_EQ((*children)[0]->serialize(true),
                               BSON("control.min.time" << BSON("$_internalExprLt" << date)));
+            ASSERT_BSONOBJ_EQ(
+                (*children)[1]->serialize(true),
+                BSON("control.max.time" << BSON("$_internalExprLt" << datePlusBucketSpan)));
 
-            auto idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[1].get());
+            auto idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[2].get());
 
             ASSERT_EQ(idPred->path(), "_id"_sd);
             ASSERT_EQ(idPred->getData().type(), BSONType::jstOID);
@@ -462,11 +466,14 @@ TEST_F(InternalUnpackBucketPredicateMappingOptimizationTest, OptimizeMapsTimePre
             auto andExpr = dynamic_cast<AndMatchExpression*>(predicate.get());
             auto children = andExpr->getChildVector();
 
-            ASSERT_EQ(children->size(), 2);
+            ASSERT_EQ(children->size(), 3);
             ASSERT_BSONOBJ_EQ((*children)[0]->serialize(true),
                               BSON("control.min.time" << BSON("$_internalExprLte" << date)));
+            ASSERT_BSONOBJ_EQ(
+                (*children)[1]->serialize(true),
+                BSON("control.max.time" << BSON("$_internalExprLte" << datePlusBucketSpan)));
 
-            auto idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[1].get());
+            auto idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[2].get());
 
             ASSERT_EQ(idPred->path(), "_id"_sd);
             ASSERT_EQ(idPred->getData().type(), BSONType::jstOID);
@@ -497,7 +504,7 @@ TEST_F(InternalUnpackBucketPredicateMappingOptimizationTest, OptimizeMapsTimePre
             auto andExpr = dynamic_cast<AndMatchExpression*>(predicate.get());
             auto children = andExpr->getChildVector();
 
-            ASSERT_EQ(children->size(), 5);
+            ASSERT_EQ(children->size(), 6);
             ASSERT_BSONOBJ_EQ((*children)[0]->serialize(true),
                               BSON("control.min.time" << BSON("$_internalExprLte" << date)));
             ASSERT_BSONOBJ_EQ(
@@ -505,8 +512,11 @@ TEST_F(InternalUnpackBucketPredicateMappingOptimizationTest, OptimizeMapsTimePre
                 BSON("control.min.time" << BSON("$_internalExprGte" << dateMinusBucketSpan)));
             ASSERT_BSONOBJ_EQ((*children)[2]->serialize(true),
                               BSON("control.max.time" << BSON("$_internalExprGte" << date)));
+            ASSERT_BSONOBJ_EQ(
+                (*children)[3]->serialize(true),
+                BSON("control.max.time" << BSON("$_internalExprLte" << datePlusBucketSpan)));
 
-            auto idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[3].get());
+            auto idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[4].get());
 
             ASSERT_EQ(idPred->path(), "_id"_sd);
             ASSERT_EQ(idPred->getData().type(), BSONType::jstOID);
@@ -515,7 +525,7 @@ TEST_F(InternalUnpackBucketPredicateMappingOptimizationTest, OptimizeMapsTimePre
             oid.init(date);
             ASSERT_TRUE(oid.compare(idPred->getData().OID()) < 0);
 
-            idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[4].get());
+            idPred = dynamic_cast<ComparisonMatchExpressionBase*>((*children)[5].get());
 
             ASSERT_EQ(idPred->path(), "_id"_sd);
             ASSERT_EQ(idPred->getData().type(), BSONType::jstOID);