1 files changed, 78 insertions, 27 deletions

diff --git a/src/mongo/db/exec/bucket_unpacker.h b/src/mongo/db/exec/bucket_unpacker.h
index 287bd9f2540..29f2f1f30d2 100644
--- a/src/mongo/db/exec/bucket_unpacker.h
+++ b/src/mongo/db/exec/bucket_unpacker.h
@@ -54,11 +54,18 @@ namespace mongo {
  */
 class BucketSpec {
 public:
+    // When unpackin buckets with kInclude we must produce measurements that contain the
+    // set of fields. Otherwise, if the kExclude option is used, the measurements will include the
+    // set difference between all fields in the bucket and the provided fields.
+    enum class Behavior { kInclude, kExclude };
+
     BucketSpec() = default;
     BucketSpec(const std::string& timeField,
                const boost::optional<std::string>& metaField,
                const std::set<std::string>& fields = {},
-               const std::set<std::string>& computedProjections = {});
+               Behavior behavior = Behavior::kExclude,
+               const std::set<std::string>& computedProjections = {},
+               bool usesExtendedRange = false);
     BucketSpec(const BucketSpec&);
     BucketSpec(BucketSpec&&);
 
@@ -92,6 +99,14 @@ public:
         return _fieldSet;
     }
 
+    void setBehavior(Behavior behavior) {
+        _behavior = behavior;
+    }
+
+    Behavior behavior() const {
+        return _behavior;
+    }
+
     void addComputedMetaProjFields(const StringData& field) {
         _computedMetaProjFields.emplace(field);
     }
@@ -104,6 +119,14 @@ public:
         _computedMetaProjFields.erase(field);
     }
 
+    void setUsesExtendedRange(bool usesExtendedRange) {
+        _usesExtendedRange = usesExtendedRange;
+    }
+
+    bool usesExtendedRange() const {
+        return _usesExtendedRange;
+    }
+
     // Returns whether 'field' depends on a pushed down $addFields or computed $project.
     bool fieldIsComputed(StringData field) const;
 
@@ -118,27 +141,45 @@ public:
         kError,
     };
 
+    struct BucketPredicate {
+        // A loose predicate is a predicate which returns true when any measures of a bucket
+        // matches.
+        std::unique_ptr<MatchExpression> loosePredicate;
+
+        // A tight predicate is a predicate which returns true when all measures of a bucket
+        // matches.
+        std::unique_ptr<MatchExpression> tightPredicate;
+    };
+
     /**
-     * Takes a predicate after $_internalUnpackBucket on a bucketed field as an argument and
-     * attempts to map it to a new predicate on the 'control' field. For example, the predicate
-     * {a: {$gt: 5}} will generate the predicate {control.max.a: {$_internalExprGt: 5}}, which will
-     * be added before the $_internalUnpackBucket stage.
+     * Takes a predicate after $_internalUnpackBucket as an argument and attempts to rewrite it as
+     * new predicates on the 'control' field. There will be a 'loose' predicate that will match if
+     * some of the event field matches, also a 'tight' predicate that will match if all of the event
+     * field matches.
      *
-     * If the original predicate is on the bucket's timeField we may also create a new predicate
-     * on the '_id' field to assist in index utilization. For example, the predicate
-     * {time: {$lt: new Date(...)}} will generate the following predicate:
+     * For example, the event level predicate {a: {$gt: 5}} will generate the loose predicate
+     * {control.max.a: {$_internalExprGt: 5}}. The loose predicate will be added before the
+     * $_internalUnpackBucket stage to filter out buckets with no match.
+     *
+     * Ideally, we'd like to add a tight predicate such as {control.min.a: {$_internalExprGt: 5}} to
+     * evaluate the filter on bucket level to avoid unnecessary event level evaluation. However, a
+     * bucket might contain events with missing fields that are skipped when computing the controls,
+     * so in reality we only add a tight predicate on timeField which is required to exist.
+     *
+     * If the original predicate is on the bucket's timeField we may also create a new loose
+     * predicate on the '_id' field (as it incorporates min time for the bucket) to assist in index
+     * utilization. For example, the predicate {time: {$lt: new Date(...)}} will generate the
+     * following predicate:
      * {$and: [
      *      {_id: {$lt: ObjectId(...)}},
      *      {control.min.time: {$_internalExprLt: new Date(...)}}
      * ]}
      *
-     * If the provided predicate is ineligible for this mapping, the function will return a nullptr.
-     * This should be interpreted as an always-true predicate.
-     *
-     * When using IneligiblePredicatePolicy::kIgnore, if the predicate can't be pushed down, it
-     * returns null. When using IneligiblePredicatePolicy::kError it raises a user error.
+     * If the provided predicate is ineligible for this mapping and using
+     * IneligiblePredicatePolicy::kIgnore, both loose and tight predicates will be set to nullptr.
+     * When using IneligiblePredicatePolicy::kError it raises a user error.
      */
-    static std::unique_ptr<MatchExpression> createPredicatesOnBucketLevelField(
+    static BucketPredicate createPredicatesOnBucketLevelField(
         const MatchExpression* matchExpr,
         const BucketSpec& bucketSpec,
         int bucketMaxSpanSeconds,
@@ -184,6 +225,7 @@ public:
 private:
     // The set of field names in the data region that should be included or excluded.
     std::set<std::string> _fieldSet;
+    Behavior _behavior = Behavior::kExclude;
 
     // Set of computed meta field projection names. Added at the end of materialized
     // measurements.
@@ -194,6 +236,7 @@ private:
 
     boost::optional<std::string> _metaField = boost::none;
     boost::optional<HashedFieldName> _metaFieldHashed = boost::none;
+    bool _usesExtendedRange = false;
 };
 
 /**
@@ -201,10 +244,6 @@ private:
  */
 class BucketUnpacker {
 public:
-    // When BucketUnpacker is created with kInclude it must produce measurements that contain the
-    // set of fields. Otherwise, if the kExclude option is used, the measurements will include the
-    // set difference between all fields in the bucket and the provided fields.
-    enum class Behavior { kInclude, kExclude };
     /**
      * Returns the number of measurements in the bucket in O(1) time.
      */
@@ -214,7 +253,7 @@ public:
     static const std::set<StringData> reservedBucketFieldNames;
 
     BucketUnpacker();
-    BucketUnpacker(BucketSpec spec, Behavior unpackerBehavior);
+    BucketUnpacker(BucketSpec spec);
     BucketUnpacker(const BucketUnpacker& other) = delete;
     BucketUnpacker(BucketUnpacker&& other);
     ~BucketUnpacker();
@@ -228,6 +267,11 @@ public:
     Document getNext();
 
     /**
+     * Similar to the previous method, but return a BSON object instead.
+     */
+    BSONObj getNextBson();
+
+    /**
      * This method will extract the j-th measurement from the bucket. A precondition of this method
      * is that j >= 0 && j <= the number of measurements within the underlying bucket.
      */
@@ -246,7 +290,6 @@ public:
      */
     BucketUnpacker copy() const {
         BucketUnpacker unpackerCopy;
-        unpackerCopy._unpackerBehavior = _unpackerBehavior;
         unpackerCopy._spec = _spec;
         unpackerCopy._includeMetaField = _includeMetaField;
         unpackerCopy._includeTimeField = _includeTimeField;
@@ -256,10 +299,10 @@ public:
     /**
      * This resets the unpacker to prepare to unpack a new bucket described by the given document.
      */
-    void reset(BSONObj&& bucket);
+    void reset(BSONObj&& bucket, bool bucketMatchedQuery = false);
 
-    Behavior behavior() const {
-        return _unpackerBehavior;
+    BucketSpec::Behavior behavior() const {
+        return _spec.behavior();
     }
 
     const BucketSpec& bucketSpec() const {
@@ -270,6 +313,10 @@ public:
         return _bucket;
     }
 
+    bool bucketMatchedQuery() const {
+        return _bucketMatchedQuery;
+    }
+
     bool includeMetaField() const {
         return _includeMetaField;
     }
@@ -306,7 +353,7 @@ public:
         return std::string{timeseries::kControlMaxFieldNamePrefix} + field;
     }
 
-    void setBucketSpecAndBehavior(BucketSpec&& bucketSpec, Behavior behavior);
+    void setBucketSpec(BucketSpec&& bucketSpec);
     void setIncludeMinTimeAsMetadata();
     void setIncludeMaxTimeAsMetadata();
 
@@ -331,12 +378,14 @@ private:
     void eraseExcludedComputedMetaProjFields();
 
     BucketSpec _spec;
-    Behavior _unpackerBehavior;
 
     std::unique_ptr<UnpackingImpl> _unpackingImpl;
 
     bool _hasNext = false;
 
+    // A flag used to mark that the entire bucket matches the following $match predicate.
+    bool _bucketMatchedQuery = false;
+
     // A flag used to mark that the timestamp value should be materialized in measurements.
     bool _includeTimeField{false};
 
@@ -357,6 +406,8 @@ private:
     // measurement.
     Value _metaValue;
 
+    BSONElement _metaBSONElem;
+
     // Since the bucket min time is the same across all materialized measurements, we can cache the
     // value in the reset phase and use it to materialize as a metadata field in each measurement
     // if required by the pipeline.
@@ -383,9 +434,9 @@ private:
  * Determines if an arbitrary field should be included in the materialized measurements.
  */
 inline bool determineIncludeField(StringData fieldName,
-                                  BucketUnpacker::Behavior unpackerBehavior,
+                                  BucketSpec::Behavior unpackerBehavior,
                                   const std::set<std::string>& unpackFieldsToIncludeExclude) {
-    const bool isInclude = unpackerBehavior == BucketUnpacker::Behavior::kInclude;
+    const bool isInclude = unpackerBehavior == BucketSpec::Behavior::kInclude;
     const bool unpackFieldsContains = unpackFieldsToIncludeExclude.find(fieldName.toString()) !=
         unpackFieldsToIncludeExclude.cend();
     return isInclude == unpackFieldsContains;