diff --git a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java index 18de4b81cbf8c..fc7b7d005af76 100644 --- a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java +++ b/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java @@ -309,51 +309,53 @@ private static int sumMaxDocValues(List l) { public T search(Query query, CollectorManager collectorManager) throws IOException { final C firstCollector = collectorManager.newCollector(); // Take advantage of the few extra rewrite rules of ConstantScoreQuery when score are not needed. - query = firstCollector.scoreMode().needsScores() ? rewrite(query) : rewrite(new ConstantScoreQuery(query)); + final ScoreMode scoreMode = firstCollector.scoreMode(); + LeafSlice[] leafSlices = getSlices(); + if (leafSlices.length == 0) { + assert leafContexts.isEmpty(); + doAggregationPostCollection(firstCollector); + return collectorManager.reduce(Collections.singletonList(firstCollector)); + } final Weight weight; try { - weight = createWeight(query, firstCollector.scoreMode(), 1); + weight = createWeight(scoreMode.needsScores() ? rewrite(query) : rewrite(new ConstantScoreQuery(query)), scoreMode, 1); } catch (@SuppressWarnings("unused") TimeExceededException e) { timeExceeded = true; doAggregationPostCollection(firstCollector); return collectorManager.reduce(Collections.singletonList(firstCollector)); } - return search(weight, collectorManager, firstCollector); + if (leafSlices.length == 1) { + search(Arrays.asList(leafSlices[0].leaves), weight, firstCollector); + return collectorManager.reduce(Collections.singletonList(firstCollector)); + } + return searchMultipleSlices(collectorManager, leafSlices, firstCollector, weight); } - /** - * Same implementation as the default one in Lucene, with an additional call to postCollection in cased there are no segments. - * The rest is a plain copy from Lucene. - */ - private T search(Weight weight, CollectorManager collectorManager, C firstCollector) throws IOException { - LeafSlice[] leafSlices = getSlices(); - if (leafSlices.length == 0) { - assert leafContexts.isEmpty(); - doAggregationPostCollection(firstCollector); - return collectorManager.reduce(Collections.singletonList(firstCollector)); - } else { - final List collectors = new ArrayList<>(leafSlices.length); - collectors.add(firstCollector); - final ScoreMode scoreMode = firstCollector.scoreMode(); - for (int i = 1; i < leafSlices.length; ++i) { - final C collector = collectorManager.newCollector(); - collectors.add(collector); + private T searchMultipleSlices( + CollectorManager collectorManager, + LeafSlice[] leafSlices, + C firstCollector, + Weight weight + ) throws IOException { + final List> listTasks = new ArrayList<>(leafSlices.length); + final ScoreMode scoreMode = firstCollector.scoreMode(); + for (int i = 0; i < leafSlices.length; ++i) { + final C collector; + if (i == 0) { + collector = firstCollector; + } else { + collector = collectorManager.newCollector(); if (scoreMode != collector.scoreMode()) { throw new IllegalStateException("CollectorManager does not always produce collectors with the same score mode"); } } - final List> listTasks = new ArrayList<>(leafSlices.length); - for (int i = 0; i < leafSlices.length; ++i) { - final LeafReaderContext[] leaves = leafSlices[i].leaves; - final C collector = collectors.get(i); - listTasks.add(() -> { - search(Arrays.asList(leaves), weight, collector); - return collector; - }); - } - List collectedCollectors = getTaskExecutor().invokeAll(listTasks); - return collectorManager.reduce(collectedCollectors); + final LeafReaderContext[] leaves = leafSlices[i].leaves; + listTasks.add(() -> { + search(Arrays.asList(leaves), weight, collector); + return collector; + }); } + return collectorManager.reduce(getTaskExecutor().invokeAll(listTasks)); } /** diff --git a/server/src/test/java/org/elasticsearch/search/query/ThrowingQueryBuilder.java b/server/src/test/java/org/elasticsearch/search/query/ThrowingQueryBuilder.java index 8c21e95f48483..f2e51b260f4fd 100644 --- a/server/src/test/java/org/elasticsearch/search/query/ThrowingQueryBuilder.java +++ b/server/src/test/java/org/elasticsearch/search/query/ThrowingQueryBuilder.java @@ -97,9 +97,7 @@ protected Query doToQuery(SearchExecutionContext context) { return new Query() { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - if (context.getShardId() == shardId || shardId < 0 || context.index().getName().equals(index)) { - throw failure; - } + maybeThrow(context); return delegate.createWeight(searcher, scoreMode, boost); } @@ -122,9 +120,21 @@ public int hashCode() { public void visit(QueryVisitor visitor) { visitor.visitLeaf(this); } + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + maybeThrow(context); + return delegate.rewrite(indexSearcher); + } }; } + private void maybeThrow(SearchExecutionContext context) { + if (context.getShardId() == shardId || shardId < 0 || context.index().getName().equals(index)) { + throw failure; + } + } + @Override protected boolean doEquals(ThrowingQueryBuilder other) { return false;