SERVER-63845 Separate variable reference tracking from pipeline field dependency analysis

1 files changed, 24 insertions, 13 deletions

diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index b258c0f57c1..633b7a2f01c 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -41,6 +41,7 @@
 #include "mongo/db/pipeline/document_source_sort.h"
 #include "mongo/db/pipeline/expression.h"
 #include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/pipeline/expression_dependencies.h"
 #include "mongo/db/pipeline/sort_reorder_helpers.h"
 #include "mongo/db/pipeline/variable_validation.h"
 #include "mongo/db/query/collation/collator_factory_interface.h"
@@ -1067,20 +1068,9 @@ DepsTracker::State DocumentSourceLookUp::getDependencies(DepsTracker* deps) cons
             source->getDependencies(&subDeps);
         }
 
-        // Add the 'let' dependencies to the tracker. Because the caller is only interested in
-        // references to external variables, filter out any subpipeline references to 'let'
-        // variables declared by this $lookup.
+        // Add the 'let' dependencies to the tracker.
         for (auto&& letVar : _letVariables) {
-            letVar.expression->addDependencies(deps);
-            subDeps.vars.erase(letVar.id);
-        }
-
-        // Add sub-pipeline variable dependencies. Do not add field dependencies, since these refer
-        // to the fields from the foreign collection rather than the local collection.
-        // Similarly, do not add SEARCH_META as a dependency, since it is scoped to one pipeline.
-        for (auto&& varId : subDeps.vars) {
-            if (varId != Variables::kSearchMetaId)
-                deps->vars.insert(varId);
+            expression::addDependencies(letVar.expression.get(), deps);
         }
     }
 
@@ -1094,6 +1084,27 @@ DepsTracker::State DocumentSourceLookUp::getDependencies(DepsTracker* deps) cons
     return DepsTracker::State::SEE_NEXT;
 }
 
+void DocumentSourceLookUp::addVariableRefs(std::set<Variables::Id>* refs) const {
+    // Do not add SEARCH_META as a reference, since it is scoped to one pipeline.
+    if (hasPipeline()) {
+        std::set<Variables::Id> subPipeRefs;
+        _resolvedIntrospectionPipeline->addVariableRefs(&subPipeRefs);
+        for (auto&& varId : subPipeRefs) {
+            if (varId != Variables::kSearchMetaId)
+                refs->insert(varId);
+        }
+    }
+
+    // Add the 'let' variable references. Because the caller is only interested in references to
+    // external variables, filter out any subpipeline references to 'let' variables declared by this
+    // $lookup. This step must happen after gathering the sub-pipeline variable references as they
+    // may refer to let variables.
+    for (auto&& letVar : _letVariables) {
+        expression::addVariableRefs(letVar.expression.get(), refs);
+        refs->erase(letVar.id);
+    }
+}
+
 boost::optional<DocumentSource::DistributedPlanLogic> DocumentSourceLookUp::distributedPlanLogic() {
     // If $lookup into a sharded foreign collection is allowed and the foreign namespace is sharded,
     // top-level $lookup stages can run in parallel on the shards.