Change ScriptFieldsFetchSubPhase to create search scripts once per segment (#26808)

Closes #26775
This commit is contained in:
Jim Ferenczi 2017-09-28 12:23:53 +02:00 committed by GitHub
parent 805437b8bc
commit f9f8856d89
2 changed files with 65 additions and 34 deletions

View File

@ -162,9 +162,15 @@ public class FetchPhase implements SearchPhase {
fetchSubPhase.hitExecute(context, hitContext);
}
}
if (context.isCancelled()) {
throw new TaskCancelledException("cancelled");
}
for (FetchSubPhase fetchSubPhase : fetchSubPhases) {
fetchSubPhase.hitsExecute(context, hits);
if (context.isCancelled()) {
throw new TaskCancelledException("cancelled");
}
}
context.fetchResult().hits(new SearchHits(hits, context.queryResult().getTotalHits(), context.queryResult().getMaxScore()));

View File

@ -18,62 +18,87 @@
*/
package org.elasticsearch.search.fetch.subphase;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.ReaderUtil;
import org.elasticsearch.common.document.DocumentField;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.fetch.FetchSubPhase;
import org.elasticsearch.search.internal.SearchContext;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
public final class ScriptFieldsFetchSubPhase implements FetchSubPhase {
@Override
public void hitExecute(SearchContext context, HitContext hitContext) {
public void hitsExecute(SearchContext context, SearchHit[] hits) throws IOException {
if (context.hasScriptFields() == false) {
return;
}
for (ScriptFieldsContext.ScriptField scriptField : context.scriptFields().fields()) {
/* Because this is called once per document we end up creating new ScriptDocValues for every document which is important because
* the values inside ScriptDocValues might be reused for different documents (Dates do this). */
SearchScript leafScript;
try {
leafScript = scriptField.script().newInstance(hitContext.readerContext());
} catch (IOException e1) {
throw new IllegalStateException("Failed to load script", e1);
}
leafScript.setDocument(hitContext.docId());
final Object value;
try {
value = leafScript.run();
} catch (RuntimeException e) {
if (scriptField.ignoreException()) {
continue;
hits = hits.clone(); // don't modify the incoming hits
Arrays.sort(hits, Comparator.comparingInt(SearchHit::docId));
int lastReaderId = -1;
SearchScript[] leafScripts = null;
List<ScriptFieldsContext.ScriptField> scriptFields = context.scriptFields().fields();
final IndexReader reader = context.searcher().getIndexReader();
for (SearchHit hit : hits) {
int readerId = ReaderUtil.subIndex(hit.docId(), reader.leaves());
LeafReaderContext leafReaderContext = reader.leaves().get(readerId);
if (readerId != lastReaderId) {
leafScripts = createLeafScripts(leafReaderContext, scriptFields);
lastReaderId = readerId;
}
int docId = hit.docId() - leafReaderContext.docBase;
for (int i = 0; i < leafScripts.length; i++) {
leafScripts[i].setDocument(docId);
final Object value;
try {
value = leafScripts[i].run();
} catch (RuntimeException e) {
if (scriptFields.get(i).ignoreException()) {
continue;
}
throw e;
}
throw e;
}
if (hitContext.hit().fieldsOrNull() == null) {
hitContext.hit().fields(new HashMap<>(2));
}
DocumentField hitField = hitContext.hit().getFields().get(scriptField.name());
if (hitField == null) {
final List<Object> values;
if (value instanceof Collection) {
// TODO: use diamond operator once JI-9019884 is fixed
values = new ArrayList<>((Collection<?>) value);
} else {
values = Collections.singletonList(value);
if (hit.fieldsOrNull() == null) {
hit.fields(new HashMap<>(2));
}
String scriptFieldName = scriptFields.get(i).name();
DocumentField hitField = hit.getFields().get(scriptFieldName);
if (hitField == null) {
final List<Object> values;
if (value instanceof Collection) {
values = new ArrayList<>((Collection<?>) value);
} else {
values = Collections.singletonList(value);
}
hitField = new DocumentField(scriptFieldName, values);
hit.getFields().put(scriptFieldName, hitField);
}
hitField = new DocumentField(scriptField.name(), values);
hitContext.hit().getFields().put(scriptField.name(), hitField);
}
}
}
private SearchScript[] createLeafScripts(LeafReaderContext context,
List<ScriptFieldsContext.ScriptField> scriptFields) {
SearchScript[] scripts = new SearchScript[scriptFields.size()];
for (int i = 0; i < scripts.length; i++) {
try {
scripts[i] = scriptFields.get(i).script().newInstance(context);
} catch (IOException e1) {
throw new IllegalStateException("Failed to load script " + scriptFields.get(i).name(), e1);
}
}
return scripts;
}
}