Move the shared fetch cache to highlighting. (#65105)

The cache is only used by highlighters, so it can be scoped to only the
highlighting context.
This commit is contained in:
Julie Tibshirani 2020-11-16 17:59:05 -08:00
parent 23a1c64347
commit 3974c3b066
11 changed files with 43 additions and 66 deletions

View File

@ -39,7 +39,6 @@ import org.elasticsearch.search.lookup.SourceLookup;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@ -104,9 +103,7 @@ final class PercolatorHighlightSubFetchPhase implements FetchSubPhase {
),
percolatorLeafReaderContext,
slot,
new SourceLookup(),
new HashMap<>()
);
new SourceLookup());
subContext.sourceLookup().setSource(document);
// force source because MemoryIndex does not store fields
SearchHighlightContext highlight = new SearchHighlightContext(fetchContext.highlight().fields(), true);

View File

@ -43,7 +43,6 @@ import org.elasticsearch.search.lookup.SourceLookup;
import org.elasticsearch.test.ESTestCase;
import java.util.Collections;
import java.util.HashMap;
import java.util.stream.IntStream;
import static org.mockito.Mockito.mock;
@ -69,9 +68,7 @@ public class PercolatorMatchedSlotSubFetchPhaseTests extends ESTestCase {
new SearchHit(0),
context,
0,
new SourceLookup(),
new HashMap<>()
);
new SourceLookup());
PercolateQuery.QueryStore queryStore = ctx -> docId -> new TermQuery(new Term("field", "value"));
MemoryIndex memoryIndex = new MemoryIndex();
memoryIndex.addField("field", "value", new WhitespaceAnalyzer());
@ -96,9 +93,7 @@ public class PercolatorMatchedSlotSubFetchPhaseTests extends ESTestCase {
new SearchHit(0),
context,
0,
new SourceLookup(),
new HashMap<>()
);
new SourceLookup());
PercolateQuery.QueryStore queryStore = ctx -> docId -> new TermQuery(new Term("field", "value"));
MemoryIndex memoryIndex = new MemoryIndex();
memoryIndex.addField("field", "value1", new WhitespaceAnalyzer());
@ -122,9 +117,7 @@ public class PercolatorMatchedSlotSubFetchPhaseTests extends ESTestCase {
new SearchHit(0),
context,
0,
new SourceLookup(),
new HashMap<>()
);
new SourceLookup());
PercolateQuery.QueryStore queryStore = ctx -> docId -> null;
MemoryIndex memoryIndex = new MemoryIndex();
memoryIndex.addField("field", "value", new WhitespaceAnalyzer());

View File

@ -119,7 +119,6 @@ public class FetchPhase {
FetchContext fetchContext = new FetchContext(context);
SearchHit[] hits = new SearchHit[context.docIdsToLoadSize()];
Map<String, Object> sharedCache = new HashMap<>();
List<FetchSubPhaseProcessor> processors = getProcessors(context.shardTarget(), fetchContext);
@ -160,9 +159,7 @@ public class FetchPhase {
docId,
storedToRequestedFields,
currentReaderContext,
fieldReader,
sharedCache
);
fieldReader);
for (FetchSubPhaseProcessor processor : processors) {
processor.process(hit);
}
@ -280,8 +277,7 @@ public class FetchPhase {
int docId,
Map<String, Set<String>> storedToRequestedFields,
LeafReaderContext subReaderContext,
CheckedBiConsumer<Integer, FieldsVisitor, IOException> storedFieldReader,
Map<String, Object> sharedCache) throws IOException {
CheckedBiConsumer<Integer, FieldsVisitor, IOException> storedFieldReader) throws IOException {
int rootDocId = findRootDocumentIfNested(context, subReaderContext, docId - subReaderContext.docBase);
if (rootDocId == -1) {
return prepareNonNestedHitContext(
@ -291,12 +287,9 @@ public class FetchPhase {
docId,
storedToRequestedFields,
subReaderContext,
storedFieldReader,
sharedCache
);
storedFieldReader);
} else {
return prepareNestedHitContext(context, docId, rootDocId, storedToRequestedFields,
subReaderContext, storedFieldReader, sharedCache);
return prepareNestedHitContext(context, docId, rootDocId, storedToRequestedFields, subReaderContext, storedFieldReader);
}
}
@ -308,20 +301,19 @@ public class FetchPhase {
* fetch subphases that use the hit context to access the preloaded source.
*/
private HitContext prepareNonNestedHitContext(SearchContext context,
SearchLookup lookup,
FieldsVisitor fieldsVisitor,
int docId,
Map<String, Set<String>> storedToRequestedFields,
LeafReaderContext subReaderContext,
CheckedBiConsumer<Integer, FieldsVisitor, IOException> fieldReader,
Map<String, Object> sharedCache) throws IOException {
SearchLookup lookup,
FieldsVisitor fieldsVisitor,
int docId,
Map<String, Set<String>> storedToRequestedFields,
LeafReaderContext subReaderContext,
CheckedBiConsumer<Integer, FieldsVisitor, IOException> fieldReader) throws IOException {
int subDocId = docId - subReaderContext.docBase;
DocumentMapper documentMapper = context.mapperService().documentMapper();
Text typeText = documentMapper.typeText();
if (fieldsVisitor == null) {
SearchHit hit = new SearchHit(docId, null, typeText, null, null);
return new HitContext(hit, subReaderContext, subDocId, lookup.source(), sharedCache);
return new HitContext(hit, subReaderContext, subDocId, lookup.source());
} else {
SearchHit hit;
loadStoredFields(context.mapperService(), fieldReader, fieldsVisitor, subDocId);
@ -335,7 +327,7 @@ public class FetchPhase {
hit = new SearchHit(docId, uid.id(), typeText, emptyMap(), emptyMap());
}
HitContext hitContext = new HitContext(hit, subReaderContext, subDocId, lookup.source(), sharedCache);
HitContext hitContext = new HitContext(hit, subReaderContext, subDocId, lookup.source());
if (fieldsVisitor.source() != null) {
hitContext.sourceLookup().setSource(fieldsVisitor.source());
}
@ -357,8 +349,8 @@ public class FetchPhase {
int rootDocId,
Map<String, Set<String>> storedToRequestedFields,
LeafReaderContext subReaderContext,
CheckedBiConsumer<Integer, FieldsVisitor, IOException> storedFieldReader,
Map<String, Object> sharedCache) throws IOException {
CheckedBiConsumer<Integer, FieldsVisitor, IOException> storedFieldReader)
throws IOException {
// Also if highlighting is requested on nested documents we need to fetch the _source from the root document,
// otherwise highlighting will attempt to fetch the _source from the nested doc, which will fail,
// because the entire _source is only stored with the root document.
@ -419,9 +411,7 @@ public class FetchPhase {
hit,
subReaderContext,
nestedDocId,
new SourceLookup(), // Use a clean, fresh SourceLookup for the nested context
sharedCache
);
new SourceLookup()); // Use a clean, fresh SourceLookup for the nested context
if (rootSourceAsMap != null) {
// Isolate the nested json array object that matches with nested hit and wrap it back into the same json

View File

@ -26,7 +26,6 @@ import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.lookup.SourceLookup;
import java.io.IOException;
import java.util.Map;
/**
* Sub phase within the fetch phase used to fetch things *about* the documents like highlighting or matched queries.
@ -38,21 +37,18 @@ public interface FetchSubPhase {
private final LeafReaderContext readerContext;
private final int docId;
private final SourceLookup sourceLookup;
private final Map<String, Object> cache;
public HitContext(
SearchHit hit,
LeafReaderContext context,
int docId,
SourceLookup sourceLookup,
Map<String, Object> cache
SourceLookup sourceLookup
) {
this.hit = hit;
this.readerContext = context;
this.docId = docId;
this.sourceLookup = sourceLookup;
sourceLookup.setSegmentAndDocument(context, docId);
this.cache = cache;
}
public SearchHit hit() {
@ -88,11 +84,6 @@ public interface FetchSubPhase {
public IndexReader topLevelReader() {
return ReaderUtil.getTopLevelContext(readerContext).reader();
}
// TODO move this into Highlighter
public Map<String, Object> cache() {
return cache;
}
}
/**

View File

@ -81,10 +81,10 @@ public class FastVectorHighlighter implements Highlighter {
Encoder encoder = field.fieldOptions().encoder().equals("html") ?
HighlightUtils.Encoders.HTML : HighlightUtils.Encoders.DEFAULT;
if (!hitContext.cache().containsKey(CACHE_KEY)) {
hitContext.cache().put(CACHE_KEY, new HighlighterEntry());
if (!fieldContext.cache.containsKey(CACHE_KEY)) {
fieldContext.cache.put(CACHE_KEY, new HighlighterEntry());
}
HighlighterEntry cache = (HighlighterEntry) hitContext.cache().get(CACHE_KEY);
HighlighterEntry cache = (HighlighterEntry) fieldContext.cache.get(CACHE_KEY);
FieldHighlightEntry entry = cache.fields.get(fieldType);
if (entry == null) {
FragListBuilder fragListBuilder;

View File

@ -23,6 +23,8 @@ import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.search.fetch.FetchContext;
import org.elasticsearch.search.fetch.FetchSubPhase;
import java.util.Map;
public class FieldHighlightContext {
public final String fieldName;
@ -32,6 +34,7 @@ public class FieldHighlightContext {
public final FetchSubPhase.HitContext hitContext;
public final Query query;
public final boolean forceSource;
public final Map<String, Object> cache;
public FieldHighlightContext(String fieldName,
SearchHighlightContext.Field field,
@ -39,7 +42,8 @@ public class FieldHighlightContext {
FetchContext context,
FetchSubPhase.HitContext hitContext,
Query query,
boolean forceSource) {
boolean forceSource,
Map<String, Object> cache) {
this.fieldName = fieldName;
this.field = field;
this.fieldType = fieldType;
@ -47,5 +51,6 @@ public class FieldHighlightContext {
this.hitContext = hitContext;
this.query = query;
this.forceSource = forceSource;
this.cache = cache;
}
}

View File

@ -56,7 +56,10 @@ public class HighlightPhase implements FetchSubPhase {
}
public FetchSubPhaseProcessor getProcessor(FetchContext context, SearchHighlightContext highlightContext, Query query) {
Map<String, Function<HitContext, FieldHighlightContext>> contextBuilders = contextBuilders(context, highlightContext, query);
Map<String, Object> sharedCache = new HashMap<>();
Map<String, Function<HitContext, FieldHighlightContext>> contextBuilders = contextBuilders(
context, highlightContext, query, sharedCache);
return new FetchSubPhaseProcessor() {
@Override
public void setNextReader(LeafReaderContext readerContext) {
@ -98,7 +101,8 @@ public class HighlightPhase implements FetchSubPhase {
private Map<String, Function<HitContext, FieldHighlightContext>> contextBuilders(FetchContext context,
SearchHighlightContext highlightContext,
Query query) {
Query query,
Map<String, Object> sharedCache) {
Map<String, Function<HitContext, FieldHighlightContext>> builders = new LinkedHashMap<>();
for (SearchHighlightContext.Field field : highlightContext.fields()) {
Highlighter highlighter = getHighlighter(field);
@ -147,7 +151,7 @@ public class HighlightPhase implements FetchSubPhase {
boolean forceSource = highlightContext.forceSource(field);
builders.put(fieldName,
hc -> new FieldHighlightContext(fieldType.name(), field, fieldType, context, hc,
highlightQuery == null ? query : highlightQuery, forceSource));
highlightQuery == null ? query : highlightQuery, forceSource, sharedCache));
}
}
return builders;

View File

@ -61,12 +61,12 @@ public class PlainHighlighter implements Highlighter {
Encoder encoder = field.fieldOptions().encoder().equals("html") ? HighlightUtils.Encoders.HTML : HighlightUtils.Encoders.DEFAULT;
if (!hitContext.cache().containsKey(CACHE_KEY)) {
hitContext.cache().put(CACHE_KEY, new HashMap<>());
if (!fieldContext.cache.containsKey(CACHE_KEY)) {
fieldContext.cache.put(CACHE_KEY, new HashMap<>());
}
@SuppressWarnings("unchecked")
Map<MappedFieldType, org.apache.lucene.search.highlight.Highlighter> cache =
(Map<MappedFieldType, org.apache.lucene.search.highlight.Highlighter>) hitContext.cache().get(CACHE_KEY);
(Map<MappedFieldType, org.apache.lucene.search.highlight.Highlighter>) fieldContext.cache.get(CACHE_KEY);
org.apache.lucene.search.highlight.Highlighter entry = cache.get(fieldType);
if (entry == null) {

View File

@ -62,7 +62,7 @@ public class UnifiedHighlighter implements Highlighter {
@Override
public HighlightField highlight(FieldHighlightContext fieldContext) throws IOException {
@SuppressWarnings("unchecked")
Map<String, CustomUnifiedHighlighter> cache = (Map<String, CustomUnifiedHighlighter>) fieldContext.hitContext.cache()
Map<String, CustomUnifiedHighlighter> cache = (Map<String, CustomUnifiedHighlighter>) fieldContext.cache
.computeIfAbsent(UnifiedHighlighter.class.getName(), k -> new HashMap<>());
if (cache.containsKey(fieldContext.fieldName) == false) {
cache.put(fieldContext.fieldName, buildHighlighter(fieldContext));

View File

@ -34,7 +34,6 @@ import org.elasticsearch.test.ESTestCase;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import static org.mockito.Mockito.mock;
@ -163,9 +162,7 @@ public class FetchSourcePhaseTests extends ESTestCase {
searchHit,
leafReaderContext,
1,
new SourceLookup(),
new HashMap<>()
);
new SourceLookup());
hitContext.sourceLookup().setSource(source == null ? null : BytesReference.bytes(source));
FetchSourcePhase phase = new FetchSourcePhase();

View File

@ -34,11 +34,11 @@ public class CustomHighlighter implements Highlighter {
@Override
public HighlightField highlight(FieldHighlightContext fieldContext) {
SearchHighlightContext.Field field = fieldContext.field;
CacheEntry cacheEntry = (CacheEntry) fieldContext.hitContext.cache().get("test-custom");
CacheEntry cacheEntry = (CacheEntry) fieldContext.cache.get("test-custom");
final int docId = fieldContext.hitContext.readerContext().docBase + fieldContext.hitContext.docId();
if (cacheEntry == null) {
cacheEntry = new CacheEntry();
fieldContext.hitContext.cache().put("test-custom", cacheEntry);
fieldContext.cache.put("test-custom", cacheEntry);
cacheEntry.docId = docId;
cacheEntry.position = 1;
} else {