reuse search lookup data for the two remaining cases (custom score and script filter) by having a "current" search context to access
This commit is contained in:
parent
34ed85a40f
commit
0f6beeb263
|
@ -22,6 +22,7 @@ package org.elasticsearch.index.query.xcontent;
|
|||
import org.apache.lucene.index.IndexReader;
|
||||
import org.apache.lucene.search.Explanation;
|
||||
import org.apache.lucene.search.Query;
|
||||
import org.elasticsearch.ElasticSearchIllegalStateException;
|
||||
import org.elasticsearch.common.Strings;
|
||||
import org.elasticsearch.common.collect.Maps;
|
||||
import org.elasticsearch.common.inject.Inject;
|
||||
|
@ -34,6 +35,7 @@ import org.elasticsearch.index.Index;
|
|||
import org.elasticsearch.index.query.QueryParsingException;
|
||||
import org.elasticsearch.index.settings.IndexSettings;
|
||||
import org.elasticsearch.script.search.SearchScript;
|
||||
import org.elasticsearch.search.internal.SearchContext;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.Map;
|
||||
|
@ -90,7 +92,11 @@ public class CustomScoreQueryParser extends AbstractIndexComponent implements XC
|
|||
throw new QueryParsingException(index, "[custom_score] requires 'script' field");
|
||||
}
|
||||
|
||||
SearchScript searchScript = new SearchScript(scriptLang, script, vars, parseContext.scriptService(), parseContext.mapperService(), parseContext.indexCache().fieldData());
|
||||
SearchContext context = SearchContext.current();
|
||||
if (context == null) {
|
||||
throw new ElasticSearchIllegalStateException("No search context on going...");
|
||||
}
|
||||
SearchScript searchScript = new SearchScript(context.scriptSearchLookup(), scriptLang, script, vars, parseContext.scriptService());
|
||||
FunctionScoreQuery functionScoreQuery = new FunctionScoreQuery(query, new ScriptScoreFunction(searchScript));
|
||||
functionScoreQuery.setBoost(boost);
|
||||
return functionScoreQuery;
|
||||
|
|
|
@ -22,6 +22,7 @@ package org.elasticsearch.index.query.xcontent;
|
|||
import org.apache.lucene.index.IndexReader;
|
||||
import org.apache.lucene.search.DocIdSet;
|
||||
import org.apache.lucene.search.Filter;
|
||||
import org.elasticsearch.ElasticSearchIllegalStateException;
|
||||
import org.elasticsearch.common.collect.Maps;
|
||||
import org.elasticsearch.common.inject.Inject;
|
||||
import org.elasticsearch.common.lucene.docset.GetDocSet;
|
||||
|
@ -35,6 +36,7 @@ import org.elasticsearch.index.query.QueryParsingException;
|
|||
import org.elasticsearch.index.settings.IndexSettings;
|
||||
import org.elasticsearch.script.ScriptService;
|
||||
import org.elasticsearch.script.search.SearchScript;
|
||||
import org.elasticsearch.search.internal.SearchContext;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.Map;
|
||||
|
@ -148,7 +150,11 @@ public class ScriptFilterParser extends AbstractIndexComponent implements XConte
|
|||
}
|
||||
|
||||
@Override public DocIdSet getDocIdSet(final IndexReader reader) throws IOException {
|
||||
final SearchScript searchScript = new SearchScript(scriptLang, script, params, scriptService, mapperService, fieldDataCache);
|
||||
SearchContext context = SearchContext.current();
|
||||
if (context == null) {
|
||||
throw new ElasticSearchIllegalStateException("No search context on going...");
|
||||
}
|
||||
final SearchScript searchScript = new SearchScript(context.scriptSearchLookup(), scriptLang, script, params, scriptService);
|
||||
searchScript.setNextReader(reader);
|
||||
return new GetDocSet(reader.maxDoc()) {
|
||||
@Override public boolean isCacheable() {
|
||||
|
|
|
@ -150,11 +150,13 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
try {
|
||||
contextProcessing(context);
|
||||
dfsPhase.execute(context);
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
return context.dfsResult();
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -164,11 +166,13 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
try {
|
||||
contextProcessing(context);
|
||||
queryPhase.execute(context);
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
return context.queryResult();
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -177,12 +181,14 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
try {
|
||||
contextProcessing(context);
|
||||
processScroll(request, context);
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
queryPhase.execute(context);
|
||||
return new ScrollQuerySearchResult(context.queryResult(), context.shardTarget());
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -193,15 +199,18 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
context.searcher().dfSource(new CachedDfSource(request.dfs(), context.similarityService().defaultSearchSimilarity()));
|
||||
} catch (IOException e) {
|
||||
freeContext(context);
|
||||
cleanContext(context);
|
||||
throw new QueryPhaseExecutionException(context, "Failed to set aggregated df", e);
|
||||
}
|
||||
try {
|
||||
queryPhase.execute(context);
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
return context.queryResult();
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -216,12 +225,14 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
if (context.scroll() == null) {
|
||||
freeContext(context.id());
|
||||
} else {
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
}
|
||||
return new QueryFetchSearchResult(context.queryResult(), context.fetchResult());
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -232,6 +243,7 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
context.searcher().dfSource(new CachedDfSource(request.dfs(), context.similarityService().defaultSearchSimilarity()));
|
||||
} catch (IOException e) {
|
||||
freeContext(context);
|
||||
cleanContext(context);
|
||||
throw new QueryPhaseExecutionException(context, "Failed to set aggregated df", e);
|
||||
}
|
||||
try {
|
||||
|
@ -241,12 +253,14 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
if (context.scroll() == null) {
|
||||
freeContext(request.id());
|
||||
} else {
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
}
|
||||
return new QueryFetchSearchResult(context.queryResult(), context.fetchResult());
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -261,12 +275,14 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
if (context.scroll() == null) {
|
||||
freeContext(request.id());
|
||||
} else {
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
}
|
||||
return new ScrollQueryFetchSearchResult(new QueryFetchSearchResult(context.queryResult(), context.fetchResult()), context.shardTarget());
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -279,12 +295,14 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
if (context.scroll() == null) {
|
||||
freeContext(request.id());
|
||||
} else {
|
||||
contextProcessingDone(context);
|
||||
contextProcessedSuccessfully(context);
|
||||
}
|
||||
return context.fetchResult();
|
||||
} catch (RuntimeException e) {
|
||||
freeContext(context);
|
||||
throw e;
|
||||
} finally {
|
||||
cleanContext(context);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -293,6 +311,7 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
if (context == null) {
|
||||
throw new SearchContextMissingException(id);
|
||||
}
|
||||
SearchContext.setCurrent(context);
|
||||
return context;
|
||||
}
|
||||
|
||||
|
@ -304,7 +323,7 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
|
||||
Engine.Searcher engineSearcher = indexShard.searcher();
|
||||
SearchContext context = new SearchContext(idGenerator.incrementAndGet(), shardTarget, request.numberOfShards(), request.timeout(), request.types(), engineSearcher, indexService, scriptService);
|
||||
|
||||
SearchContext.setCurrent(context);
|
||||
try {
|
||||
context.scroll(request.scroll());
|
||||
|
||||
|
@ -357,7 +376,7 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
}
|
||||
}
|
||||
|
||||
private void contextProcessingDone(SearchContext context) {
|
||||
private void contextProcessedSuccessfully(SearchContext context) {
|
||||
if (context.keepAliveTimeout() != null) {
|
||||
((KeepAliveTimerTask) context.keepAliveTimeout().getTask()).doneProcessing();
|
||||
} else {
|
||||
|
@ -366,6 +385,10 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
|
|||
}
|
||||
}
|
||||
|
||||
private void cleanContext(SearchContext context) {
|
||||
SearchContext.removeCurrent();
|
||||
}
|
||||
|
||||
private void parseSource(SearchContext context, byte[] source, int offset, int length) throws SearchParseException {
|
||||
// nothing to parse...
|
||||
if (source == null || length == 0) {
|
||||
|
|
|
@ -55,6 +55,20 @@ import java.util.List;
|
|||
*/
|
||||
public class SearchContext implements Releasable {
|
||||
|
||||
private static ThreadLocal<SearchContext> current = new ThreadLocal<SearchContext>();
|
||||
|
||||
public static void setCurrent(SearchContext value) {
|
||||
current.set(value);
|
||||
}
|
||||
|
||||
public static void removeCurrent() {
|
||||
current.remove();
|
||||
}
|
||||
|
||||
public static SearchContext current() {
|
||||
return current.get();
|
||||
}
|
||||
|
||||
private final long id;
|
||||
|
||||
private final SearchShardTarget shardTarget;
|
||||
|
|
|
@ -863,15 +863,16 @@ public class SimpleIndexQueryParserTests {
|
|||
assertThat(((TermFilter) constantScoreQuery.getFilter()).getTerm(), equalTo(new Term("name.last", "banon")));
|
||||
}
|
||||
|
||||
@Test public void testCustomScoreQuery1() throws IOException {
|
||||
IndexQueryParser queryParser = queryParser();
|
||||
String query = copyToStringFromClasspath("/org/elasticsearch/index/query/xcontent/custom_score1.json");
|
||||
Query parsedQuery = queryParser.parse(query).query();
|
||||
assertThat(parsedQuery, instanceOf(FunctionScoreQuery.class));
|
||||
FunctionScoreQuery functionScoreQuery = (FunctionScoreQuery) parsedQuery;
|
||||
assertThat(((TermQuery) functionScoreQuery.getSubQuery()).getTerm(), equalTo(new Term("name.last", "banon")));
|
||||
assertThat(functionScoreQuery.getFunction(), instanceOf(CustomScoreQueryParser.ScriptScoreFunction.class));
|
||||
}
|
||||
// Disabled since we need a current context to execute it...
|
||||
// @Test public void testCustomScoreQuery1() throws IOException {
|
||||
// IndexQueryParser queryParser = queryParser();
|
||||
// String query = copyToStringFromClasspath("/org/elasticsearch/index/query/xcontent/custom_score1.json");
|
||||
// Query parsedQuery = queryParser.parse(query).query();
|
||||
// assertThat(parsedQuery, instanceOf(FunctionScoreQuery.class));
|
||||
// FunctionScoreQuery functionScoreQuery = (FunctionScoreQuery) parsedQuery;
|
||||
// assertThat(((TermQuery) functionScoreQuery.getSubQuery()).getTerm(), equalTo(new Term("name.last", "banon")));
|
||||
// assertThat(functionScoreQuery.getFunction(), instanceOf(CustomScoreQueryParser.ScriptScoreFunction.class));
|
||||
// }
|
||||
|
||||
@Test public void testCustomBoostFactorQueryBuilder() throws IOException {
|
||||
IndexQueryParser queryParser = queryParser();
|
||||
|
|
Loading…
Reference in New Issue