Return an empty suggestion when suggest phase times out (#122575)

We recently fixed timeout handling in the suggest phase. A test failure on SearchTimeoutIT surfaced an issue with the current approach. In case partial results are allowed, it may happen that some shards time out while executing the suggest phase and some don't.

SearchPhaseController assumes that if one shard has suggest results, all of the other shards will have suggest results too. We could address that assertion and check is the search timed out, instead this commit changes timeout handling in the suggest phase to return an empty suggestion instead of null. This seems appropriate in terms of providing some results and makes the assertion about non null suggestions in SearchPhaseController happy.

Relates to #122357

Closes #122548

Author: javanna

docs/changelog/122575.yaml

@@ -0,0 +1,6 @@
+pr: 122575
+summary: Return an empty suggestion when suggest phase times out
+area: Suggesters
+type: bug
+issues:
+ - 122548

muted-tests.yml

@@ -326,9 +326,6 @@ tests:
 - class: org.elasticsearch.compute.operator.exchange.ExchangeServiceTests
   method: testExchangeSourceContinueOnFailure
   issue: https://github.com/elastic/elasticsearch/issues/122408
-- class: org.elasticsearch.search.SearchTimeoutIT
-  method: testSuggestTimeoutWithPartialResults
-  issue: https://github.com/elastic/elasticsearch/issues/122548
 
 # Examples:
 #

server/src/internalClusterTest/java/org/elasticsearch/search/SearchTimeoutIT.java

@@ -484,8 +484,7 @@ protected TermSuggestion innerExecute(
             CharsRefBuilder spare
         ) {
             contextIndexSearcher.throwTimeExceededException();
-            assert false;
-            return new TermSuggestion(name, suggestion.getSize(), SortBy.SCORE);
+            throw new AssertionError("should have thrown TimeExceededException");
         }
 
         @Override

server/src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -126,15 +126,7 @@ static void executeRank(SearchContext searchContext) throws QueryPhaseExecutionE
 
     static void executeQuery(SearchContext searchContext) throws QueryPhaseExecutionException {
         if (searchContext.hasOnlySuggest()) {
-            try {
-                SuggestPhase.execute(searchContext);
-            } catch (ContextIndexSearcher.TimeExceededException timeExceededException) {
-                SearchTimeoutException.handleTimeout(
-                    searchContext.request().allowPartialSearchResults(),
-                    searchContext.shardTarget(),
-                    searchContext.queryResult()
-                );
-            }
+            SuggestPhase.execute(searchContext);
             searchContext.queryResult().topDocs(new TopDocsAndMaxScore(Lucene.EMPTY_TOP_DOCS, Float.NaN), new DocValueFormat[0]);
             return;
         }
@@ -150,18 +142,10 @@ static void executeQuery(SearchContext searchContext) throws QueryPhaseExecution
 
         addCollectorsAndSearch(searchContext);
 
-        try {
-            RescorePhase.execute(searchContext);
-            SuggestPhase.execute(searchContext);
-            if (searchContext.getProfilers() != null) {
-                searchContext.queryResult().profileResults(searchContext.getProfilers().buildQueryPhaseResults());
-            }
-        } catch (ContextIndexSearcher.TimeExceededException timeExceededException) {
-            SearchTimeoutException.handleTimeout(
-                searchContext.request().allowPartialSearchResults(),
-                searchContext.shardTarget(),
-                searchContext.queryResult()
-            );
+        RescorePhase.execute(searchContext);
+        SuggestPhase.execute(searchContext);
+        if (searchContext.getProfilers() != null) {
+            searchContext.queryResult().profileResults(searchContext.getProfilers().buildQueryPhaseResults());
         }
     }
 

server/src/main/java/org/elasticsearch/search/rescore/RescorePhase.java

@@ -18,7 +18,9 @@
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.lucene.grouping.TopFieldGroups;
+import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.query.SearchTimeoutException;
 import org.elasticsearch.search.sort.ShardDocSortField;
 import org.elasticsearch.search.sort.SortAndFormats;
 
@@ -84,6 +86,13 @@ public static void execute(SearchContext context) {
                 .topDocs(new TopDocsAndMaxScore(topDocs, topDocs.scoreDocs[0].score), context.queryResult().sortValueFormats());
         } catch (IOException e) {
             throw new ElasticsearchException("Rescore Phase Failed", e);
+        } catch (ContextIndexSearcher.TimeExceededException timeExceededException) {
+            SearchTimeoutException.handleTimeout(
+                context.request().allowPartialSearchResults(),
+                context.shardTarget(),
+                context.queryResult()
+            );
+            // if the rescore phase times out and partial results are allowed, the returned top docs from this shard won't be rescored
         }
     }
 

server/src/main/java/org/elasticsearch/search/suggest/SuggestPhase.java

@@ -10,7 +10,9 @@
 
 import org.apache.lucene.util.CharsRefBuilder;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.query.SearchTimeoutException;
 import org.elasticsearch.search.suggest.Suggest.Suggestion;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry.Option;
@@ -40,12 +42,17 @@ public static void execute(SearchContext context) {
             for (Map.Entry<String, SuggestionSearchContext.SuggestionContext> entry : suggest.suggestions().entrySet()) {
                 SuggestionSearchContext.SuggestionContext suggestion = entry.getValue();
                 Suggester<SuggestionContext> suggester = suggestion.getSuggester();
-                Suggestion<? extends Entry<? extends Option>> result = suggester.execute(
-                    entry.getKey(),
-                    suggestion,
-                    context.searcher(),
-                    spare
-                );
+                Suggestion<? extends Entry<? extends Option>> result;
+                try {
+                    result = suggester.execute(entry.getKey(), suggestion, context.searcher(), spare);
+                } catch (ContextIndexSearcher.TimeExceededException timeExceededException) {
+                    SearchTimeoutException.handleTimeout(
+                        context.request().allowPartialSearchResults(),
+                        context.shardTarget(),
+                        context.queryResult()
+                    );
+                    result = suggester.emptySuggestion(entry.getKey(), suggestion, spare);
+                }
                 if (result != null) {
                     assert entry.getKey().equals(result.name);
                     suggestions.add(result);

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

@@ -297,7 +297,8 @@ public void testSuggestOnlyWithTimeout() throws Exception {
             assertTrue(context.hasOnlySuggest());
             QueryPhase.execute(context);
             assertTrue(context.queryResult().searchTimedOut());
-            assertNull(context.queryResult().suggest());
+            assertEquals(1, context.queryResult().suggest().size());
+            assertEquals(0, context.queryResult().suggest().getSuggestion("suggestion").getEntries().size());
             assertNotNull(context.queryResult().topDocs());
             assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value());
         }
@@ -311,7 +312,8 @@ public void testSuggestAndQueryWithSuggestTimeout() throws Exception {
             QueryPhase.execute(context);
             assertThat(context.queryResult().topDocs().topDocs.totalHits.value(), Matchers.greaterThan(0L));
             assertTrue(context.queryResult().searchTimedOut());
-            assertNull(context.queryResult().suggest());
+            assertEquals(1, context.queryResult().suggest().size());
+            assertEquals(0, context.queryResult().suggest().getSuggestion("suggestion").getEntries().size());
         }
     }