Run sort optimization when from+size>0 (#57250)

elastic · May 29, 2020 · 4ef75bb · 4ef75bb
1 parent 387b181
commit 4ef75bb
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java
@@ -244,7 +244,7 @@ static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExe
 
             CheckedConsumer<List<LeafReaderContext>, IOException> leafSorter = l -> {};
             // try to rewrite numeric or date sort to the optimized distanceFeatureQuery
-            if ((searchContext.sort() != null) && (searchContext.size() > 0) && SYS_PROP_REWRITE_SORT) {
+            if ((searchContext.sort() != null) && SYS_PROP_REWRITE_SORT) {
                 Query rewrittenQuery = tryRewriteLongSort(searchContext, searcher.getIndexReader(), query, hasFilterCollector);
                 if (rewrittenQuery != null) {
                     query = rewrittenQuery;
@@ -415,6 +415,7 @@ private static boolean searchWithCollectorManager(SearchContext searchContext, C
 
     private static Query tryRewriteLongSort(SearchContext searchContext, IndexReader reader,
                                             Query query, boolean hasFilterCollector) throws IOException {
+        if ((searchContext.from() + searchContext.size()) <= 0) return null;
         if (searchContext.searchAfter() != null) return null; //TODO: handle sort optimization with search after
         if (searchContext.scrollContext() != null) return null;
         if (searchContext.collapse() != null) return null;

diff --git a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java
@@ -712,7 +712,17 @@ public void testNumericLongOrDateSortOptimization() throws Exception {
         QueryPhase.executeInternal(searchContext);
         assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, true);
 
-        // 5. Test that sort optimization is NOT run with size 0
+        // 5. Test that sort optimization is run when from > 0 and size = 0
+        {
+            sortAndFormats = new SortAndFormats(longSort, new DocValueFormat[]{DocValueFormat.RAW});
+            searchContext.sort(sortAndFormats);
+            searchContext.from(5);
+            searchContext.setSize(0);
+            QueryPhase.executeInternal(searchContext);
+            assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, false);
+        }
+
+        // 6. Test that sort optimization is NOT run with from = 0 and size= 0
         {
             sortAndFormats = new SortAndFormats(longSort, new DocValueFormat[]{DocValueFormat.RAW});
             searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader)));

diff --git a/test/framework/src/main/java/org/elasticsearch/test/TestSearchContext.java b/test/framework/src/main/java/org/elasticsearch/test/TestSearchContext.java
@@ -86,6 +86,7 @@ public class TestSearchContext extends SearchContext {
     int trackTotalHitsUpTo = SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO;
 
     ContextIndexSearcher searcher;
+    int from;
     int size;
     private int terminateAfter = DEFAULT_TERMINATE_AFTER;
     private SearchContextAggregations aggregations;
@@ -432,12 +433,13 @@ public Query query() {
 
     @Override
     public int from() {
-        return 0;
+        return from;
     }
 
     @Override
     public SearchContext from(int from) {
-        return null;
+        this.from = from;
+        return this;
     }
 
     @Override