Skip approcimation when track total hits sets to true

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

Author: prudhvigodithi

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

@@ -36,6 +36,10 @@ protected boolean canApproximate(SearchContext context) {
         if (context.aggregations() != null) {
             return false;
         }
+        // Exclude approximation when "track_total_hits": true
+        if (context.trackTotalHitsUpTo() == SearchContext.TRACK_TOTAL_HITS_ACCURATE) {
+            return false;
+        }
 
         if (context.request() != null && context.request().source() != null && context.innerHits().getInnerHits().isEmpty()) {
             FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source());

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

@@ -440,6 +440,11 @@ public boolean canApproximate(SearchContext context) {
         if (context.aggregations() != null) {
             return false;
         }
+        // Exclude approximation when "track_total_hits": true
+        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);

server/src/test/java/org/opensearch/search/approximate/ApproximateMatchAllQueryTests.java

@@ -20,6 +20,7 @@
 import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.aggregations.SearchContextAggregations;
 import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.internal.ShardSearchRequest;
 import org.opensearch.search.sort.FieldSortBuilder;
 import org.opensearch.search.sort.SortOrder;
@@ -105,4 +106,60 @@ public ShardSearchRequest request() {
         assertThrows(IllegalStateException.class, () -> approximateMatchAllQuery.rewrite(null));
     }
 
+    public void testCannotApproximateWithTrackTotalHits() {
+        ApproximateMatchAllQuery approximateMatchAllQuery = new ApproximateMatchAllQuery();
+
+        ShardSearchRequest[] shardSearchRequest = new ShardSearchRequest[1];
+
+        MapperService mockMapper = mock(MapperService.class);
+        String sortfield = "myfield";
+        MappedFieldType myFieldType = new NumberFieldMapper.NumberFieldType(sortfield, NumberFieldMapper.NumberType.LONG);
+        when(mockMapper.fieldType(sortfield)).thenReturn(myFieldType);
+
+        Settings settings = Settings.builder()
+            .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+            .build();
+        IndexMetadata indexMetadata = new IndexMetadata.Builder("index").settings(settings).build();
+        QueryShardContext queryShardContext = new QueryShardContext(
+            0,
+            new IndexSettings(indexMetadata, settings),
+            BigArrays.NON_RECYCLING_INSTANCE,
+            null,
+            null,
+            mockMapper,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null
+        );
+        TestSearchContext searchContext = new TestSearchContext(queryShardContext) {
+            @Override
+            public ShardSearchRequest request() {
+                return shardSearchRequest[0];
+            }
+        };
+
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        shardSearchRequest[0] = new ShardSearchRequest(null, System.currentTimeMillis(), null);
+        shardSearchRequest[0].source(source);
+        source.sort(sortfield, SortOrder.ASC);
+
+        assertTrue(approximateMatchAllQuery.canApproximate(searchContext));
+
+        searchContext.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_ACCURATE);
+        assertFalse("Should not approximate when track_total_hits is accurate", approximateMatchAllQuery.canApproximate(searchContext));
+
+        searchContext.trackTotalHitsUpTo(SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO);
+        assertTrue("Should approximate when track_total_hits is not accurate", approximateMatchAllQuery.canApproximate(searchContext));
+    }
+
 }

server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java

@@ -33,6 +33,7 @@
 import static java.util.Arrays.asList;
 import static org.apache.lucene.document.LongPoint.pack;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class ApproximatePointRangeQueryTests extends OpenSearchTestCase {
 
@@ -372,4 +373,23 @@ public boolean canApproximate(SearchContext context) {
         SearchContext searchContext = mock(SearchContext.class);
         assertTrue(queryCanApproximate.canApproximate(searchContext));
     }
+
+    public void testCannotApproximateWithTrackTotalHits() {
+        ApproximatePointRangeQuery query = new ApproximatePointRangeQuery(
+            "point",
+            pack(0).bytes,
+            pack(20).bytes,
+            1,
+            ApproximatePointRangeQuery.LONG_FORMAT
+        );
+        SearchContext mockContext = mock(SearchContext.class);
+        when(mockContext.trackTotalHitsUpTo()).thenReturn(SearchContext.TRACK_TOTAL_HITS_ACCURATE);
+        assertFalse(query.canApproximate(mockContext));
+        when(mockContext.trackTotalHitsUpTo()).thenReturn(SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO);
+        when(mockContext.aggregations()).thenReturn(null);
+        when(mockContext.from()).thenReturn(0);
+        when(mockContext.size()).thenReturn(10);
+        when(mockContext.request()).thenReturn(null);
+        assertTrue(query.canApproximate(mockContext));
+    }
 }