Examples with QuerySearchRequest org.opensearch.search.query.QuerySearchRequest used on opensource projects

Example 1 with QuerySearchRequest

use of org.opensearch.search.query.QuerySearchRequest 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 2 with QuerySearchRequest

use of org.opensearch.search.query.QuerySearchRequest 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 3 with QuerySearchRequest

use of org.opensearch.search.query.QuerySearchRequest in project OpenSearch by opensearch-project.

the class DfsQueryPhaseTests method testDfsWith1ShardFailed.

public void testDfsWith1ShardFailed() 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) {
                listener.onFailure(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());
    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().topDocs.totalHits.value);
    assertEquals(42, responseRef.get().get(0).queryResult().topDocs().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(new ShardSearchContextId("", 2L)));
    assertNull(responseRef.get().get(1));
}

Example 4 with QuerySearchRequest

use of org.opensearch.search.query.QuerySearchRequest in project OpenSearch by opensearch-project.

the class DfsQueryPhaseTests method testDfsWith2Shards.

public void testDfsWith2Shards() 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) {
                QuerySearchResult queryResult = new QuerySearchResult(new ShardSearchContextId("", 123), new SearchShardTarget("node2", 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(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.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());
    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().topDocs.totalHits.value);
    assertEquals(42, responseRef.get().get(0).queryResult().topDocs().topDocs.scoreDocs[0].doc);
    assertNotNull(responseRef.get().get(1));
    assertNull(responseRef.get().get(1).fetchResult());
    assertEquals(1, responseRef.get().get(1).queryResult().topDocs().topDocs.totalHits.value);
    assertEquals(84, responseRef.get().get(1).queryResult().topDocs().topDocs.scoreDocs[0].doc);
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
    assertEquals(2, mockSearchPhaseContext.numSuccess.get());
}