Examples with DfsSearchResult org.opensearch.search.dfs.DfsSearchResult used on opensource projects

Example 1 with DfsSearchResult

use of org.opensearch.search.dfs.DfsSearchResult in project OpenSearch by opensearch-project.

the class SearchPhaseController method aggregateDfs.

public AggregatedDfs aggregateDfs(Collection<DfsSearchResult> results) {
    ObjectObjectHashMap<Term, TermStatistics> termStatistics = HppcMaps.newNoNullKeysMap();
    ObjectObjectHashMap<String, CollectionStatistics> fieldStatistics = HppcMaps.newNoNullKeysMap();
    long aggMaxDoc = 0;
    for (DfsSearchResult lEntry : results) {
        final Term[] terms = lEntry.terms();
        final TermStatistics[] stats = lEntry.termStatistics();
        assert terms.length == stats.length;
        for (int i = 0; i < terms.length; i++) {
            assert terms[i] != null;
            if (stats[i] == null) {
                continue;
            }
            TermStatistics existing = termStatistics.get(terms[i]);
            if (existing != null) {
                assert terms[i].bytes().equals(existing.term());
                termStatistics.put(terms[i], new TermStatistics(existing.term(), existing.docFreq() + stats[i].docFreq(), existing.totalTermFreq() + stats[i].totalTermFreq()));
            } else {
                termStatistics.put(terms[i], stats[i]);
            }
        }
        assert !lEntry.fieldStatistics().containsKey(null);
        final Object[] keys = lEntry.fieldStatistics().keys;
        final Object[] values = lEntry.fieldStatistics().values;
        for (int i = 0; i < keys.length; i++) {
            if (keys[i] != null) {
                String key = (String) keys[i];
                CollectionStatistics value = (CollectionStatistics) values[i];
                if (value == null) {
                    continue;
                }
                assert key != null;
                CollectionStatistics existing = fieldStatistics.get(key);
                if (existing != null) {
                    CollectionStatistics merged = new CollectionStatistics(key, existing.maxDoc() + value.maxDoc(), existing.docCount() + value.docCount(), existing.sumTotalTermFreq() + value.sumTotalTermFreq(), existing.sumDocFreq() + value.sumDocFreq());
                    fieldStatistics.put(key, merged);
                } else {
                    fieldStatistics.put(key, value);
                }
            }
        }
        aggMaxDoc += lEntry.maxDoc();
    }
    return new AggregatedDfs(termStatistics, fieldStatistics, aggMaxDoc);
}

Example 2 with DfsSearchResult

use of org.opensearch.search.dfs.DfsSearchResult in project OpenSearch by opensearch-project.

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 CountedCollector<SearchPhaseResult> counter = new CountedCollector<>(queryResult, searchResults.size(), () -> context.executeNextPhase(this, nextPhaseFactory.apply(queryResult)), context);
    for (final DfsSearchResult dfsResult : searchResults) {
        final SearchShardTarget searchShardTarget = dfsResult.getSearchShardTarget();
        Transport.Connection connection = context.getConnection(searchShardTarget.getClusterAlias(), searchShardTarget.getNodeId());
        QuerySearchRequest querySearchRequest = new QuerySearchRequest(searchShardTarget.getOriginalIndices(), dfsResult.getContextId(), dfsResult.getShardSearchRequest(), dfs);
        final int shardIndex = dfsResult.getShardIndex();
        searchTransportService.sendExecuteQuery(connection, querySearchRequest, context.getTask(), new SearchActionListener<QuerySearchResult>(searchShardTarget, shardIndex) {

            @Override
            protected void innerOnResponse(QuerySearchResult response) {
                try {
                    counter.onResult(response);
                } catch (Exception e) {
                    context.onPhaseFailure(DfsQueryPhase.this, "", e);
                }
            }

            @Override
            public void onFailure(Exception exception) {
                try {
                    context.getLogger().debug(() -> new ParameterizedMessage("[{}] Failed to execute query phase", querySearchRequest.contextId()), exception);
                    progressListener.notifyQueryFailure(shardIndex, searchShardTarget, exception);
                    counter.onFailure(shardIndex, searchShardTarget, exception);
                } finally {
                    if (context.getRequest().pointInTimeBuilder() == null) {
                        // the query might not have been executed at all (for example because thread pool rejected
                        // execution) and the search context that was created in dfs phase might not be released.
                        // release it again to be in the safe side
                        context.sendReleaseSearchContext(querySearchRequest.contextId(), connection, searchShardTarget.getOriginalIndices());
                    }
                }
            }
        });
    }
}

Example 3 with DfsSearchResult

use of org.opensearch.search.dfs.DfsSearchResult in project OpenSearch by opensearch-project.

the class CountedCollectorTests method testCollect.

public void testCollect() throws InterruptedException {
    ArraySearchPhaseResults<SearchPhaseResult> consumer = new ArraySearchPhaseResults<>(randomIntBetween(1, 100));
    List<Integer> state = new ArrayList<>();
    int numResultsExpected = randomIntBetween(1, consumer.getAtomicArray().length());
    MockSearchPhaseContext context = new MockSearchPhaseContext(consumer.getAtomicArray().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<>(consumer, 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(() -> {
                    DfsSearchResult dfsSearchResult = new DfsSearchResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), shardID), null, null);
                    dfsSearchResult.setShardIndex(shardID);
                    dfsSearchResult.setSearchShardTarget(new SearchShardTarget("foo", new ShardId("bar", "baz", shardID), null, OriginalIndices.NONE));
                    collector.onResult(dfsSearchResult);
                });
                break;
            case 2:
                state.add(2);
                executor.execute(() -> collector.onFailure(shardID, new SearchShardTarget("foo", new ShardId("bar", "baz", shardID), null, OriginalIndices.NONE), new RuntimeException("boom")));
                break;
            default:
                fail("unknown state");
        }
    }
    latch.await();
    assertEquals(numResultsExpected, state.size());
    AtomicArray<SearchPhaseResult> results = consumer.getAtomicArray();
    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).getContextId().getId());
                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));
    }
}

Example 4 with DfsSearchResult

use of org.opensearch.search.dfs.DfsSearchResult in project OpenSearch by opensearch-project.

the class DfsQueryPhaseTests method testFailPhaseOnException.

public void testFailPhaseOnException() throws IOException {
    AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
    AtomicReference<AtomicArray<SearchPhaseResult>> responseRef = new AtomicReference<>();
    results.set(0, newSearchResult(0, new ShardSearchContextId("", 1), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.set(1, newSearchResult(1, new ShardSearchContextId("", 2), new SearchShardTarget("node2", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
    results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
    SearchTransportService searchTransportService = new SearchTransportService(null, null) {

        @Override
        public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, SearchActionListener<QuerySearchResult> listener) {
            if (request.contextId().getId() == 1) {
                QuerySearchResult queryResult = new QuerySearchResult(new ShardSearchContextId("", 123), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE), null);
                queryResult.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), 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.contextId().getId() == 2) {
                throw new UncheckedIOException(new MockDirectoryWrapper.FakeIOException());
            } else {
                fail("no such request ID: " + request.contextId());
            }
        }
    };
    SearchPhaseController searchPhaseController = searchPhaseController();
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    mockSearchPhaseContext.searchTransport = searchTransportService;
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.searchRequest, results.length(), exc -> {
    });
    DfsQueryPhase phase = new DfsQueryPhase(results.asList(), null, consumer, (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());
}

Example 5 with DfsSearchResult

use of org.opensearch.search.dfs.DfsSearchResult in project OpenSearch by opensearch-project.

the class DfsQueryPhaseTests method newSearchResult.

private static DfsSearchResult newSearchResult(int shardIndex, ShardSearchContextId contextId, SearchShardTarget target) {
    DfsSearchResult result = new DfsSearchResult(contextId, target, null);
    result.setShardIndex(shardIndex);
    return result;
}