Don't run sort optimization on size=0 (#57044)

Sort optimization creates TopFieldCollector that errors
when size=0. This ensures that sort optimization is not
run when size=0.

Closes #56923

Author: mayya-sharipova

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) && SYS_PROP_REWRITE_SORT) {
+            if ((searchContext.sort() != null) && (searchContext.size() > 0) && SYS_PROP_REWRITE_SORT) {
                 Query rewrittenQuery = tryRewriteLongSort(searchContext, searcher.getIndexReader(), query, hasFilterCollector);
                 if (rewrittenQuery != null) {
                     query = rewrittenQuery;

server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java

@@ -703,6 +703,23 @@ public void testNumericLongOrDateSortOptimization() throws Exception {
         searchContext.sort(sortAndFormats);
         QueryPhase.executeInternal(searchContext);
         assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, true);
+
+        // 5. Test that sort optimization is NOT run with size 0
+        {
+            sortAndFormats = new SortAndFormats(longSort, new DocValueFormat[]{DocValueFormat.RAW});
+            searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader)));
+            when(searchContext.mapperService()).thenReturn(mapperService);
+            searchContext.sort(sortAndFormats);
+            searchContext.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
+            searchContext.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()));
+            searchContext.setSize(0);
+
+            QueryPhase.executeInternal(searchContext);
+            TotalHits totalHits = searchContext.queryResult().topDocs().topDocs.totalHits;
+            assertEquals(TotalHits.Relation.EQUAL_TO, totalHits.relation);
+            assertEquals(numDocs, totalHits.value);
+        }
+
         reader.close();
         dir.close();
     }