Backport of 18189 PR

Signed-off-by: Prudhvi Godithi <[email protected]>

Author: prudhvigodithi

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 3.0.x]
 ### Added
+Improve sort-query performance by retaining the default `totalHitsThreshold` for approximated `match_all` queries ([#18189](https://github.com/opensearch-project/OpenSearch/pull/18189))
 
 ### Changed
 

server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java

@@ -1145,7 +1145,7 @@ public void testSortMissingNumbersMinMax() throws Exception {
             .get();
         assertNoFailures(searchResponse);
 
-        assertThat(searchResponse.getHits().getTotalHits().value(), equalTo(2L));
+        assertThat(searchResponse.getHits().getTotalHits().value(), equalTo(3L));
         assertThat(searchResponse.getHits().getAt(0).getId(), equalTo("1"));
         // The order here could be unstable (depends on document order) since missing == field value
         assertThat(searchResponse.getHits().getAt(1).getId(), is(oneOf("3", "2")));

server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java

@@ -444,11 +444,12 @@ public boolean canApproximate(SearchContext context) {
         if (context.trackTotalHitsUpTo() == SearchContext.TRACK_TOTAL_HITS_ACCURATE) {
             return false;
         }
+
         // size 0 could be set for caching
         if (context.from() + context.size() == 0) {
             this.setSize(SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO);
         } else {
-            this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo() + 1));
+            this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo()));
         }
         if (context.request() != null && context.request().source() != null) {
             FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source());

server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Weight;
 import org.opensearch.search.internal.SearchContext;
 
@@ -47,7 +48,11 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
             // Default to the original query. This suggests that we were not called from ContextIndexSearcher.
             return originalQuery.rewrite(indexSearcher);
         }
-        return resolvedQuery.rewrite(indexSearcher);
+        Query rewritten = resolvedQuery.rewrite(indexSearcher);
+        if (rewritten != resolvedQuery) {
+            resolvedQuery = rewritten;
+        }
+        return this;
     }
 
     public void setContext(SearchContext context) {
@@ -78,6 +83,15 @@ public boolean equals(Object o) {
         return true;
     }
 
+    @Override
+    public Weight createWeight(IndexSearcher indexSearcher, ScoreMode scoreMode, float boost) throws IOException {
+        if (resolvedQuery == null) {
+            // Default to the original query.
+            return originalQuery.createWeight(indexSearcher, scoreMode, boost);
+        }
+        return resolvedQuery.createWeight(indexSearcher, scoreMode, boost);
+    }
+
     @Override
     public int hashCode() {
         int h = classHash();

server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java

@@ -74,6 +74,7 @@
 import org.opensearch.common.util.CachedSupplier;
 import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery;
 import org.opensearch.search.DocValueFormat;
+import org.opensearch.search.approximate.ApproximateScoreQuery;
 import org.opensearch.search.collapse.CollapseContext;
 import org.opensearch.search.internal.ScrollContext;
 import org.opensearch.search.internal.SearchContext;
@@ -724,6 +725,8 @@ static int shortcutTotalHitCount(IndexReader reader, Query query) throws IOExcep
                 query = ((ConstantScoreQuery) query).getQuery();
             } else if (query instanceof BoostQuery) {
                 query = ((BoostQuery) query).getQuery();
+            } else if (query instanceof ApproximateScoreQuery) {
+                query = ((ApproximateScoreQuery) query).getOriginalQuery();
             } else {
                 break;
             }

server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java

@@ -287,13 +287,20 @@ public void testRangeQuery() throws IOException {
         String date2 = "2016-04-28T11:33:52";
         long instant1 = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date1)).toInstant().toEpochMilli();
         long instant2 = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date2)).toInstant().toEpochMilli() + 999;
-        Query expected = new ApproximatePointRangeQuery(
+        ApproximatePointRangeQuery approximatePointRangeQuery = new ApproximatePointRangeQuery(
             "field",
             pack(new long[] { instant1 }).bytes,
             pack(new long[] { instant2 }).bytes,
             new long[] { instant1 }.length,
             ApproximatePointRangeQuery.LONG_FORMAT
         );
+        Query expected = new ApproximateScoreQuery(
+            new IndexOrDocValuesQuery(
+                LongPoint.newRangeQuery("field", instant1, instant2),
+                SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
+            ),
+            approximatePointRangeQuery
+        );
         Query rangeQuery = ft.rangeQuery(date1, date2, true, true, null, null, null, context);
         assertTrue(rangeQuery instanceof ApproximateScoreQuery);
         ((ApproximateScoreQuery) rangeQuery).setContext(new TestSearchContext(context));