iter

javanna · Apr 17, 2024 · 5f4cc87 · 5f4cc87
1 parent 411fbd5
commit 5f4cc87
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 14 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
@@ -118,7 +118,7 @@ private SearchHits buildSearchHits(SearchContext context, int[] docIdsToLoad, Pr
         // the getProcessors call to before the setLookupProviders call, which causes weird issues in InnerHitsPhase.
         // setLookupProviders resets the SearchLookup used throughout the rest of the fetch phase, which StoredValueFetchers rely on
         // to retrieve stored fields, and InnerHitsPhase is the last sub-fetch phase and re-runs the entire fetch phase.
-        fieldLookupProvider.setPreloadedStoredFields(storedFieldsSpec.requiredStoredFields());
+        fieldLookupProvider.setPreloadedStoredFieldNames(storedFieldsSpec.requiredStoredFields());
 
         StoredFieldLoader storedFieldLoader = profiler.storedFields(StoredFieldLoader.fromSpec(storedFieldsSpec));
         IdLoader idLoader = context.newIdLoader();
@@ -168,7 +168,7 @@ protected SearchHit nextDoc(int doc) throws IOException {
                     leafIdLoader
                 );
                 sourceProvider.source = hit.source();
-                fieldLookupProvider.setStoredFields(hit.hit().getId(), hit.loadedFields());
+                fieldLookupProvider.setPreloadedStoredFieldValues(hit.hit().getId(), hit.loadedFields());
                 for (FetchSubPhaseProcessor processor : processors) {
                     processor.process(hit);
                 }

diff --git a/server/src/main/java/org/elasticsearch/search/fetch/PreloadedFieldLookupProvider.java b/server/src/main/java/org/elasticsearch/search/fetch/PreloadedFieldLookupProvider.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.search.fetch;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.index.mapper.IdFieldMapper;
 import org.elasticsearch.search.lookup.FieldLookup;
 import org.elasticsearch.search.lookup.LeafFieldLookupProvider;
@@ -29,8 +30,8 @@
  */
 class PreloadedFieldLookupProvider implements LeafFieldLookupProvider {
 
-    private Set<String> preloadedStoredFields;
-    private Map<String, List<Object>> storedFields;
+    private final SetOnce<Set<String>> preloadedStoredFieldNames = new SetOnce<>();
+    private Map<String, List<Object>> preloadedStoredFieldValues;
     private String id;
     private LeafFieldLookupProvider backUpLoader;
     private Supplier<LeafFieldLookupProvider> loaderSupplier;
@@ -43,8 +44,8 @@ public void populateFieldLookup(FieldLookup fieldLookup, int doc) throws IOExcep
             fieldLookup.setValues(Collections.singletonList(id));
             return;
         }
-        if (preloadedStoredFields.contains(field)) {
-            fieldLookup.setValues(storedFields.get(field));
+        if (preloadedStoredFieldNames.get().contains(field)) {
+            fieldLookup.setValues(preloadedStoredFieldValues.get(field));
             return;
         }
         // stored field not preloaded, go and get it directly
@@ -54,14 +55,17 @@ public void populateFieldLookup(FieldLookup fieldLookup, int doc) throws IOExcep
         backUpLoader.populateFieldLookup(fieldLookup, doc);
     }
 
-    void setPreloadedStoredFields(Set<String> preloadedStoredFields) {
-        this.preloadedStoredFields = preloadedStoredFields;
+    void setPreloadedStoredFieldNames(Set<String> preloadedStoredFieldNames) {
+        this.preloadedStoredFieldNames.set(preloadedStoredFieldNames);
     }
 
-    void setStoredFields(String id, Map<String, List<Object>> storedFields) {
-        assert preloadedStoredFields.containsAll(storedFields.keySet())
-            : "Provided stored field that was not expected to be preloaded? " + storedFields.keySet() + " - " + preloadedStoredFields;
-        this.storedFields = storedFields;
+    void setPreloadedStoredFieldValues(String id, Map<String, List<Object>> preloadedStoredFieldValues) {
+        assert preloadedStoredFieldNames.get().containsAll(preloadedStoredFieldValues.keySet())
+            : "Provided stored field that was not expected to be preloaded? "
+                + preloadedStoredFieldValues.keySet()
+                + " - "
+                + preloadedStoredFieldNames;
+        this.preloadedStoredFieldValues = preloadedStoredFieldValues;
         this.id = id;
     }
 

diff --git a/server/src/test/java/org/elasticsearch/search/fetch/PreloadedFieldLookupProviderTests.java b/server/src/test/java/org/elasticsearch/search/fetch/PreloadedFieldLookupProviderTests.java
@@ -32,8 +32,8 @@ public class PreloadedFieldLookupProviderTests extends ESTestCase {
 
     public void testFallback() throws IOException {
         PreloadedFieldLookupProvider lookup = new PreloadedFieldLookupProvider();
-        lookup.setPreloadedStoredFields(Collections.singleton("foo"));
-        lookup.setStoredFields("id", Map.of("foo", List.of("bar")));
+        lookup.setPreloadedStoredFieldNames(Collections.singleton("foo"));
+        lookup.setPreloadedStoredFieldValues("id", Map.of("foo", List.of("bar")));
 
         MappedFieldType idFieldType = mock(MappedFieldType.class);
         when(idFieldType.name()).thenReturn(IdFieldMapper.NAME);