Disable shard/segment level search_after short cutting if track_total…

…_hits != false (opensearch-project#9683) Signed-off-by: gashutos <gashutos@amazon.com> Signed-off-by: Chaitanya Gohel <104654647+gashutos@users.noreply.github.com> Signed-off-by: Shivansh Arora <hishiv@amazon.com>
Gaurav614 · Apr 25, 2024 · 335fbd2 · 335fbd2
1 parent 87906a1
commit 335fbd2
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 12 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -189,6 +189,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix condition to remove index create block ([#9437](https://github.com/opensearch-project/OpenSearch/pull/9437))
 - Add support to clear archived index setting ([#9019](https://github.com/opensearch-project/OpenSearch/pull/9019))
 - [Segment Replication] Fixed bug where replica shard temporarily serves stale data during an engine reset ([#9495](https://github.com/opensearch-project/OpenSearch/pull/9495))
+- Disable shard/segment level search_after short cutting if track_total_hits != false ([#9683](https://github.com/opensearch-project/OpenSearch/pull/9683))
 - [Segment Replication] Fixed bug where bytes behind metric is not accurate ([#9686](https://github.com/opensearch-project/OpenSearch/pull/9686))
 
 ### Security

diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -146,6 +146,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
@@ -1558,17 +1559,29 @@ private CanMatchResponse canMatch(ShardSearchRequest request, boolean checkRefre
                     canMatch = aliasFilterCanMatch;
                 }
                 final FieldDoc searchAfterFieldDoc = getSearchAfterFieldDoc(request, context);
-                canMatch = canMatch && canMatchSearchAfter(searchAfterFieldDoc, minMax, sortBuilder);
+                final Integer trackTotalHitsUpto = request.source() == null ? null : request.source().trackTotalHitsUpTo();
+                canMatch = canMatch && canMatchSearchAfter(searchAfterFieldDoc, minMax, sortBuilder, trackTotalHitsUpto);
 
                 return new CanMatchResponse(canMatch || hasRefreshPending, minMax);
             }
         }
     }
 
-    public static boolean canMatchSearchAfter(FieldDoc searchAfter, MinAndMax<?> minMax, FieldSortBuilder primarySortField) {
+    public static boolean canMatchSearchAfter(
+        FieldDoc searchAfter,
+        MinAndMax<?> minMax,
+        FieldSortBuilder primarySortField,
+        Integer trackTotalHitsUpto
+    ) {
         // Check for sort.missing == null, since in case of missing values sort queries, if segment/shard's min/max
         // is out of search_after range, it still should be printed and hence we should not skip segment/shard.
-        if (searchAfter != null && minMax != null && primarySortField != null && primarySortField.missing() == null) {
+        // Skipping search on shard/segment entirely can cause mismatch on total_tracking_hits, hence skip only if
+        // track_total_hits is false.
+        if (searchAfter != null
+            && minMax != null
+            && primarySortField != null
+            && primarySortField.missing() == null
+            && Objects.equals(trackTotalHitsUpto, SearchContext.TRACK_TOTAL_HITS_DISABLED)) {
             final Object searchAfterPrimary = searchAfter.fields[0];
             if (primarySortField.order() == SortOrder.DESC) {
                 if (minMax.compareMin(searchAfterPrimary) > 0) {

diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java
@@ -511,7 +511,12 @@ private boolean canMatchSearchAfter(LeafReaderContext ctx) throws IOException {
                     ctx,
                     primarySortField
                 );
-                return SearchService.canMatchSearchAfter(searchContext.searchAfter(), minMax, primarySortField);
+                return SearchService.canMatchSearchAfter(
+                    searchContext.searchAfter(),
+                    minMax,
+                    primarySortField,
+                    searchContext.trackTotalHitsUpTo()
+                );
             }
         }
         return true;

diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -1756,7 +1756,7 @@ public void testCanMatchSearchAfterAscGreaterThanMax() throws IOException {
         MinAndMax<?> minMax = new MinAndMax<Long>(0L, 9L);
         FieldSortBuilder primarySort = new FieldSortBuilder("test");
         primarySort.order(SortOrder.ASC);
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), false);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), false);
     }
 
     /**
@@ -1769,7 +1769,7 @@ public void testCanMatchSearchAfterAscLessThanMax() throws IOException {
         MinAndMax<?> minMax = new MinAndMax<Long>(0L, 9L);
         FieldSortBuilder primarySort = new FieldSortBuilder("test");
         primarySort.order(SortOrder.ASC);
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), true);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), true);
     }
 
     /**
@@ -1782,7 +1782,7 @@ public void testCanMatchSearchAfterAscEqualMax() throws IOException {
         MinAndMax<?> minMax = new MinAndMax<Long>(0L, 9L);
         FieldSortBuilder primarySort = new FieldSortBuilder("test");
         primarySort.order(SortOrder.ASC);
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), true);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), true);
     }
 
     /**
@@ -1795,7 +1795,7 @@ public void testCanMatchSearchAfterDescGreaterThanMin() throws IOException {
         MinAndMax<?> minMax = new MinAndMax<Long>(0L, 9L);
         FieldSortBuilder primarySort = new FieldSortBuilder("test");
         primarySort.order(SortOrder.DESC);
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), true);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), true);
     }
 
     /**
@@ -1808,7 +1808,7 @@ public void testCanMatchSearchAfterDescLessThanMin() throws IOException {
         MinAndMax<?> minMax = new MinAndMax<Long>(0L, 9L);
         FieldSortBuilder primarySort = new FieldSortBuilder("test");
         primarySort.order(SortOrder.DESC);
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), false);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), false);
     }
 
     /**
@@ -1821,7 +1821,7 @@ public void testCanMatchSearchAfterDescEqualMin() throws IOException {
         MinAndMax<?> minMax = new MinAndMax<Long>(0L, 9L);
         FieldSortBuilder primarySort = new FieldSortBuilder("test");
         primarySort.order(SortOrder.DESC);
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), true);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), true);
     }
 
     /**
@@ -1835,9 +1835,24 @@ public void testCanMatchSearchAfterWithMissing() throws IOException {
         FieldSortBuilder primarySort = new FieldSortBuilder("test");
         primarySort.order(SortOrder.DESC);
         // Should be false without missing values
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), false);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), false);
         primarySort.missing("_last");
         // Should be true with missing values
-        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort), true);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, SearchContext.TRACK_TOTAL_HITS_DISABLED), true);
+    }
+
+    /**
+     * Test for DESC order search_after query with track_total_hits=true.
+     * Min = 0L, Max = 9L, search_after = -1L
+     * With above min/max and search_after, it should not match, but since
+     * track_total_hits = true,
+     * Expected result is canMatch = true
+     */
+    public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IOException {
+        FieldDoc searchAfter = new FieldDoc(0, 0, new Long[] { -1L });
+        MinAndMax<?> minMax = new MinAndMax<Long>(0L, 9L);
+        FieldSortBuilder primarySort = new FieldSortBuilder("test");
+        primarySort.order(SortOrder.DESC);
+        assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, 1000), true);
     }
 }