Speedup SearchContext.hasOnlySuggest check and related logic (elastic…

…#112747) Random find from looking at what we do on the transport thread for search: This check was quite heavy by building the list of sub-searches needlessly. Also, we were needlessly allocating capturing lambdas in the stats logic that used this call.
pxsalehi · Sep 12, 2024 · c1f09d3 · c1f09d3
1 parent f7a0196
commit c1f09d3
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 26 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/search/stats/ShardSearchStats.java b/server/src/main/java/org/elasticsearch/index/search/stats/ShardSearchStats.java
@@ -56,36 +56,28 @@ public SearchStats stats(String... groups) {
 
     @Override
     public void onPreQueryPhase(SearchContext searchContext) {
-        computeStats(searchContext, statsHolder -> {
-            if (searchContext.hasOnlySuggest()) {
-                statsHolder.suggestCurrent.inc();
-            } else {
-                statsHolder.queryCurrent.inc();
-            }
-        });
+        computeStats(
+            searchContext,
+            searchContext.hasOnlySuggest() ? statsHolder -> statsHolder.suggestCurrent.inc() : statsHolder -> statsHolder.queryCurrent.inc()
+        );
     }
 
     @Override
     public void onFailedQueryPhase(SearchContext searchContext) {
-        computeStats(searchContext, statsHolder -> {
-            if (searchContext.hasOnlySuggest()) {
-                statsHolder.suggestCurrent.dec();
-            } else {
-                statsHolder.queryCurrent.dec();
-            }
-        });
+        computeStats(
+            searchContext,
+            searchContext.hasOnlySuggest() ? statsHolder -> statsHolder.suggestCurrent.dec() : statsHolder -> statsHolder.queryCurrent.dec()
+        );
     }
 
     @Override
     public void onQueryPhase(SearchContext searchContext, long tookInNanos) {
-        computeStats(searchContext, statsHolder -> {
-            if (searchContext.hasOnlySuggest()) {
-                statsHolder.suggestMetric.inc(tookInNanos);
-                statsHolder.suggestCurrent.dec();
-            } else {
-                statsHolder.queryMetric.inc(tookInNanos);
-                statsHolder.queryCurrent.dec();
-            }
+        computeStats(searchContext, searchContext.hasOnlySuggest() ? statsHolder -> {
+            statsHolder.suggestMetric.inc(tookInNanos);
+            statsHolder.suggestCurrent.dec();
+        } : statsHolder -> {
+            statsHolder.queryMetric.inc(tookInNanos);
+            statsHolder.queryCurrent.dec();
         });
     }
 
@@ -109,8 +101,9 @@ public void onFetchPhase(SearchContext searchContext, long tookInNanos) {
 
     private void computeStats(SearchContext searchContext, Consumer<StatsHolder> consumer) {
         consumer.accept(totalStats);
-        if (searchContext.groupStats() != null) {
-            for (String group : searchContext.groupStats()) {
+        var groupStats = searchContext.groupStats();
+        if (groupStats != null) {
+            for (String group : groupStats) {
                 consumer.accept(groupStats(group));
             }
         }

diff --git a/server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java b/server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java
@@ -1104,7 +1104,7 @@ public List<SearchExtBuilder> ext() {
      * @return true if the source only has suggest
      */
     public boolean isSuggestOnly() {
-        return suggestBuilder != null && query() == null && knnSearch.isEmpty() && aggregations == null;
+        return suggestBuilder != null && knnSearch.isEmpty() && aggregations == null && subSearchSourceBuilders.isEmpty();
     }
 
     /**

diff --git a/server/src/main/java/org/elasticsearch/search/internal/SearchContext.java b/server/src/main/java/org/elasticsearch/search/internal/SearchContext.java
@@ -364,7 +364,8 @@ public void addReleasable(Releasable releasable) {   // TODO most Releasables ar
      * @return true if the request contains only suggest
      */
     public final boolean hasOnlySuggest() {
-        return request().source() != null && request().source().isSuggestOnly();
+        var source = request().source();
+        return source != null && source.isSuggestOnly();
     }
 
     /**