Examples with ShardFetchSearchRequest org.elasticsearch.search.fetch.ShardFetchSearchRequest used on opensource projects

Example 1 with ShardFetchSearchRequest

use of org.elasticsearch.search.fetch.ShardFetchSearchRequest in project elasticsearch by elastic.

the class FetchSearchPhase method innerRun.

private void innerRun() throws IOException {
    final int numShards = context.getNumShards();
    final boolean isScrollSearch = context.getRequest().scroll() != null;
    ScoreDoc[] sortedShardDocs = searchPhaseController.sortDocs(isScrollSearch, queryResults);
    String scrollId = isScrollSearch ? TransportSearchHelper.buildScrollId(queryResults) : null;
    List<AtomicArray.Entry<QuerySearchResultProvider>> queryResultsAsList = queryResults.asList();
    final SearchPhaseController.ReducedQueryPhase reducedQueryPhase = resultConsumer.reduce();
    final boolean queryAndFetchOptimization = queryResults.length() == 1;
    final Runnable finishPhase = () -> moveToNextPhase(searchPhaseController, sortedShardDocs, scrollId, reducedQueryPhase, queryAndFetchOptimization ? queryResults : fetchResults);
    if (queryAndFetchOptimization) {
        assert queryResults.get(0) == null || queryResults.get(0).fetchResult() != null;
        // query AND fetch optimization
        finishPhase.run();
    } else {
        final IntArrayList[] docIdsToLoad = searchPhaseController.fillDocIdsToLoad(numShards, sortedShardDocs);
        if (sortedShardDocs.length == 0) {
            // no docs to fetch -- sidestep everything and return
            queryResultsAsList.stream().map(e -> e.value.queryResult()).forEach(// we have to release contexts here to free up resources
            this::releaseIrrelevantSearchContext);
            finishPhase.run();
        } else {
            final ScoreDoc[] lastEmittedDocPerShard = isScrollSearch ? searchPhaseController.getLastEmittedDocPerShard(reducedQueryPhase, sortedShardDocs, numShards) : null;
            final CountedCollector<FetchSearchResult> counter = new CountedCollector<>(fetchResults::set, // we count down every shard in the result no matter if we got any results or not
            docIdsToLoad.length, finishPhase, context);
            for (int i = 0; i < docIdsToLoad.length; i++) {
                IntArrayList entry = docIdsToLoad[i];
                QuerySearchResultProvider queryResult = queryResults.get(i);
                if (entry == null) {
                    // no results for this shard ID
                    if (queryResult != null) {
                        // if we got some hits from this shard we have to release the context there
                        // we do this as we go since it will free up resources and passing on the request on the
                        // transport layer is cheap.
                        releaseIrrelevantSearchContext(queryResult.queryResult());
                    }
                    // in any case we count down this result since we don't talk to this shard anymore
                    counter.countDown();
                } else {
                    Transport.Connection connection = context.getConnection(queryResult.shardTarget().getNodeId());
                    ShardFetchSearchRequest fetchSearchRequest = createFetchRequest(queryResult.queryResult().id(), i, entry, lastEmittedDocPerShard);
                    executeFetch(i, queryResult.shardTarget(), counter, fetchSearchRequest, queryResult.queryResult(), connection);
                }
            }
        }
    }
}

Example 2 with ShardFetchSearchRequest

use of org.elasticsearch.search.fetch.ShardFetchSearchRequest in project elasticsearch by elastic.

the class FetchSearchPhaseTests method testFetchTwoDocument.

public void testFetchTwoDocument() throws IOException {
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
    InitialSearchPhase.SearchPhaseResults<QuerySearchResultProvider> results = controller.newSearchPhaseResults(mockSearchPhaseContext.getRequest(), 2);
    AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
    int resultSetSize = randomIntBetween(2, 10);
    QuerySearchResult queryResult = new QuerySearchResult(123, new SearchShardTarget("node1", new Index("test", "na"), 0));
    queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(42, 1.0F) }, 2.0F), new DocValueFormat[0]);
    // the size of the result set
    queryResult.size(resultSetSize);
    results.consumeResult(0, queryResult);
    queryResult = new QuerySearchResult(321, new SearchShardTarget("node2", new Index("test", "na"), 1));
    queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(84, 2.0F) }, 2.0F), new DocValueFormat[0]);
    queryResult.size(resultSetSize);
    results.consumeResult(1, queryResult);
    SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {

        @Override
        public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRequest request, SearchTask task, ActionListener<FetchSearchResult> listener) {
            FetchSearchResult fetchResult = new FetchSearchResult();
            if (request.id() == 321) {
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(84) }, 1, 2.0F));
            } else {
                assertEquals(123, request.id());
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(42) }, 1, 1.0F));
            }
            listener.onResponse(fetchResult);
        }
    };
    mockSearchPhaseContext.searchTransport = searchTransportService;
    FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext, (searchResponse) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(searchResponse);
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    assertNotNull(responseRef.get());
    assertEquals(2, responseRef.get().getHits().totalHits);
    assertEquals(84, responseRef.get().getHits().getAt(0).docId());
    assertEquals(42, responseRef.get().getHits().getAt(1).docId());
    assertEquals(0, responseRef.get().getFailedShards());
    assertEquals(2, responseRef.get().getSuccessfulShards());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
}

Example 3 with ShardFetchSearchRequest

use of org.elasticsearch.search.fetch.ShardFetchSearchRequest in project elasticsearch by elastic.

the class FetchSearchPhaseTests method testFailFetchOneDoc.

public void testFailFetchOneDoc() throws IOException {
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
    InitialSearchPhase.SearchPhaseResults<QuerySearchResultProvider> results = controller.newSearchPhaseResults(mockSearchPhaseContext.getRequest(), 2);
    AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
    int resultSetSize = randomIntBetween(2, 10);
    QuerySearchResult queryResult = new QuerySearchResult(123, new SearchShardTarget("node1", new Index("test", "na"), 0));
    queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(42, 1.0F) }, 2.0F), new DocValueFormat[0]);
    // the size of the result set
    queryResult.size(resultSetSize);
    results.consumeResult(0, queryResult);
    queryResult = new QuerySearchResult(321, new SearchShardTarget("node2", new Index("test", "na"), 1));
    queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(84, 2.0F) }, 2.0F), new DocValueFormat[0]);
    queryResult.size(resultSetSize);
    results.consumeResult(1, queryResult);
    SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {

        @Override
        public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRequest request, SearchTask task, ActionListener<FetchSearchResult> listener) {
            if (request.id() == 321) {
                FetchSearchResult fetchResult = new FetchSearchResult();
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(84) }, 1, 2.0F));
                listener.onResponse(fetchResult);
            } else {
                listener.onFailure(new MockDirectoryWrapper.FakeIOException());
            }
        }
    };
    mockSearchPhaseContext.searchTransport = searchTransportService;
    FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext, (searchResponse) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(searchResponse);
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    assertNotNull(responseRef.get());
    assertEquals(2, responseRef.get().getHits().totalHits);
    assertEquals(84, responseRef.get().getHits().getAt(0).docId());
    assertEquals(1, responseRef.get().getFailedShards());
    assertEquals(1, responseRef.get().getSuccessfulShards());
    assertEquals(1, responseRef.get().getShardFailures().length);
    assertTrue(responseRef.get().getShardFailures()[0].getCause() instanceof MockDirectoryWrapper.FakeIOException);
    assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(123L));
}

Example 4 with ShardFetchSearchRequest

use of org.elasticsearch.search.fetch.ShardFetchSearchRequest in project elasticsearch by elastic.

the class FetchSearchPhaseTests method testCleanupIrrelevantContexts.

public void testCleanupIrrelevantContexts() throws IOException {
    // contexts that are not fetched should be cleaned up
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
    InitialSearchPhase.SearchPhaseResults<QuerySearchResultProvider> results = controller.newSearchPhaseResults(mockSearchPhaseContext.getRequest(), 2);
    AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
    int resultSetSize = 1;
    QuerySearchResult queryResult = new QuerySearchResult(123, new SearchShardTarget("node1", new Index("test", "na"), 0));
    queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(42, 1.0F) }, 2.0F), new DocValueFormat[0]);
    // the size of the result set
    queryResult.size(resultSetSize);
    results.consumeResult(0, queryResult);
    queryResult = new QuerySearchResult(321, new SearchShardTarget("node2", new Index("test", "na"), 1));
    queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(84, 2.0F) }, 2.0F), new DocValueFormat[0]);
    queryResult.size(resultSetSize);
    results.consumeResult(1, queryResult);
    SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {

        @Override
        public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRequest request, SearchTask task, ActionListener<FetchSearchResult> listener) {
            FetchSearchResult fetchResult = new FetchSearchResult();
            if (request.id() == 321) {
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(84) }, 1, 2.0F));
            } else {
                fail("requestID 123 should not be fetched but was");
            }
            listener.onResponse(fetchResult);
        }
    };
    mockSearchPhaseContext.searchTransport = searchTransportService;
    FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext, (searchResponse) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(searchResponse);
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    assertNotNull(responseRef.get());
    assertEquals(2, responseRef.get().getHits().totalHits);
    assertEquals(1, responseRef.get().getHits().internalHits().length);
    assertEquals(84, responseRef.get().getHits().getAt(0).docId());
    assertEquals(0, responseRef.get().getFailedShards());
    assertEquals(2, responseRef.get().getSuccessfulShards());
    assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(123L));
}

Example 5 with ShardFetchSearchRequest

use of org.elasticsearch.search.fetch.ShardFetchSearchRequest in project elasticsearch by elastic.

the class FetchSearchPhaseTests method testFetchDocsConcurrently.

public void testFetchDocsConcurrently() throws IOException, InterruptedException {
    int resultSetSize = randomIntBetween(0, 100);
    // we use at least 2 hits otherwise this is subject to single shard optimization and we trip an assert...
    // also numshards --> 1 hit per shard
    int numHits = randomIntBetween(2, 100);
    SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(numHits);
    InitialSearchPhase.SearchPhaseResults<QuerySearchResultProvider> results = controller.newSearchPhaseResults(mockSearchPhaseContext.getRequest(), numHits);
    AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
    for (int i = 0; i < numHits; i++) {
        QuerySearchResult queryResult = new QuerySearchResult(i, new SearchShardTarget("node1", new Index("test", "na"), 0));
        queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(i + 1, i) }, i), new DocValueFormat[0]);
        // the size of the result set
        queryResult.size(resultSetSize);
        results.consumeResult(i, queryResult);
    }
    SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {

        @Override
        public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRequest request, SearchTask task, ActionListener<FetchSearchResult> listener) {
            new Thread(() -> {
                FetchSearchResult fetchResult = new FetchSearchResult();
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit((int) (request.id() + 1)) }, 1, 100F));
                listener.onResponse(fetchResult);
            }).start();
        }
    };
    mockSearchPhaseContext.searchTransport = searchTransportService;
    CountDownLatch latch = new CountDownLatch(1);
    FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext, (searchResponse) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(searchResponse);
            latch.countDown();
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    latch.await();
    mockSearchPhaseContext.assertNoFailure();
    assertNotNull(responseRef.get());
    assertEquals(numHits, responseRef.get().getHits().totalHits);
    assertEquals(Math.min(numHits, resultSetSize), responseRef.get().getHits().getHits().length);
    SearchHit[] hits = responseRef.get().getHits().getHits();
    for (int i = 0; i < hits.length; i++) {
        assertNotNull(hits[i]);
        assertEquals("index: " + i, numHits - i, hits[i].docId());
        assertEquals("index: " + i, numHits - 1 - i, (int) hits[i].getScore());
    }
    assertEquals(0, responseRef.get().getFailedShards());
    assertEquals(numHits, responseRef.get().getSuccessfulShards());
    // all non fetched results will be freed
    int sizeReleasedContexts = Math.max(0, numHits - resultSetSize);
    assertEquals(mockSearchPhaseContext.releasedSearchContexts.toString(), sizeReleasedContexts, mockSearchPhaseContext.releasedSearchContexts.size());
}