save more

Author: original-brownbear

server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java

@@ -40,55 +40,45 @@
  */
 class DfsQueryPhase extends SearchPhase {
     private final SearchPhaseResults<SearchPhaseResult> queryResult;
-    private final List<DfsSearchResult> searchResults;
-    private final AggregatedDfs dfs;
-    private final List<DfsKnnResults> knnResults;
     private final Client client;
     private final AbstractSearchAsyncAction<?> context;
-    private final SearchTransportService searchTransportService;
     private final SearchProgressListener progressListener;
 
-    DfsQueryPhase(
-        List<DfsSearchResult> searchResults,
-        List<DfsKnnResults> knnResults,
-        SearchPhaseResults<SearchPhaseResult> queryResult,
-        Client client,
-        AbstractSearchAsyncAction<?> context
-    ) {
+    DfsQueryPhase(SearchPhaseResults<SearchPhaseResult> queryResult, Client client, AbstractSearchAsyncAction<?> context) {
         super("dfs_query");
         this.progressListener = context.getTask().getProgressListener();
         this.queryResult = queryResult;
-        this.searchResults = searchResults;
-        this.dfs = SearchPhaseController.aggregateDfs(searchResults);
-        this.knnResults = knnResults;
         this.client = client;
         this.context = context;
-        this.searchTransportService = context.getSearchTransport();
     }
 
     // protected for testing
-    protected SearchPhase nextPhase() {
+    protected SearchPhase nextPhase(AggregatedDfs dfs) {
         return SearchQueryThenFetchAsyncAction.nextPhase(client, context, queryResult, dfs);
     }
 
+    @SuppressWarnings("unchecked")
     @Override
     public void run() {
+        List<DfsSearchResult> searchResults = (List<DfsSearchResult>) context.results.getAtomicArray().asList();
+        AggregatedDfs dfs = SearchPhaseController.aggregateDfs(searchResults);
         // TODO we can potentially also consume the actual per shard results from the initial phase here in the aggregateDfs
         // to free up memory early
         final CountedCollector<SearchPhaseResult> counter = new CountedCollector<>(
             queryResult,
             searchResults.size(),
-            () -> context.executeNextPhase(this, this::nextPhase),
+            () -> context.executeNextPhase(this, () -> nextPhase(dfs)),
             context
         );
 
+        List<DfsKnnResults> knnResults = SearchPhaseController.mergeKnnResults(context.getRequest(), searchResults);
         for (final DfsSearchResult dfsResult : searchResults) {
             final SearchShardTarget shardTarget = dfsResult.getSearchShardTarget();
             final int shardIndex = dfsResult.getShardIndex();
             QuerySearchRequest querySearchRequest = new QuerySearchRequest(
                 context.getOriginalIndices(shardIndex),
                 dfsResult.getContextId(),
-                rewriteShardSearchRequest(dfsResult.getShardSearchRequest()),
+                rewriteShardSearchRequest(knnResults, dfsResult.getShardSearchRequest()),
                 dfs
             );
             final Transport.Connection connection;
@@ -98,11 +88,8 @@ public void run() {
                 shardFailure(e, querySearchRequest, shardIndex, shardTarget, counter);
                 continue;
             }
-            searchTransportService.sendExecuteQuery(
-                connection,
-                querySearchRequest,
-                context.getTask(),
-                new SearchActionListener<>(shardTarget, shardIndex) {
+            context.getSearchTransport()
+                .sendExecuteQuery(connection, querySearchRequest, context.getTask(), new SearchActionListener<>(shardTarget, shardIndex) {
 
                     @Override
                     protected void innerOnResponse(QuerySearchResult response) {
@@ -127,8 +114,7 @@ public void onFailure(Exception exception) {
                             }
                         }
                     }
-                }
-            );
+                });
         }
     }
 
@@ -145,7 +131,7 @@ private void shardFailure(
     }
 
     // package private for testing
-    ShardSearchRequest rewriteShardSearchRequest(ShardSearchRequest request) {
+    ShardSearchRequest rewriteShardSearchRequest(List<DfsKnnResults> knnResults, ShardSearchRequest request) {
         SearchSourceBuilder source = request.source();
         if (source == null || source.knnSearch().isEmpty()) {
             return request;

server/src/main/java/org/elasticsearch/action/search/SearchDfsQueryThenFetchAsyncAction.java

@@ -17,12 +17,10 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchShardTarget;
-import org.elasticsearch.search.dfs.DfsKnnResults;
 import org.elasticsearch.search.dfs.DfsSearchResult;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.transport.Transport;
 
-import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
 import java.util.function.BiFunction;
@@ -91,9 +89,7 @@ protected void executePhaseOnShard(
 
     @Override
     protected SearchPhase getNextPhase() {
-        final List<DfsSearchResult> dfsSearchResults = results.getAtomicArray().asList();
-        final List<DfsKnnResults> mergedKnnResults = SearchPhaseController.mergeKnnResults(getRequest(), dfsSearchResults);
-        return new DfsQueryPhase(dfsSearchResults, mergedKnnResults, queryPhaseResultConsumer, client, this);
+        return new DfsQueryPhase(queryPhaseResultConsumer, client, this);
     }
 
     @Override

server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java

@@ -27,6 +27,7 @@
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.dfs.AggregatedDfs;
 import org.elasticsearch.search.dfs.DfsKnnResults;
 import org.elasticsearch.search.dfs.DfsSearchResult;
 import org.elasticsearch.search.internal.AliasFilter;
@@ -319,9 +320,13 @@ private static DfsQueryPhase makeDfsPhase(
         MockSearchPhaseContext mockSearchPhaseContext,
         AtomicReference<AtomicArray<SearchPhaseResult>> responseRef
     ) {
-        return new DfsQueryPhase(results.asList(), null, consumer, null, mockSearchPhaseContext) {
+        int shards = mockSearchPhaseContext.numShards;
+        for (int i = 0; i < shards; i++) {
+            mockSearchPhaseContext.results.getAtomicArray().set(i, results.get(i));
+        }
+        return new DfsQueryPhase(consumer, null, mockSearchPhaseContext) {
             @Override
-            protected SearchPhase nextPhase() {
+            protected SearchPhase nextPhase(AggregatedDfs dfs) {
                 return new SearchPhase("test") {
                     @Override
                     public void run() {
@@ -342,7 +347,7 @@ public void testRewriteShardSearchRequestWithRank() {
         );
         MockSearchPhaseContext mspc = new MockSearchPhaseContext(2);
         mspc.searchTransport = new SearchTransportService(null, null, null);
-        DfsQueryPhase dqp = new DfsQueryPhase(List.of(), dkrs, mock(QueryPhaseResultConsumer.class), null, mspc);
+        DfsQueryPhase dqp = new DfsQueryPhase(mock(QueryPhaseResultConsumer.class), null, mspc);
 
         QueryBuilder bm25 = new TermQueryBuilder("field", "term");
         SearchSourceBuilder ssb = new SearchSourceBuilder().query(bm25)
@@ -356,7 +361,7 @@ public void testRewriteShardSearchRequestWithRank() {
         SearchRequest sr = new SearchRequest().allowPartialSearchResults(true).source(ssb);
         ShardSearchRequest ssr = new ShardSearchRequest(null, sr, new ShardId("test", "testuuid", 1), 1, 1, null, 1.0f, 0, null);
 
-        dqp.rewriteShardSearchRequest(ssr);
+        dqp.rewriteShardSearchRequest(dkrs, ssr);
 
         KnnScoreDocQueryBuilder ksdqb0 = new KnnScoreDocQueryBuilder(
             new ScoreDoc[] { new ScoreDoc(1, 3.0f, 1), new ScoreDoc(4, 1.5f, 1) },