From aebb78bf5c74ce82aaf9b48800cd78ea8411e77b Mon Sep 17 00:00:00 2001
From: Mayya Sharipova <mayya.sharipova@elastic.co>
Date: Fri, 29 May 2020 11:26:25 -0400
Subject: [PATCH] Run sort optimization when from+size>0 (#57250)

---
 .../org/elasticsearch/search/query/QueryPhase.java   |  3 ++-
 .../elasticsearch/search/query/QueryPhaseTests.java  | 12 +++++++++++-
 .../org/elasticsearch/test/TestSearchContext.java    |  6 ++++--
 3 files changed, 17 insertions(+), 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
index f36a5ccac6c..9e94f1ddf11 100644
--- 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 @@ public class QueryPhase implements SearchPhase {
 
             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 @@ public class QueryPhase implements SearchPhase {
 
     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
index d0c2d7def0c..6eaef7ec4f5 100644
--- a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java
+++ b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java
@@ -713,7 +713,17 @@ public class QueryPhaseTests extends IndexShardTestCase {
         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
index 12f06bd6a9d..f7f410e6389 100644
--- 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 class TestSearchContext extends SearchContext {
 
     @Override
     public int from() {
-        return 0;
+        return from;
     }
 
     @Override
     public SearchContext from(int from) {
-        return null;
+        this.from = from;
+        return this;
     }
 
     @Override