Search in sources :

Example 1 with DfsSearchResult

use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.

the class DfsQueryPhase method run.

@Override
public void run() throws IOException {
    // 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 AggregatedDfs dfs = searchPhaseController.aggregateDfs(dfsSearchResults);
    final CountedCollector<QuerySearchResultProvider> counter = new CountedCollector<>(queryResult::consumeResult, dfsSearchResults.asList().size(), () -> {
        context.executeNextPhase(this, nextPhaseFactory.apply(queryResult));
    }, context);
    for (final AtomicArray.Entry<DfsSearchResult> entry : dfsSearchResults.asList()) {
        DfsSearchResult dfsResult = entry.value;
        final int shardIndex = entry.index;
        final SearchShardTarget searchShardTarget = dfsResult.shardTarget();
        Transport.Connection connection = context.getConnection(searchShardTarget.getNodeId());
        QuerySearchRequest querySearchRequest = new QuerySearchRequest(context.getRequest(), dfsResult.id(), dfs);
        searchTransportService.sendExecuteQuery(connection, querySearchRequest, context.getTask(), ActionListener.wrap(result -> counter.onResult(shardIndex, result, searchShardTarget), exception -> {
            try {
                if (context.getLogger().isDebugEnabled()) {
                    context.getLogger().debug((Supplier<?>) () -> new ParameterizedMessage("[{}] Failed to execute query phase", querySearchRequest.id()), exception);
                }
                counter.onFailure(shardIndex, searchShardTarget, exception);
            } finally {
                context.sendReleaseSearchContext(querySearchRequest.id(), connection);
            }
        }));
    }
}
Also used : SearchShardTarget(org.elasticsearch.search.SearchShardTarget) Transport(org.elasticsearch.transport.Transport) Supplier(org.apache.logging.log4j.util.Supplier) AtomicArray(org.elasticsearch.common.util.concurrent.AtomicArray) AggregatedDfs(org.elasticsearch.search.dfs.AggregatedDfs) IOException(java.io.IOException) ParameterizedMessage(org.apache.logging.log4j.message.ParameterizedMessage) QuerySearchRequest(org.elasticsearch.search.query.QuerySearchRequest) Function(java.util.function.Function) ActionListener(org.elasticsearch.action.ActionListener) DfsSearchResult(org.elasticsearch.search.dfs.DfsSearchResult) QuerySearchResultProvider(org.elasticsearch.search.query.QuerySearchResultProvider) AtomicArray(org.elasticsearch.common.util.concurrent.AtomicArray) QuerySearchResultProvider(org.elasticsearch.search.query.QuerySearchResultProvider) DfsSearchResult(org.elasticsearch.search.dfs.DfsSearchResult) AggregatedDfs(org.elasticsearch.search.dfs.AggregatedDfs) SearchShardTarget(org.elasticsearch.search.SearchShardTarget) QuerySearchRequest(org.elasticsearch.search.query.QuerySearchRequest) Supplier(org.apache.logging.log4j.util.Supplier) ParameterizedMessage(org.apache.logging.log4j.message.ParameterizedMessage) Transport(org.elasticsearch.transport.Transport)

Example 2 with DfsSearchResult

use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.

the class SearchPhaseController method aggregateDfs.

public AggregatedDfs aggregateDfs(AtomicArray<DfsSearchResult> results) {
    ObjectObjectHashMap<Term, TermStatistics> termStatistics = HppcMaps.newNoNullKeysMap();
    ObjectObjectHashMap<String, CollectionStatistics> fieldStatistics = HppcMaps.newNoNullKeysMap();
    long aggMaxDoc = 0;
    for (AtomicArray.Entry<DfsSearchResult> lEntry : results.asList()) {
        final Term[] terms = lEntry.value.terms();
        final TermStatistics[] stats = lEntry.value.termStatistics();
        assert terms.length == stats.length;
        for (int i = 0; i < terms.length; i++) {
            assert terms[i] != null;
            TermStatistics existing = termStatistics.get(terms[i]);
            if (existing != null) {
                assert terms[i].bytes().equals(existing.term());
                // totalTermFrequency is an optional statistic we need to check if either one or both
                // are set to -1 which means not present and then set it globally to -1
                termStatistics.put(terms[i], new TermStatistics(existing.term(), existing.docFreq() + stats[i].docFreq(), optionalSum(existing.totalTermFreq(), stats[i].totalTermFreq())));
            } else {
                termStatistics.put(terms[i], stats[i]);
            }
        }
        assert !lEntry.value.fieldStatistics().containsKey(null);
        final Object[] keys = lEntry.value.fieldStatistics().keys;
        final Object[] values = lEntry.value.fieldStatistics().values;
        for (int i = 0; i < keys.length; i++) {
            if (keys[i] != null) {
                String key = (String) keys[i];
                CollectionStatistics value = (CollectionStatistics) values[i];
                assert key != null;
                CollectionStatistics existing = fieldStatistics.get(key);
                if (existing != null) {
                    CollectionStatistics merged = new CollectionStatistics(key, existing.maxDoc() + value.maxDoc(), optionalSum(existing.docCount(), value.docCount()), optionalSum(existing.sumTotalTermFreq(), value.sumTotalTermFreq()), optionalSum(existing.sumDocFreq(), value.sumDocFreq()));
                    fieldStatistics.put(key, merged);
                } else {
                    fieldStatistics.put(key, value);
                }
            }
        }
        aggMaxDoc += lEntry.value.maxDoc();
    }
    return new AggregatedDfs(termStatistics, fieldStatistics, aggMaxDoc);
}
Also used : AtomicArray(org.elasticsearch.common.util.concurrent.AtomicArray) DfsSearchResult(org.elasticsearch.search.dfs.DfsSearchResult) Term(org.apache.lucene.index.Term) TermStatistics(org.apache.lucene.search.TermStatistics) CollectionStatistics(org.apache.lucene.search.CollectionStatistics) AggregatedDfs(org.elasticsearch.search.dfs.AggregatedDfs)

Example 3 with DfsSearchResult

use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.

the class DfsQueryPhaseTests method testFailPhaseOnException.

public void testFailPhaseOnException() 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) {
                throw new UncheckedIOException(new MockDirectoryWrapper.FakeIOException());
            } 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());
    expectThrows(UncheckedIOException.class, () -> phase.run());
    // phase execution will clean up on the contexts
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
}
Also used : AtomicArray(org.elasticsearch.common.util.concurrent.AtomicArray) DfsSearchResult(org.elasticsearch.search.dfs.DfsSearchResult) AtomicReference(java.util.concurrent.atomic.AtomicReference) Index(org.elasticsearch.index.Index) UncheckedIOException(java.io.UncheckedIOException) IOException(java.io.IOException) UncheckedIOException(java.io.UncheckedIOException) ScoreDoc(org.apache.lucene.search.ScoreDoc) TopDocs(org.apache.lucene.search.TopDocs) ActionListener(org.elasticsearch.action.ActionListener) QuerySearchResult(org.elasticsearch.search.query.QuerySearchResult) SearchShardTarget(org.elasticsearch.search.SearchShardTarget) QuerySearchRequest(org.elasticsearch.search.query.QuerySearchRequest)

Example 4 with DfsSearchResult

use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.

the class CountedCollectorTests method testCollect.

public void testCollect() throws InterruptedException {
    AtomicArray<SearchPhaseResult> results = new AtomicArray<>(randomIntBetween(1, 100));
    List<Integer> state = new ArrayList<>();
    int numResultsExpected = randomIntBetween(1, results.length());
    MockSearchPhaseContext context = new MockSearchPhaseContext(results.length());
    CountDownLatch latch = new CountDownLatch(1);
    boolean maybeFork = randomBoolean();
    Executor executor = (runnable) -> {
        if (randomBoolean() && maybeFork) {
            new Thread(runnable).start();
        } else {
            runnable.run();
        }
    };
    CountedCollector<SearchPhaseResult> collector = new CountedCollector<>(results::set, numResultsExpected, latch::countDown, context);
    for (int i = 0; i < numResultsExpected; i++) {
        int shardID = i;
        switch(randomIntBetween(0, 2)) {
            case 0:
                state.add(0);
                executor.execute(() -> collector.countDown());
                break;
            case 1:
                state.add(1);
                executor.execute(() -> collector.onResult(shardID, new DfsSearchResult(shardID, null), new SearchShardTarget("foo", new Index("bar", "baz"), shardID)));
                break;
            case 2:
                state.add(2);
                executor.execute(() -> collector.onFailure(shardID, new SearchShardTarget("foo", new Index("bar", "baz"), shardID), new RuntimeException("boom")));
                break;
            default:
                fail("unknown state");
        }
    }
    latch.await();
    assertEquals(numResultsExpected, state.size());
    for (int i = 0; i < numResultsExpected; i++) {
        switch(state.get(i)) {
            case 0:
                assertNull(results.get(i));
                break;
            case 1:
                assertNotNull(results.get(i));
                assertEquals(i, results.get(i).id());
                break;
            case 2:
                final int shardId = i;
                assertEquals(1, context.failures.stream().filter(f -> f.shardId() == shardId).count());
                break;
            default:
                fail("unknown state");
        }
    }
    for (int i = numResultsExpected; i < results.length(); i++) {
        assertNull("index: " + i, results.get(i));
    }
}
Also used : SearchShardTarget(org.elasticsearch.search.SearchShardTarget) CountDownLatch(java.util.concurrent.CountDownLatch) List(java.util.List) SearchPhaseResult(org.elasticsearch.search.SearchPhaseResult) Executor(java.util.concurrent.Executor) AtomicArray(org.elasticsearch.common.util.concurrent.AtomicArray) Index(org.elasticsearch.index.Index) ESTestCase(org.elasticsearch.test.ESTestCase) DfsSearchResult(org.elasticsearch.search.dfs.DfsSearchResult) ArrayList(java.util.ArrayList) AtomicArray(org.elasticsearch.common.util.concurrent.AtomicArray) DfsSearchResult(org.elasticsearch.search.dfs.DfsSearchResult) ArrayList(java.util.ArrayList) Index(org.elasticsearch.index.Index) CountDownLatch(java.util.concurrent.CountDownLatch) Executor(java.util.concurrent.Executor) SearchPhaseResult(org.elasticsearch.search.SearchPhaseResult) SearchShardTarget(org.elasticsearch.search.SearchShardTarget)

Example 5 with DfsSearchResult

use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.

the class DfsQueryPhaseTests method testDfsWith1ShardFailed.

public void testDfsWith1ShardFailed() 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) {
                listener.onFailure(new MockDirectoryWrapper.FakeIOException());
            } 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);
    assertNull(responseRef.get().get(1));
    assertEquals(1, mockSearchPhaseContext.numSuccess.get());
    assertEquals(1, mockSearchPhaseContext.failures.size());
    assertTrue(mockSearchPhaseContext.failures.get(0).getCause() instanceof MockDirectoryWrapper.FakeIOException);
    assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(2L));
    assertNull(responseRef.get().get(1));
}
Also used : MockDirectoryWrapper(org.apache.lucene.store.MockDirectoryWrapper) AtomicArray(org.elasticsearch.common.util.concurrent.AtomicArray) DfsSearchResult(org.elasticsearch.search.dfs.DfsSearchResult) AtomicReference(java.util.concurrent.atomic.AtomicReference) Index(org.elasticsearch.index.Index) IOException(java.io.IOException) UncheckedIOException(java.io.UncheckedIOException) ScoreDoc(org.apache.lucene.search.ScoreDoc) TopDocs(org.apache.lucene.search.TopDocs) ActionListener(org.elasticsearch.action.ActionListener) QuerySearchResult(org.elasticsearch.search.query.QuerySearchResult) SearchShardTarget(org.elasticsearch.search.SearchShardTarget) QuerySearchRequest(org.elasticsearch.search.query.QuerySearchRequest)

Aggregations

AtomicArray (org.elasticsearch.common.util.concurrent.AtomicArray)6 DfsSearchResult (org.elasticsearch.search.dfs.DfsSearchResult)6 SearchShardTarget (org.elasticsearch.search.SearchShardTarget)5 IOException (java.io.IOException)4 ActionListener (org.elasticsearch.action.ActionListener)4 Index (org.elasticsearch.index.Index)4 QuerySearchRequest (org.elasticsearch.search.query.QuerySearchRequest)4 UncheckedIOException (java.io.UncheckedIOException)3 AtomicReference (java.util.concurrent.atomic.AtomicReference)3 ScoreDoc (org.apache.lucene.search.ScoreDoc)3 TopDocs (org.apache.lucene.search.TopDocs)3 QuerySearchResult (org.elasticsearch.search.query.QuerySearchResult)3 AggregatedDfs (org.elasticsearch.search.dfs.AggregatedDfs)2 ArrayList (java.util.ArrayList)1 List (java.util.List)1 CountDownLatch (java.util.concurrent.CountDownLatch)1 Executor (java.util.concurrent.Executor)1 Function (java.util.function.Function)1 ParameterizedMessage (org.apache.logging.log4j.message.ParameterizedMessage)1 Supplier (org.apache.logging.log4j.util.Supplier)1