optimize single shard search (or search result) by not doing another round of sorting

modules/elasticsearch/src/main/java/org/elasticsearch/search/controller/SearchPhaseController.java

@@ -27,7 +27,9 @@ import org.elasticsearch.common.collect.Iterables;
 import org.elasticsearch.common.collect.Lists;
 import org.elasticsearch.common.collect.Maps;
 import org.elasticsearch.common.collect.Ordering;
+import org.elasticsearch.common.component.AbstractComponent;
 import org.elasticsearch.common.inject.Inject;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.trove.ExtTIntArrayList;
 import org.elasticsearch.common.trove.impl.Constants;
 import org.elasticsearch.common.trove.map.hash.TObjectIntHashMap;
@@ -53,7 +55,7 @@ import java.util.Map;
 /**
  * @author kimchy (shay.banon)
  */
-public class SearchPhaseController {
+public class SearchPhaseController extends AbstractComponent {
 
     public static Ordering<QuerySearchResultProvider> QUERY_RESULT_ORDERING = new Ordering<QuerySearchResultProvider>() {
         @Override public int compare(@Nullable QuerySearchResultProvider o1, @Nullable QuerySearchResultProvider o2) {
@@ -69,8 +71,16 @@ public class SearchPhaseController {
 
     private final FacetProcessors facetProcessors;
 
-    @Inject public SearchPhaseController(FacetProcessors facetProcessors) {
+    private final boolean optimizeSingleShard;
+
+    @Inject public SearchPhaseController(Settings settings, FacetProcessors facetProcessors) {
+        super(settings);
         this.facetProcessors = facetProcessors;
+        this.optimizeSingleShard = componentSettings.getAsBoolean("optimize_single_shard", true);
+    }
+
+    public boolean optimizeSingleShard() {
+        return optimizeSingleShard;
     }
 
     public AggregatedDfs aggregateDfs(Iterable<DfsSearchResult> results) {
@@ -90,6 +100,52 @@ public class SearchPhaseController {
             return EMPTY;
         }
 
+        if (optimizeSingleShard) {
+            boolean canOptimize = false;
+            QuerySearchResult result = null;
+            if (results1.size() == 1) {
+                canOptimize = true;
+                result = results1.iterator().next().queryResult();
+            } else {
+                // lets see if we only got hits from a single shard, if so, we can optimize...
+                for (QuerySearchResultProvider queryResult : results1) {
+                    if (queryResult.queryResult().topDocs().scoreDocs.length > 0) {
+                        if (result != null) { // we already have one, can't really optimize
+                            canOptimize = false;
+                            break;
+                        }
+                        canOptimize = true;
+                        result = queryResult.queryResult();
+                    }
+                }
+            }
+            if (canOptimize) {
+                ScoreDoc[] scoreDocs = result.topDocs().scoreDocs;
+                if (scoreDocs.length < result.from()) {
+                    return EMPTY;
+                }
+                int resultDocsSize = result.size();
+                if ((scoreDocs.length - result.from()) < resultDocsSize) {
+                    resultDocsSize = scoreDocs.length - result.from();
+                }
+                if (result.topDocs() instanceof TopFieldDocs) {
+                    ShardDoc[] docs = new ShardDoc[resultDocsSize];
+                    for (int i = 0; i < resultDocsSize; i++) {
+                        ScoreDoc scoreDoc = scoreDocs[result.from() + i];
+                        docs[i] = new ShardFieldDoc(result.shardTarget(), scoreDoc.doc, scoreDoc.score, ((FieldDoc) scoreDoc).fields);
+                    }
+                    return docs;
+                } else {
+                    ShardDoc[] docs = new ShardDoc[resultDocsSize];
+                    for (int i = 0; i < resultDocsSize; i++) {
+                        ScoreDoc scoreDoc = scoreDocs[result.from() + i];
+                        docs[i] = new ShardScoreDoc(result.shardTarget(), scoreDoc.doc, scoreDoc.score);
+                    }
+                    return docs;
+                }
+            }
+        }
+
         List<? extends QuerySearchResultProvider> results = QUERY_RESULT_ORDERING.sortedCopy(results1);
 
         QuerySearchResultProvider queryResultProvider = results.get(0);

modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/embedded/SingleShardEmbeddedSearchTests.java

@@ -66,7 +66,12 @@ public class SingleShardEmbeddedSearchTests extends AbstractNodesTests {
 
     private Client client;
 
+    protected boolean optimizeSingleShard() {
+        return true;
+    }
+
     @BeforeClass public void createNodeAndInitWithData() throws Exception {
+        putDefaultSettings(settingsBuilder().put("search.controller.optimize_single_shard", optimizeSingleShard()));
         startNode("server1");
         client = client("server1");
 
@@ -89,6 +94,10 @@ public class SingleShardEmbeddedSearchTests extends AbstractNodesTests {
         closeAllNodes();
     }
 
+    @Test public void verifyOptimizeSingleShardSetting() {
+        assertThat(searchPhaseController.optimizeSingleShard(), equalTo(optimizeSingleShard()));
+    }
+
     @Test public void testDirectDfs() throws Exception {
         DfsSearchResult dfsResult = searchService.executeDfsPhase(searchRequest(searchSource().query(termQuery("name", "test1"))));
 

modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/embedded/SingleShardNoOptimizationEmbeddedSearchTests.java

@@ -0,0 +1,27 @@
+/*
+ * Licensed to Elastic Search and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Elastic Search licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.test.integration.search.embedded;
+
+public class SingleShardNoOptimizationEmbeddedSearchTests extends SingleShardEmbeddedSearchTests {
+
+    @Override protected boolean optimizeSingleShard() {
+        return false;
+    }
+}