SERVER-45176 apply "early return" optimization for normal projection path

3 files changed, 41 insertions, 0 deletions

diff --git a/src/mongo/db/exec/inclusion_projection_executor.h b/src/mongo/db/exec/inclusion_projection_executor.h
index 08dd86f7969..f17b1e9e259 100644
--- a/src/mongo/db/exec/inclusion_projection_executor.h
+++ b/src/mongo/db/exec/inclusion_projection_executor.h
@@ -71,6 +71,10 @@ public:
         }
     }
 
+    boost::optional<size_t> maxFieldsToProject() const override {
+        return _children.size() + _projectedFields.size();
+    }
+
 protected:
     // For inclusions, we can apply an optimization here by simply appending to the output document
     // via MutableDocument::addField, rather than always checking for existing fields via setField.
diff --git a/src/mongo/db/exec/projection_node.cpp b/src/mongo/db/exec/projection_node.cpp
index b66c43b0fc4..43e0f530b65 100644
--- a/src/mongo/db/exec/projection_node.cpp
+++ b/src/mongo/db/exec/projection_node.cpp
@@ -40,6 +40,7 @@ ProjectionNode::ProjectionNode(ProjectionPolicies policies, std::string pathToNo
     : _policies(policies), _pathToNode(std::move(pathToNode)) {}
 
 void ProjectionNode::addProjectionForPath(const FieldPath& path) {
+    makeOptimizationsStale();
     if (path.getPathLength() == 1) {
         _projectedFields.insert(path.fullPath());
         return;
@@ -50,6 +51,7 @@ void ProjectionNode::addProjectionForPath(const FieldPath& path) {
 
 void ProjectionNode::addExpressionForPath(const FieldPath& path,
                                           boost::intrusive_ptr<Expression> expr) {
+    makeOptimizationsStale();
     // If the computed fields policy is 'kBanComputedFields', we should never reach here.
     invariant(_policies.computedFieldsPolicy == ComputedFieldsPolicy::kAllowComputedFields);
 
@@ -84,11 +86,13 @@ boost::intrusive_ptr<Expression> ProjectionNode::getExpressionForPath(const Fiel
 }
 
 ProjectionNode* ProjectionNode::addOrGetChild(const std::string& field) {
+    makeOptimizationsStale();
     auto child = getChild(field);
     return child ? child : addChild(field);
 }
 
 ProjectionNode* ProjectionNode::addChild(const std::string& field) {
+    makeOptimizationsStale();
     invariant(!str::contains(field, "."));
     _orderToProcessAdditionsAndChildren.push_back(field);
     auto insertedPair = _children.emplace(std::make_pair(field, makeChild(field)));
@@ -119,6 +123,8 @@ Document ProjectionNode::applyToDocument(const Document& inputDoc) const {
 void ProjectionNode::applyProjections(const Document& inputDoc, MutableDocument* outputDoc) const {
     // Iterate over the input document so that the projected document retains its field ordering.
     auto it = inputDoc.fieldIterator();
+    size_t projectedFields = 0;
+
     while (it.more()) {
         auto fieldName = it.fieldName();
         absl::string_view fieldNameKey{fieldName.rawData(), fieldName.size()};
@@ -126,12 +132,19 @@ void ProjectionNode::applyProjections(const Document& inputDoc, MutableDocument*
         if (_projectedFields.find(fieldNameKey) != _projectedFields.end()) {
             outputProjectedField(
                 fieldName, applyLeafProjectionToValue(it.next().second), outputDoc);
+            ++projectedFields;
         } else if (auto childIt = _children.find(fieldNameKey); childIt != _children.end()) {
             outputProjectedField(
                 fieldName, childIt->second->applyProjectionsToValue(it.next().second), outputDoc);
+            ++projectedFields;
         } else {
             it.advance();
         }
+
+        // Check if we can avoid reading from the document any further.
+        if (_maxFieldsToProject && _maxFieldsToProject <= projectedFields) {
+            break;
+        }
     }
 
     // Ensure we project all specified fields, including those not present in the input document.
@@ -249,6 +262,8 @@ void ProjectionNode::optimize() {
     for (auto&& childPair : _children) {
         childPair.second->optimize();
     }
+
+    _maxFieldsToProject = maxFieldsToProject();
 }
 
 Document ProjectionNode::serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
diff --git a/src/mongo/db/exec/projection_node.h b/src/mongo/db/exec/projection_node.h
index 491611ab910..95ef83350cb 100644
--- a/src/mongo/db/exec/projection_node.h
+++ b/src/mongo/db/exec/projection_node.h
@@ -104,6 +104,14 @@ public:
     void reportProjectedPaths(std::set<std::string>* preservedPaths) const;
 
     /**
+     * Return an optional number, x, which indicates that it is safe to stop reading the document
+     * being projected once x fields have been projected.
+     */
+    virtual boost::optional<size_t> maxFieldsToProject() const {
+        return boost::none;
+    }
+
+    /**
      * Recursively reports all computed paths in this projection, adding them into 'computedPaths'.
      *
      * Computed paths that are identified as the result of a simple rename are instead filled out in
@@ -184,6 +192,15 @@ private:
     // Returns nullptr if no such child exists.
     ProjectionNode* getChild(const std::string& field) const;
 
+    /**
+     * Indicates that metadata computed by previous calls to optimize() is now stale and must be
+     * recomputed. This must be called any time the tree is updated (an expression added or child
+     * node added).
+     */
+    void makeOptimizationsStale() {
+        _maxFieldsToProject = boost::none;
+    }
+
     // Our projection semantics are such that all field additions need to be processed in the order
     // specified. '_orderToProcessAdditionsAndChildren' tracks that order.
     //
@@ -193,5 +210,10 @@ private:
     // '_expressions', and "b.c" will be tracked as a child ProjectionNode in '_children'. For the
     // example above, '_orderToProcessAdditionsAndChildren' would be ["a", "b", "d"].
     std::vector<std::string> _orderToProcessAdditionsAndChildren;
+
+    // Maximum number of fields that need to be projected. This allows for an "early" return
+    // optimization which means we don't have to iterate over an entire document. The value is
+    // stored here to avoid re-computation for each document.
+    boost::optional<size_t> _maxFieldsToProject;
 };
 }  // namespace mongo::projection_executor