From c14fdf63f43d649568c193241b4eb0addc94567a Mon Sep 17 00:00:00 2001
From: Alan Woodward <romseygeek@apache.org>
Date: Mon, 10 Jul 2023 12:44:55 +0100
Subject: [PATCH] Set new providers before building FetchSubPhaseProcessors
 (#97460) (#97500)

The FetchPhase sets source and field providers on its SearchLookup instance
to ensure that we only have to load stored fields once per document. However,
these were being set after subprocessors were constructed, and the
FetchFieldsPhase subprocessor pulls a reference to the SearchLookup at
construction time. This meant that the FieldsFetcher was not using these
providers, instead using the standard source and fields lookup, in effect meaning
that fields access would load source and fields separately.

This would mostly just be an issue with performance; however, the top hits
aggregation calls the fetch phase multiple times, with the result that the second
and subsequent buckets in a top hits aggregation would see the source lookup
from the previous bucket when building their FetchFieldsPhase subprocessor.

This commit moves the call to SearchExecutionContext#setLookupProviders()
in FetchPhase#buildSearchHits() before the call to build fetch phase subprocessors,
ensuring that the fetch fields phase sees the correct SearchLookup.

Fixes #96284
---
 docs/changelog/97460.yaml                     |  6 ++++
 .../test/aggregations/top_hits.yml            | 30 +++++++++++++++++++
 .../search/fetch/FetchPhase.java              |  8 ++---
 3 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 docs/changelog/97460.yaml

diff --git a/docs/changelog/97460.yaml b/docs/changelog/97460.yaml
new file mode 100644
index 0000000000000..55bcacaaded11
--- /dev/null
+++ b/docs/changelog/97460.yaml
@@ -0,0 +1,6 @@
+pr: 97460
+summary: Set new providers before building `FetchSubPhaseProcessors`
+area: "Search"
+type: bug
+issues:
+ - 96284
diff --git a/modules/aggregations/src/yamlRestTest/resources/rest-api-spec/test/aggregations/top_hits.yml b/modules/aggregations/src/yamlRestTest/resources/rest-api-spec/test/aggregations/top_hits.yml
index 14b2abac443fa..f23ce936e4c9a 100644
--- a/modules/aggregations/src/yamlRestTest/resources/rest-api-spec/test/aggregations/top_hits.yml
+++ b/modules/aggregations/src/yamlRestTest/resources/rest-api-spec/test/aggregations/top_hits.yml
@@ -720,3 +720,33 @@ synthetic _source:
                     buckets_path:
                       ts: top_salary_hits[salary]
                     script: "params.ts < 8000"
+
+---
+runtime fields:
+  - do:
+      search:
+        index: test
+        body:
+          runtime_mappings:
+            runtime_field:
+              type: keyword
+              script:
+                lang: painless
+                source: "emit(params['_source'].toString())"
+          fields: [runtime_field]
+          size: 0
+          aggs:
+            page:
+              terms:
+                field: page
+              aggs:
+                top_hits:
+                  top_hits: {}
+
+  - match:  { hits.total.value: 3 }
+  - length: { aggregations.page.buckets: 2 }
+  - match:  { aggregations.page.buckets.0.key: 1 }
+  - match:  { aggregations.page.buckets.0.top_hits.hits.hits.0.fields.runtime_field: ["{page=1, text=the quick brown fox}"] }
+  - match:  { aggregations.page.buckets.0.top_hits.hits.hits.1.fields.runtime_field: ["{page=1, text=jumped over the lazy dog}"] }
+  - match:  { aggregations.page.buckets.1.key: 2 }
+  - match:  { aggregations.page.buckets.1.top_hits.hits.hits.0.fields.runtime_field: ["{page=2, text=The vorpal blade went snicker-snack!}"] }
diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
index 405512551bf1d..c043e02e1d4e7 100644
--- a/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
+++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
@@ -99,6 +99,10 @@ private SearchHits buildSearchHits(SearchContext context, Profiler profiler) {
         FetchContext fetchContext = new FetchContext(context);
         SourceLoader sourceLoader = context.newSourceLoader();
 
+        PreloadedSourceProvider sourceProvider = new PreloadedSourceProvider();
+        PreloadedFieldLookupProvider fieldLookupProvider = new PreloadedFieldLookupProvider();
+        context.getSearchExecutionContext().setLookupProviders(sourceProvider, ctx -> fieldLookupProvider);
+
         List<FetchSubPhaseProcessor> processors = getProcessors(context.shardTarget(), fetchContext, profiler);
 
         StoredFieldsSpec storedFieldsSpec = StoredFieldsSpec.build(processors, FetchSubPhaseProcessor::storedFieldsSpec);
@@ -109,10 +113,6 @@ private SearchHits buildSearchHits(SearchContext context, Profiler profiler) {
 
         NestedDocuments nestedDocuments = context.getSearchExecutionContext().getNestedDocuments();
 
-        PreloadedSourceProvider sourceProvider = new PreloadedSourceProvider();
-        PreloadedFieldLookupProvider fieldLookupProvider = new PreloadedFieldLookupProvider();
-        context.getSearchExecutionContext().setLookupProviders(sourceProvider, ctx -> fieldLookupProvider);
-
         FetchPhaseDocsIterator docsIterator = new FetchPhaseDocsIterator() {
 
             LeafReaderContext ctx;