Examples with AtomicArray org.elasticsearch.common.util.concurrent.AtomicArray used on opensource projects

Example 16 with AtomicArray

use of org.elasticsearch.common.util.concurrent.AtomicArray in project elasticsearch by elastic.

the class DfsQueryPhaseTests method testDfsWith2Shards.

public void testDfsWith2Shards() throws IOException {
    AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
    AtomicReference<AtomicArray<QuerySearchResultProvider>> responseRef = new AtomicReference<>();
    results.set(0, new DfsSearchResult(1, new SearchShardTarget("node1", new Index("test", "na"), 0)));
    results.set(1, new DfsSearchResult(2, new SearchShardTarget("node2", new Index("test", "na"), 0)));
    results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
    results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
    SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
    SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {

        @Override
        public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, ActionListener<QuerySearchResult> listener) {
            if (request.id() == 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(2);
                listener.onResponse(queryResult);
            } else if (request.id() == 2) {
                QuerySearchResult queryResult = new QuerySearchResult(123, new SearchShardTarget("node2", new Index("test", "na"), 0));
                queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(84, 2.0F) }, 2.0F), new DocValueFormat[0]);
                // the size of the result set
                queryResult.size(2);
                listener.onResponse(queryResult);
            } else {
                fail("no such request ID: " + request.id());
            }
        }
    };
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    mockSearchPhaseContext.searchTransport = searchTransportService;
    DfsQueryPhase phase = new DfsQueryPhase(results, controller, (response) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(response.results);
        }
    }, mockSearchPhaseContext);
    assertEquals("dfs_query", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    assertNotNull(responseRef.get());
    assertNotNull(responseRef.get().get(0));
    assertNull(responseRef.get().get(0).fetchResult());
    assertEquals(1, responseRef.get().get(0).queryResult().topDocs().totalHits);
    assertEquals(42, responseRef.get().get(0).queryResult().topDocs().scoreDocs[0].doc);
    assertNotNull(responseRef.get().get(1));
    assertNull(responseRef.get().get(1).fetchResult());
    assertEquals(1, responseRef.get().get(1).queryResult().topDocs().totalHits);
    assertEquals(84, responseRef.get().get(1).queryResult().topDocs().scoreDocs[0].doc);
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
    assertEquals(2, mockSearchPhaseContext.numSuccess.get());
}

Example 17 with AtomicArray

use of org.elasticsearch.common.util.concurrent.AtomicArray in project elasticsearch by elastic.

the class SearchPhaseControllerTests method generateQueryResults.

private AtomicArray<QuerySearchResultProvider> generateQueryResults(int nShards, List<CompletionSuggestion> suggestions, int searchHitsSize, boolean useConstantScore) {
    AtomicArray<QuerySearchResultProvider> queryResults = new AtomicArray<>(nShards);
    for (int shardIndex = 0; shardIndex < nShards; shardIndex++) {
        QuerySearchResult querySearchResult = new QuerySearchResult(shardIndex, new SearchShardTarget("", new Index("", ""), shardIndex));
        TopDocs topDocs = new TopDocs(0, new ScoreDoc[0], 0);
        if (searchHitsSize > 0) {
            int nDocs = randomIntBetween(0, searchHitsSize);
            ScoreDoc[] scoreDocs = new ScoreDoc[nDocs];
            float maxScore = 0F;
            for (int i = 0; i < nDocs; i++) {
                float score = useConstantScore ? 1.0F : Math.abs(randomFloat());
                scoreDocs[i] = new ScoreDoc(i, score);
                if (score > maxScore) {
                    maxScore = score;
                }
            }
            topDocs = new TopDocs(scoreDocs.length, scoreDocs, maxScore);
        }
        List<CompletionSuggestion> shardSuggestion = new ArrayList<>();
        for (CompletionSuggestion completionSuggestion : suggestions) {
            CompletionSuggestion suggestion = new CompletionSuggestion(completionSuggestion.getName(), completionSuggestion.getSize());
            final CompletionSuggestion.Entry completionEntry = new CompletionSuggestion.Entry(new Text(""), 0, 5);
            suggestion.addTerm(completionEntry);
            int optionSize = randomIntBetween(1, suggestion.getSize());
            float maxScore = randomIntBetween(suggestion.getSize(), (int) Float.MAX_VALUE);
            for (int i = 0; i < optionSize; i++) {
                completionEntry.addOption(new CompletionSuggestion.Entry.Option(i, new Text(""), maxScore, Collections.emptyMap()));
                float dec = randomIntBetween(0, optionSize);
                if (dec <= maxScore) {
                    maxScore -= dec;
                }
            }
            suggestion.setShardIndex(shardIndex);
            shardSuggestion.add(suggestion);
        }
        querySearchResult.topDocs(topDocs, null);
        querySearchResult.size(searchHitsSize);
        querySearchResult.suggest(new Suggest(new ArrayList<>(shardSuggestion)));
        queryResults.set(shardIndex, querySearchResult);
    }
    return queryResults;
}

Example 18 with AtomicArray

use of org.elasticsearch.common.util.concurrent.AtomicArray in project elasticsearch by elastic.

the class TransportMultiSearchTemplateAction method doExecute.

@Override
protected void doExecute(MultiSearchTemplateRequest request, ActionListener<MultiSearchTemplateResponse> listener) {
    final AtomicArray<MultiSearchTemplateResponse.Item> responses = new AtomicArray<>(request.requests().size());
    final AtomicInteger counter = new AtomicInteger(responses.length());
    for (int i = 0; i < responses.length(); i++) {
        final int index = i;
        searchTemplateAction.execute(request.requests().get(i), new ActionListener<SearchTemplateResponse>() {

            @Override
            public void onResponse(SearchTemplateResponse searchTemplateResponse) {
                responses.set(index, new MultiSearchTemplateResponse.Item(searchTemplateResponse, null));
                if (counter.decrementAndGet() == 0) {
                    finishHim();
                }
            }

            @Override
            public void onFailure(Exception e) {
                responses.set(index, new MultiSearchTemplateResponse.Item(null, e));
                if (counter.decrementAndGet() == 0) {
                    finishHim();
                }
            }

            private void finishHim() {
                MultiSearchTemplateResponse.Item[] items = responses.toArray(new MultiSearchTemplateResponse.Item[responses.length()]);
                listener.onResponse(new MultiSearchTemplateResponse(items));
            }
        });
    }
}