Speed up empty highlighting many fields (#61860)

Kibana often highlights *everything* like this:
```
POST /_search
{
  "query": ...,
  "size": 500,
  "highlight": {
    "fields": {
      "*": { ... }
    }
  }
}
```

This can get slow when there are hundreds of mapped fields. I tested
this locally and unscientifically and it took a request from 20ms to
150ms when there are 100 fields. I've seen clusters with 2000 fields
where simple search go from 500ms to 1500ms just by turning on this sort
of highlighting. Even when the query is just a `range` that and the
fields are all numbers and stuff so it won't highlight anything.

This speeds up the `unified` highlighter in this case in a few ways:
1. Build the highlighting infrastructure once field rather than once pre
   document per field. This cuts out a *ton* of work analyzing the query
   over and over and over again.
2. Bail out of the highlighter before loading values if we can't produce
   any results.

Combined these take that local 150ms case down to 65ms. This is unlikely
to be really useful when there are only a few fetched docs and only a
few fields, but we often end up having many fields with many fetched
docs.

Author: nik9000

plugins/mapper-annotated-text/src/main/java/org/elasticsearch/index/mapper/annotatedtext/AnnotatedTextFieldMapper.java

@@ -46,7 +46,6 @@
 import org.elasticsearch.index.mapper.ValueFetcher;
 import org.elasticsearch.index.mapper.annotatedtext.AnnotatedTextFieldMapper.AnnotatedText.AnnotationToken;
 import org.elasticsearch.index.similarity.SimilarityProvider;
-import org.elasticsearch.search.fetch.FetchSubPhase.HitContext;
 
 import java.io.IOException;
 import java.io.Reader;
@@ -299,22 +298,25 @@ public AnnotationToken getAnnotation(int index) {
     // original markup form in order to inject annotations.
     public static final class AnnotatedHighlighterAnalyzer extends AnalyzerWrapper {
         private final Analyzer delegate;
-        private final HitContext hitContext;
-        public AnnotatedHighlighterAnalyzer(Analyzer delegate, HitContext hitContext){
+        private AnnotatedText[] annotations;
+
+        public AnnotatedHighlighterAnalyzer(Analyzer delegate){
             super(delegate.getReuseStrategy());
             this.delegate = delegate;
-            this.hitContext = hitContext;
         }
 
         @Override
         public Analyzer getWrappedAnalyzer(String fieldName) {
           return delegate;
         }
 
+        public void setAnnotations(AnnotatedText[] annotations) {
+            this.annotations = annotations;
+        }
+
         @Override
         protected TokenStreamComponents wrapComponents(String fieldName, TokenStreamComponents components) {
             AnnotationsInjector injector = new AnnotationsInjector(components.getTokenStream());
-            AnnotatedText[] annotations = (AnnotatedText[]) hitContext.cache().get(AnnotatedText.class.getName());
             AtomicInteger readerNum = new AtomicInteger(0);
             return new TokenStreamComponents(r -> {
                 String plainText = readToString(r);

plugins/mapper-annotated-text/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/AnnotatedPassageFormatter.java

@@ -44,8 +44,11 @@ public class AnnotatedPassageFormatter extends PassageFormatter {
     private final Encoder encoder;
     AnnotatedText[] annotations;
 
-    public AnnotatedPassageFormatter(AnnotatedText[] annotations, Encoder encoder) {
+    public AnnotatedPassageFormatter(Encoder encoder) {
         this.encoder = encoder;
+    }
+
+    void setAnnotations(AnnotatedText[] annotations) {
         this.annotations = annotations;
     }
 

plugins/mapper-annotated-text/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/AnnotatedTextHighlighter.java

@@ -21,6 +21,7 @@
 
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.search.highlight.Encoder;
+import org.apache.lucene.search.uhighlight.CustomUnifiedHighlighter;
 import org.apache.lucene.search.uhighlight.PassageFormatter;
 import org.elasticsearch.index.mapper.DocumentMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
@@ -37,39 +38,37 @@ public class AnnotatedTextHighlighter extends UnifiedHighlighter {
 
     public static final String NAME = "annotated";
 
-    @Override
-    protected Analyzer getAnalyzer(DocumentMapper docMapper, HitContext hitContext) {
-        return new AnnotatedHighlighterAnalyzer(super.getAnalyzer(docMapper, hitContext), hitContext);
-    }
-
     // Convert the marked-up values held on-disk to plain-text versions for highlighting
     @Override
-    protected List<Object> loadFieldValues(MappedFieldType fieldType,
-                                           Field field,
-                                           HitContext hitContext,
-                                           boolean forceSource) throws IOException {
-        List<Object> fieldValues = super.loadFieldValues(fieldType, field, hitContext, forceSource);
-        String[] fieldValuesAsString = fieldValues.toArray(new String[fieldValues.size()]);
+    protected List<Object> loadFieldValues(
+        CustomUnifiedHighlighter highlighter,
+        MappedFieldType fieldType,
+        Field field,
+        HitContext hitContext,
+        boolean forceSource
+    ) throws IOException {
+        List<Object> fieldValues = super.loadFieldValues(highlighter, fieldType, field, hitContext, forceSource);
 
-        AnnotatedText[] annotations = new AnnotatedText[fieldValuesAsString.length];
-        for (int i = 0; i < fieldValuesAsString.length; i++) {
-            annotations[i] = AnnotatedText.parse(fieldValuesAsString[i]);
+        List<Object> strings = new ArrayList<>(fieldValues.size());
+        AnnotatedText[] annotations = new AnnotatedText[fieldValues.size()];
+        for (int i = 0; i < fieldValues.size(); i++) {
+            annotations[i] = AnnotatedText.parse(fieldValues.get(i).toString());
+            strings.add(annotations[i].textMinusMarkup);
         }
-        // Store the annotations in the hitContext
-        hitContext.cache().put(AnnotatedText.class.getName(), annotations);
+        // Store the annotations in the formatter and analyzer
+        ((AnnotatedPassageFormatter) highlighter.getFormatter()).setAnnotations(annotations);
+        ((AnnotatedHighlighterAnalyzer) highlighter.getIndexAnalyzer()).setAnnotations(annotations);
+        return strings;
+    }
 
-        ArrayList<Object> result = new ArrayList<>(annotations.length);
-        for (int i = 0; i < annotations.length; i++) {
-            result.add(annotations[i].textMinusMarkup);
-        }
-        return result;
+    @Override
+    protected Analyzer getAnalyzer(DocumentMapper docMapper) {
+        return new AnnotatedHighlighterAnalyzer(super.getAnalyzer(docMapper));
     }
 
     @Override
     protected PassageFormatter getPassageFormatter(HitContext hitContext, SearchHighlightContext.Field field, Encoder encoder) {
-        // Retrieve the annotations from the hitContext
-        AnnotatedText[] annotations = (AnnotatedText[]) hitContext.cache().get(AnnotatedText.class.getName());
-        return new AnnotatedPassageFormatter(annotations, encoder);
+        return new AnnotatedPassageFormatter(encoder);
     }
 
 }

plugins/mapper-annotated-text/src/test/java/org/elasticsearch/search/highlight/AnnotatedTextHighlighterTests.java renamed to plugins/mapper-annotated-text/src/test/java/org/elasticsearch/search/fetch/subphase/highlight/AnnotatedTextHighlighterTests.java

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-package org.elasticsearch.search.highlight;
+package org.elasticsearch.search.fetch.subphase.highlight;
 
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.analysis.standard.StandardAnalyzer;
@@ -28,7 +28,6 @@
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexOptions;
 import org.apache.lucene.index.IndexWriterConfig;
-import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.RandomIndexWriter;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.IndexSearcher;
@@ -48,14 +47,11 @@
 import org.elasticsearch.index.mapper.annotatedtext.AnnotatedTextFieldMapper.AnnotatedHighlighterAnalyzer;
 import org.elasticsearch.index.mapper.annotatedtext.AnnotatedTextFieldMapper.AnnotatedText;
 import org.elasticsearch.index.mapper.annotatedtext.AnnotatedTextFieldMapper.AnnotationAnalyzerWrapper;
-import org.elasticsearch.search.fetch.FetchSubPhase.HitContext;
-import org.elasticsearch.search.fetch.subphase.highlight.AnnotatedPassageFormatter;
 import org.elasticsearch.test.ESTestCase;
 
 import java.net.URLEncoder;
 import java.text.BreakIterator;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.Locale;
 
 import static org.apache.lucene.search.uhighlight.CustomUnifiedHighlighter.MULTIVAL_SEP_CHAR;
@@ -70,7 +66,6 @@ private void assertHighlightOneDoc(String fieldName, String []markedUpInputs,
 
         // Annotated fields wrap the usual analyzer with one that injects extra tokens
         Analyzer wrapperAnalyzer = new AnnotationAnalyzerWrapper(new StandardAnalyzer());
-
         Directory dir = newDirectory();
         IndexWriterConfig iwc = newIndexWriterConfig(wrapperAnalyzer);
         iwc.setMergePolicy(newTieredMergePolicy(random()));
@@ -93,17 +88,14 @@ private void assertHighlightOneDoc(String fieldName, String []markedUpInputs,
         IndexSearcher searcher = newSearcher(reader);
         iw.close();
 
-        LeafReaderContext context = searcher.getIndexReader().leaves().get(0);
-        HitContext mockHitContext = new HitContext(null, context, 0, null, new HashMap<>());
-        AnnotatedHighlighterAnalyzer hiliteAnalyzer = new AnnotatedHighlighterAnalyzer(wrapperAnalyzer, mockHitContext);
-
         AnnotatedText[] annotations = new AnnotatedText[markedUpInputs.length];
         for (int i = 0; i < markedUpInputs.length; i++) {
             annotations[i] = AnnotatedText.parse(markedUpInputs[i]);
         }
-        mockHitContext.cache().put(AnnotatedText.class.getName(), annotations);
-
-        AnnotatedPassageFormatter passageFormatter = new AnnotatedPassageFormatter(annotations,new DefaultEncoder());
+        AnnotatedHighlighterAnalyzer hiliteAnalyzer = new AnnotatedHighlighterAnalyzer(wrapperAnalyzer);
+        hiliteAnalyzer.setAnnotations(annotations);
+        AnnotatedPassageFormatter passageFormatter = new AnnotatedPassageFormatter(new DefaultEncoder());
+        passageFormatter.setAnnotations(annotations);
 
         ArrayList<Object> plainTextForHighlighter = new ArrayList<>(annotations.length);
         for (int i = 0; i < annotations.length; i++) {
@@ -113,13 +105,24 @@ private void assertHighlightOneDoc(String fieldName, String []markedUpInputs,
         TopDocs topDocs = searcher.search(new MatchAllDocsQuery(), 1, Sort.INDEXORDER);
         assertThat(topDocs.totalHits.value, equalTo(1L));
         String rawValue = Strings.collectionToDelimitedString(plainTextForHighlighter, String.valueOf(MULTIVAL_SEP_CHAR));
-
-        CustomUnifiedHighlighter highlighter = new CustomUnifiedHighlighter(searcher, hiliteAnalyzer, null,
-                passageFormatter, locale,
-                breakIterator, rawValue, noMatchSize);
+        CustomUnifiedHighlighter highlighter = new CustomUnifiedHighlighter(
+            searcher,
+            hiliteAnalyzer,
+            null,
+            passageFormatter,
+            locale,
+            breakIterator,
+            "index",
+            "text",
+            query,
+            noMatchSize,
+            expectedPassages.length,
+            name -> "text".equals(name),
+            Integer.MAX_VALUE,
+            Integer.MAX_VALUE
+        );
         highlighter.setFieldMatcher((name) -> "text".equals(name));
-        final Snippet[] snippets =
-            highlighter.highlightField("text", query, topDocs.scoreDocs[0].doc, expectedPassages.length);
+        final Snippet[] snippets = highlighter.highlightField(getOnlyLeafReader(reader), topDocs.scoreDocs[0].doc, () -> rawValue);
         assertEquals(expectedPassages.length, snippets.length);
         for (int i = 0; i < snippets.length; i++) {
             assertEquals(expectedPassages[i], snippets[i].getText());

server/src/main/java/org/apache/lucene/search/uhighlight/CustomFieldHighlighter.java

@@ -19,6 +19,7 @@
 
 package org.apache.lucene.search.uhighlight;
 
+import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.util.BytesRef;
 
 import java.io.IOException;
@@ -39,17 +40,27 @@ class CustomFieldHighlighter extends FieldHighlighter {
 
     private final Locale breakIteratorLocale;
     private final int noMatchSize;
-    private final String fieldValue;
+    private String fieldValue;
 
     CustomFieldHighlighter(String field, FieldOffsetStrategy fieldOffsetStrategy,
                            Locale breakIteratorLocale, BreakIterator breakIterator,
                            PassageScorer passageScorer, int maxPassages, int maxNoHighlightPassages,
-                           PassageFormatter passageFormatter, int noMatchSize, String fieldValue) {
+                           PassageFormatter passageFormatter, int noMatchSize) {
         super(field, fieldOffsetStrategy, breakIterator, passageScorer, maxPassages,
             maxNoHighlightPassages, passageFormatter);
         this.breakIteratorLocale = breakIteratorLocale;
         this.noMatchSize = noMatchSize;
-        this.fieldValue = fieldValue;
+    }
+
+    @Override
+    public Object highlightFieldForDoc(LeafReader reader, int docId, String content) throws IOException {
+        this.fieldValue = content;
+        try {
+            return super.highlightFieldForDoc(reader, docId, content);
+        } finally {
+            // Clear the reference to the field value in case it is large
+            fieldValue = null;
+        }
     }
 
     @Override