Skip to content

Commit e3a9271

Browse files
kimchykimchy
committed
unify more count and search implementation
1 parent bc28873 commit e3a9271

File tree

16 files changed

+113
-133
lines changed

16 files changed

+113
-133
lines changed

‎src/main/java/org/elasticsearch/action/count/TransportCountAction.java

Lines changed: 48 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,18 @@
3030
import org.elasticsearch.cluster.block.ClusterBlockLevel;
3131
import org.elasticsearch.cluster.routing.GroupShardsIterator;
3232
import org.elasticsearch.cluster.routing.ShardRouting;
33+
import org.elasticsearch.common.bytes.BytesReference;
3334
import org.elasticsearch.common.inject.Inject;
35+
import org.elasticsearch.common.lucene.Lucene;
3436
import org.elasticsearch.common.settings.Settings;
37+
import org.elasticsearch.index.query.QueryParseContext;
38+
import org.elasticsearch.index.service.IndexService;
3539
import org.elasticsearch.index.shard.service.IndexShard;
3640
import org.elasticsearch.indices.IndicesService;
41+
import org.elasticsearch.script.ScriptService;
42+
import org.elasticsearch.search.internal.InternalSearchRequest;
43+
import org.elasticsearch.search.internal.SearchContext;
44+
import org.elasticsearch.search.query.QueryPhaseExecutionException;
3745
import org.elasticsearch.threadpool.ThreadPool;
3846
import org.elasticsearch.transport.TransportService;
3947

@@ -51,10 +59,14 @@ public class TransportCountAction extends TransportBroadcastOperationAction<Coun
5159

5260
private final IndicesService indicesService;
5361

62+
private final ScriptService scriptService;
63+
5464
@Inject
55-
public TransportCountAction(Settings settings, ThreadPool threadPool, ClusterService clusterService, TransportService transportService, IndicesService indicesService) {
65+
public TransportCountAction(Settings settings, ThreadPool threadPool, ClusterService clusterService, TransportService transportService,
66+
IndicesService indicesService, ScriptService scriptService) {
5667
super(settings, threadPool, clusterService, transportService);
5768
this.indicesService = indicesService;
69+
this.scriptService = scriptService;
5870
}
5971

6072
@Override
@@ -130,9 +142,40 @@ protected CountResponse newResponse(CountRequest request, AtomicReferenceArray s
130142

131143
@Override
132144
protected ShardCountResponse shardOperation(ShardCountRequest request) throws ElasticSearchException {
133-
IndexShard indexShard = indicesService.indexServiceSafe(request.index()).shardSafe(request.shardId());
134-
long count = indexShard.count(request.minScore(), request.querySource(),
135-
request.filteringAliases(), request.types());
136-
return new ShardCountResponse(request.index(), request.shardId(), count);
145+
IndexService indexService = indicesService.indexServiceSafe(request.index());
146+
IndexShard indexShard = indexService.shardSafe(request.shardId());
147+
148+
SearchContext context = new SearchContext(0,
149+
new InternalSearchRequest().types(request.types()).filteringAliases(request.filteringAliases()),
150+
null, indexShard.searcher(), indexService, indexShard,
151+
scriptService);
152+
SearchContext.setCurrent(context);
153+
154+
try {
155+
// TODO: min score should move to be "null" as a value that is not initialized...
156+
if (request.minScore() != -1) {
157+
context.minimumScore(request.minScore());
158+
}
159+
BytesReference querySource = request.querySource();
160+
if (querySource != null && querySource.length() > 0) {
161+
try {
162+
QueryParseContext.setTypes(request.types());
163+
context.parsedQuery(indexService.queryParserService().parse(querySource));
164+
} finally {
165+
QueryParseContext.removeTypes();
166+
}
167+
}
168+
context.preProcess();
169+
try {
170+
long count = Lucene.count(context.searcher(), context.query());
171+
return new ShardCountResponse(request.index(), request.shardId(), count);
172+
} catch (Exception e) {
173+
throw new QueryPhaseExecutionException(context, "failed to execute count", e);
174+
}
175+
} finally {
176+
// this will also release the index searcher
177+
context.release();
178+
SearchContext.removeCurrent();
179+
}
137180
}
138181
}

‎src/main/java/org/elasticsearch/action/support/DefaultShardOperationFailedException.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,4 +96,9 @@ public void writeTo(StreamOutput out) throws IOException {
9696
out.writeVInt(shardId);
9797
out.writeUTF(reason);
9898
}
99+
100+
@Override
101+
public String toString() {
102+
return "[" + index + "][" + shardId + "] failed, reason [" + reason + "]";
103+
}
99104
}

‎src/main/java/org/elasticsearch/common/lucene/Lucene.java

Lines changed: 8 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -82,14 +82,14 @@ public static Version parseVersion(@Nullable String version, Version defaultVers
8282
return defaultVersion;
8383
}
8484

85-
public static long count(IndexSearcher searcher, Query query, float minScore) throws IOException {
86-
return count(searcher, query, null, minScore);
87-
}
88-
89-
public static long count(IndexSearcher searcher, Query query, Filter filter, float minScore) throws IOException {
90-
CountCollector countCollector = new CountCollector(minScore);
91-
searcher.search(query, filter, countCollector);
92-
return countCollector.count();
85+
public static long count(IndexSearcher searcher, Query query) throws IOException {
86+
TotalHitCountCollector countCollector = new TotalHitCountCollector();
87+
// we don't need scores, so wrap it in a constant score query
88+
if (!(query instanceof ConstantScoreQuery)) {
89+
query = new ConstantScoreQuery(query);
90+
}
91+
searcher.search(query, countCollector);
92+
return countCollector.getTotalHits();
9393
}
9494

9595
public static int docId(IndexReader reader, Term term) throws IOException {
@@ -320,42 +320,6 @@ public static SegmentInfo getSegmentInfo(SegmentReader reader) {
320320
}
321321
}
322322

323-
public static class CountCollector extends Collector {
324-
325-
private final float minScore;
326-
private Scorer scorer;
327-
private long count;
328-
329-
public CountCollector(float minScore) {
330-
this.minScore = minScore;
331-
}
332-
333-
public long count() {
334-
return this.count;
335-
}
336-
337-
@Override
338-
public void setScorer(Scorer scorer) throws IOException {
339-
this.scorer = scorer;
340-
}
341-
342-
@Override
343-
public void collect(int doc) throws IOException {
344-
if (scorer.score() > minScore) {
345-
count++;
346-
}
347-
}
348-
349-
@Override
350-
public void setNextReader(IndexReader reader, int docBase) throws IOException {
351-
}
352-
353-
@Override
354-
public boolean acceptsDocsOutOfOrder() {
355-
return true;
356-
}
357-
}
358-
359323
public static class ExistsCollector extends Collector {
360324

361325
private boolean exists;

‎src/main/java/org/elasticsearch/common/lucene/search/Queries.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,9 @@ public static Query fixNegativeQueryIfNeeded(Query q) {
110110
}
111111

112112
public static boolean isMatchAllQuery(Query query) {
113+
if (query == Queries.MATCH_ALL_QUERY) {
114+
return true;
115+
}
113116
if (query instanceof MatchAllDocsQuery) {
114117
return true;
115118
}

‎src/main/java/org/elasticsearch/index/mapper/MapperService.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,7 @@ public DocumentMapper documentMapperWithAutoCreate(String type) {
398398
/**
399399
* A filter for search. If a filter is required, will return it, otherwise, will return <tt>null</tt>.
400400
*/
401+
@Nullable
401402
public Filter searchFilter(String... types) {
402403
if (types == null || types.length == 0) {
403404
if (hasNested) {

‎src/main/java/org/elasticsearch/index/shard/service/IndexShard.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,6 @@ public interface IndexShard extends IndexShardComponent {
9696

9797
Engine.GetResult get(Engine.Get get) throws ElasticSearchException;
9898

99-
long count(float minScore, BytesReference querySource, @Nullable String[] filteringAliases, String... types) throws ElasticSearchException;
100-
10199
void refresh(Engine.Refresh refresh) throws ElasticSearchException;
102100

103101
void flush(Engine.Flush flush) throws ElasticSearchException;

‎src/main/java/org/elasticsearch/index/shard/service/InternalIndexShard.java

Lines changed: 0 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,6 @@
3737
import org.elasticsearch.common.bytes.BytesReference;
3838
import org.elasticsearch.common.inject.Inject;
3939
import org.elasticsearch.common.io.FastByteArrayOutputStream;
40-
import org.elasticsearch.common.lucene.Lucene;
41-
import org.elasticsearch.common.lucene.search.Queries;
4240
import org.elasticsearch.common.metrics.MeanMetric;
4341
import org.elasticsearch.common.settings.Settings;
4442
import org.elasticsearch.common.unit.TimeValue;
@@ -54,7 +52,6 @@
5452
import org.elasticsearch.index.merge.MergeStats;
5553
import org.elasticsearch.index.merge.scheduler.MergeSchedulerProvider;
5654
import org.elasticsearch.index.query.IndexQueryParserService;
57-
import org.elasticsearch.index.query.QueryParseContext;
5855
import org.elasticsearch.index.refresh.RefreshStats;
5956
import org.elasticsearch.index.search.nested.IncludeAllChildrenQuery;
6057
import org.elasticsearch.index.search.nested.NonNestedDocsFilter;
@@ -402,39 +399,6 @@ public Engine.GetResult get(Engine.Get get) throws ElasticSearchException {
402399
return engine.get(get);
403400
}
404401

405-
@Override
406-
public long count(float minScore, BytesReference querySource, @Nullable String[] filteringAliases, String... types) throws ElasticSearchException {
407-
readAllowed();
408-
Query query;
409-
if (querySource == null || querySource.length() == 0) {
410-
query = Queries.MATCH_ALL_QUERY;
411-
} else {
412-
try {
413-
QueryParseContext.setTypes(types);
414-
query = queryParserService.parse(querySource).query();
415-
} finally {
416-
QueryParseContext.removeTypes();
417-
}
418-
}
419-
// wrap it in filter, cache it, and constant score it
420-
// Don't cache it, since it might be very different queries each time...
421-
// query = new ConstantScoreQuery(filterCache.cache(new QueryWrapperFilter(query)));
422-
query = filterQueryIfNeeded(query, types);
423-
Filter aliasFilter = indexAliasesService.aliasFilter(filteringAliases);
424-
Engine.Searcher searcher = engine.searcher();
425-
try {
426-
long count = Lucene.count(searcher.searcher(), query, aliasFilter, minScore);
427-
if (logger.isTraceEnabled()) {
428-
logger.trace("count of [{}] is [{}]", query, count);
429-
}
430-
return count;
431-
} catch (IOException e) {
432-
throw new ElasticSearchException("Failed to count query [" + query + "]", e);
433-
} finally {
434-
searcher.release();
435-
}
436-
}
437-
438402
@Override
439403
public void refresh(Engine.Refresh refresh) throws ElasticSearchException {
440404
verifyStarted();

‎src/main/java/org/elasticsearch/search/SearchService.java

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
package org.elasticsearch.search;
2121

2222
import com.google.common.collect.ImmutableMap;
23-
import org.apache.lucene.search.Filter;
2423
import org.apache.lucene.search.TopDocs;
2524
import org.elasticsearch.ElasticSearchException;
2625
import org.elasticsearch.action.search.SearchType;
@@ -490,9 +489,6 @@ SearchContext createContext(InternalSearchRequest request, @Nullable Engine.Sear
490489
context.size(10);
491490
}
492491

493-
Filter aliasFilter = indexService.aliasesService().aliasFilter(request.filteringAliases());
494-
context.aliasFilter(aliasFilter);
495-
496492
// pre process
497493
dfsPhase.preProcess(context);
498494
queryPhase.preProcess(context);

‎src/main/java/org/elasticsearch/search/internal/InternalSearchRequest.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,9 @@ public String[] filteringAliases() {
140140
return filteringAliases;
141141
}
142142

143-
public void filteringAliases(String[] filteringAliases) {
143+
public InternalSearchRequest filteringAliases(String[] filteringAliases) {
144144
this.filteringAliases = filteringAliases;
145+
return this;
145146
}
146147

147148
public String[] types() {

‎src/main/java/org/elasticsearch/search/internal/SearchContext.java

Lines changed: 29 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,14 @@
2121

2222
import com.google.common.collect.ImmutableList;
2323
import com.google.common.collect.Lists;
24-
import org.apache.lucene.search.Filter;
25-
import org.apache.lucene.search.Query;
26-
import org.apache.lucene.search.Sort;
24+
import org.apache.lucene.search.*;
2725
import org.elasticsearch.ElasticSearchException;
2826
import org.elasticsearch.action.search.SearchType;
2927
import org.elasticsearch.common.Nullable;
3028
import org.elasticsearch.common.lease.Releasable;
29+
import org.elasticsearch.common.lucene.search.Queries;
30+
import org.elasticsearch.common.lucene.search.function.BoostScoreFunction;
31+
import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
3132
import org.elasticsearch.index.analysis.AnalysisService;
3233
import org.elasticsearch.index.cache.field.data.FieldDataCache;
3334
import org.elasticsearch.index.cache.filter.FilterCache;
@@ -182,6 +183,9 @@ public SearchContext(long id, InternalSearchRequest request, SearchShardTarget s
182183
this.indexService = indexService;
183184

184185
this.searcher = new ContextIndexSearcher(this, engineSearcher);
186+
187+
// initialize the filtering alias based on the provided filters
188+
aliasFilter = indexService.aliasesService().aliasFilter(request.filteringAliases());
185189
}
186190

187191
@Override
@@ -205,6 +209,28 @@ public boolean release() throws ElasticSearchException {
205209
return true;
206210
}
207211

212+
/**
213+
* Should be called before executing the main query and after all other parameters have been set.
214+
*/
215+
public void preProcess() {
216+
if (query() == null) {
217+
parsedQuery(ParsedQuery.MATCH_ALL_PARSED_QUERY);
218+
}
219+
if (queryBoost() != 1.0f) {
220+
parsedQuery(new ParsedQuery(new FunctionScoreQuery(query(), new BoostScoreFunction(queryBoost)), parsedQuery()));
221+
}
222+
Filter searchFilter = mapperService().searchFilter(types());
223+
if (searchFilter != null) {
224+
if (Queries.isMatchAllQuery(query())) {
225+
Query q = new DeletionAwareConstantScoreQuery(filterCache().cache(searchFilter));
226+
q.setBoost(query().getBoost());
227+
parsedQuery(new ParsedQuery(q, parsedQuery()));
228+
} else {
229+
parsedQuery(new ParsedQuery(new FilteredQuery(query(), filterCache().cache(searchFilter)), parsedQuery()));
230+
}
231+
}
232+
}
233+
208234
public long id() {
209235
return this.id;
210236
}
@@ -383,11 +409,6 @@ public Filter parsedFilter() {
383409
return this.filter;
384410
}
385411

386-
public SearchContext aliasFilter(Filter aliasFilter) {
387-
this.aliasFilter = aliasFilter;
388-
return this;
389-
}
390-
391412
public Filter aliasFilter() {
392413
return aliasFilter;
393414
}

‎src/main/java/org/elasticsearch/search/query/QueryPhase.java

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,13 @@
2020
package org.elasticsearch.search.query;
2121

2222
import com.google.common.collect.ImmutableMap;
23-
import org.apache.lucene.search.*;
23+
import org.apache.lucene.search.Collector;
24+
import org.apache.lucene.search.Query;
25+
import org.apache.lucene.search.TopDocs;
26+
import org.apache.lucene.search.TotalHitCountCollector;
2427
import org.elasticsearch.action.search.SearchType;
2528
import org.elasticsearch.common.inject.Inject;
2629
import org.elasticsearch.common.lucene.Lucene;
27-
import org.elasticsearch.common.lucene.search.Queries;
28-
import org.elasticsearch.common.lucene.search.function.BoostScoreFunction;
29-
import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
30-
import org.elasticsearch.index.query.ParsedQuery;
3130
import org.elasticsearch.search.SearchParseElement;
3231
import org.elasticsearch.search.SearchPhase;
3332
import org.elasticsearch.search.facet.FacetPhase;
@@ -75,22 +74,7 @@ public QueryPhase(FacetPhase facetPhase) {
7574

7675
@Override
7776
public void preProcess(SearchContext context) {
78-
if (context.query() == null) {
79-
context.parsedQuery(ParsedQuery.MATCH_ALL_PARSED_QUERY);
80-
}
81-
if (context.queryBoost() != 1.0f) {
82-
context.parsedQuery(new ParsedQuery(new FunctionScoreQuery(context.query(), new BoostScoreFunction(context.queryBoost())), context.parsedQuery()));
83-
}
84-
Filter searchFilter = context.mapperService().searchFilter(context.types());
85-
if (searchFilter != null) {
86-
if (Queries.isMatchAllQuery(context.query())) {
87-
Query q = new DeletionAwareConstantScoreQuery(context.filterCache().cache(searchFilter));
88-
q.setBoost(context.query().getBoost());
89-
context.parsedQuery(new ParsedQuery(q, context.parsedQuery()));
90-
} else {
91-
context.parsedQuery(new ParsedQuery(new FilteredQuery(context.query(), context.filterCache().cache(searchFilter)), context.parsedQuery()));
92-
}
93-
}
77+
context.preProcess();
9478
facetPhase.preProcess(context);
9579
}
9680

‎src/test/java/org/elasticsearch/test/unit/common/lucene/search/MatchAllDocsFilterTests.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public void testMatchAllDocsFilter() throws Exception {
5252
IndexSearcher searcher = new IndexSearcher(reader);
5353

5454
DeletionAwareConstantScoreQuery query = new DeletionAwareConstantScoreQuery(Queries.MATCH_ALL_FILTER);
55-
long count = Lucene.count(searcher, query, -1);
55+
long count = Lucene.count(searcher, query);
5656
assertThat(count, equalTo(2l));
5757

5858
reader.close();

‎src/test/java/org/elasticsearch/test/unit/common/lucene/search/MoreLikeThisQueryTests.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ public void testSimple() throws Exception {
5757
mltQuery.setLikeText("lucene");
5858
mltQuery.setMinTermFrequency(1);
5959
mltQuery.setMinDocFreq(1);
60-
long count = Lucene.count(searcher, mltQuery, -1);
60+
long count = Lucene.count(searcher, mltQuery);
6161
assertThat(count, equalTo(2l));
6262

6363
reader.close();

0 commit comments

Comments
 (0)