Examples with FetchSearchResult org.opensearch.search.fetch.FetchSearchResult used on opensource projects

Example 1 with FetchSearchResult

use of org.opensearch.search.fetch.FetchSearchResult in project OpenSearch by opensearch-project.

the class SearchPhaseController method getHits.

private SearchHits getHits(ReducedQueryPhase reducedQueryPhase, boolean ignoreFrom, Collection<? extends SearchPhaseResult> fetchResults, IntFunction<SearchPhaseResult> resultsLookup) {
    SortedTopDocs sortedTopDocs = reducedQueryPhase.sortedTopDocs;
    int sortScoreIndex = -1;
    if (sortedTopDocs.isSortedByField) {
        SortField[] sortFields = sortedTopDocs.sortFields;
        for (int i = 0; i < sortFields.length; i++) {
            if (sortFields[i].getType() == SortField.Type.SCORE) {
                sortScoreIndex = i;
            }
        }
    }
    // clean the fetch counter
    for (SearchPhaseResult entry : fetchResults) {
        entry.fetchResult().initCounter();
    }
    int from = ignoreFrom ? 0 : reducedQueryPhase.from;
    int numSearchHits = (int) Math.min(reducedQueryPhase.fetchHits - from, reducedQueryPhase.size);
    // with collapsing we can have more fetch hits than sorted docs
    numSearchHits = Math.min(sortedTopDocs.scoreDocs.length, numSearchHits);
    // merge hits
    List<SearchHit> hits = new ArrayList<>();
    if (!fetchResults.isEmpty()) {
        for (int i = 0; i < numSearchHits; i++) {
            ScoreDoc shardDoc = sortedTopDocs.scoreDocs[i];
            SearchPhaseResult fetchResultProvider = resultsLookup.apply(shardDoc.shardIndex);
            if (fetchResultProvider == null) {
                // TODO it would be nice to assert this in the future
                continue;
            }
            FetchSearchResult fetchResult = fetchResultProvider.fetchResult();
            final int index = fetchResult.counterGetAndIncrement();
            assert index < fetchResult.hits().getHits().length : "not enough hits fetched. index [" + index + "] length: " + fetchResult.hits().getHits().length;
            SearchHit searchHit = fetchResult.hits().getHits()[index];
            searchHit.shard(fetchResult.getSearchShardTarget());
            if (sortedTopDocs.isSortedByField) {
                FieldDoc fieldDoc = (FieldDoc) shardDoc;
                searchHit.sortValues(fieldDoc.fields, reducedQueryPhase.sortValueFormats);
                if (sortScoreIndex != -1) {
                    searchHit.score(((Number) fieldDoc.fields[sortScoreIndex]).floatValue());
                }
            } else {
                searchHit.score(shardDoc.score);
            }
            hits.add(searchHit);
        }
    }
    return new SearchHits(hits.toArray(new SearchHit[0]), reducedQueryPhase.totalHits, reducedQueryPhase.maxScore, sortedTopDocs.sortFields, sortedTopDocs.collapseField, sortedTopDocs.collapseValues);
}

Example 2 with FetchSearchResult

use of org.opensearch.search.fetch.FetchSearchResult in project OpenSearch by opensearch-project.

the class FetchSearchPhaseTests method testFailFetchOneDoc.

public void testFailFetchOneDoc() {
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    SearchPhaseController controller = new SearchPhaseController(writableRegistry(), s -> InternalAggregationTestCase.emptyReduceContextBuilder());
    QueryPhaseResultConsumer results = controller.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.getRequest(), 2, exc -> {
    });
    int resultSetSize = randomIntBetween(2, 10);
    final ShardSearchContextId ctx = new ShardSearchContextId(UUIDs.base64UUID(), 123);
    QuerySearchResult queryResult = new QuerySearchResult(ctx, 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(resultSetSize);
    queryResult.setShardIndex(0);
    results.consumeResult(queryResult, () -> {
    });
    queryResult = new QuerySearchResult(new ShardSearchContextId("", 321), new SearchShardTarget("node2", new ShardId("test", "na", 1), 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]);
    queryResult.size(resultSetSize);
    queryResult.setShardIndex(1);
    results.consumeResult(queryResult, () -> {
    });
    mockSearchPhaseContext.searchTransport = new SearchTransportService(null, null) {

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

        @Override
        public void run() {
            mockSearchPhaseContext.sendSearchResponse(searchResponse, null);
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    SearchResponse searchResponse = mockSearchPhaseContext.searchResponse.get();
    assertNotNull(searchResponse);
    assertEquals(2, searchResponse.getHits().getTotalHits().value);
    assertEquals(84, searchResponse.getHits().getAt(0).docId());
    assertEquals(1, searchResponse.getFailedShards());
    assertEquals(1, searchResponse.getSuccessfulShards());
    assertEquals(1, searchResponse.getShardFailures().length);
    assertTrue(searchResponse.getShardFailures()[0].getCause() instanceof MockDirectoryWrapper.FakeIOException);
    assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(ctx));
}

Example 3 with FetchSearchResult

use of org.opensearch.search.fetch.FetchSearchResult in project OpenSearch by opensearch-project.

the class FetchSearchPhaseTests method testFetchTwoDocument.

public void testFetchTwoDocument() {
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    SearchPhaseController controller = new SearchPhaseController(writableRegistry(), s -> InternalAggregationTestCase.emptyReduceContextBuilder());
    QueryPhaseResultConsumer results = controller.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.getRequest(), 2, exc -> {
    });
    int resultSetSize = randomIntBetween(2, 10);
    ShardSearchContextId ctx1 = new ShardSearchContextId(UUIDs.base64UUID(), 123);
    QuerySearchResult queryResult = new QuerySearchResult(ctx1, 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(resultSetSize);
    queryResult.setShardIndex(0);
    results.consumeResult(queryResult, () -> {
    });
    final ShardSearchContextId ctx2 = new ShardSearchContextId(UUIDs.base64UUID(), 321);
    queryResult = new QuerySearchResult(ctx2, new SearchShardTarget("node2", new ShardId("test", "na", 1), 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]);
    queryResult.size(resultSetSize);
    queryResult.setShardIndex(1);
    results.consumeResult(queryResult, () -> {
    });
    mockSearchPhaseContext.searchTransport = new SearchTransportService(null, null) {

        @Override
        public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRequest request, SearchTask task, SearchActionListener<FetchSearchResult> listener) {
            FetchSearchResult fetchResult = new FetchSearchResult();
            if (request.contextId().equals(ctx2)) {
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(84) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 2.0F));
            } else {
                assertEquals(ctx1, request.contextId());
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(42) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0F));
            }
            listener.onResponse(fetchResult);
        }
    };
    FetchSearchPhase phase = new FetchSearchPhase(results, controller, null, mockSearchPhaseContext, (searchResponse, scrollId) -> new SearchPhase("test") {

        @Override
        public void run() {
            mockSearchPhaseContext.sendSearchResponse(searchResponse, null);
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    SearchResponse searchResponse = mockSearchPhaseContext.searchResponse.get();
    assertNotNull(searchResponse);
    assertEquals(2, searchResponse.getHits().getTotalHits().value);
    assertEquals(84, searchResponse.getHits().getAt(0).docId());
    assertEquals(42, searchResponse.getHits().getAt(1).docId());
    assertEquals(0, searchResponse.getFailedShards());
    assertEquals(2, searchResponse.getSuccessfulShards());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
}

Example 4 with FetchSearchResult

use of org.opensearch.search.fetch.FetchSearchResult in project OpenSearch by opensearch-project.

the class FetchSearchPhaseTests method testExceptionFailsPhase.

public void testExceptionFailsPhase() {
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    SearchPhaseController controller = new SearchPhaseController(writableRegistry(), s -> InternalAggregationTestCase.emptyReduceContextBuilder());
    QueryPhaseResultConsumer results = controller.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.getRequest(), 2, exc -> {
    });
    int resultSetSize = randomIntBetween(2, 10);
    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(resultSetSize);
    queryResult.setShardIndex(0);
    results.consumeResult(queryResult, () -> {
    });
    queryResult = new QuerySearchResult(new ShardSearchContextId("", 321), new SearchShardTarget("node2", new ShardId("test", "na", 1), 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]);
    queryResult.size(resultSetSize);
    queryResult.setShardIndex(1);
    results.consumeResult(queryResult, () -> {
    });
    AtomicInteger numFetches = new AtomicInteger(0);
    mockSearchPhaseContext.searchTransport = new SearchTransportService(null, null) {

        @Override
        public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRequest request, SearchTask task, SearchActionListener<FetchSearchResult> listener) {
            FetchSearchResult fetchResult = new FetchSearchResult();
            if (numFetches.incrementAndGet() == 1) {
                throw new RuntimeException("BOOM");
            }
            if (request.contextId().getId() == 321) {
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(84) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 2.0F));
            } else {
                assertEquals(request, 123);
                fetchResult.hits(new SearchHits(new SearchHit[] { new SearchHit(42) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0F));
            }
            listener.onResponse(fetchResult);
        }
    };
    FetchSearchPhase phase = new FetchSearchPhase(results, controller, null, mockSearchPhaseContext, (searchResponse, scrollId) -> new SearchPhase("test") {

        @Override
        public void run() {
            mockSearchPhaseContext.sendSearchResponse(searchResponse, null);
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    assertNotNull(mockSearchPhaseContext.phaseFailure.get());
    assertEquals(mockSearchPhaseContext.phaseFailure.get().getMessage(), "BOOM");
    assertNull(mockSearchPhaseContext.searchResponse.get());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
}

Example 5 with FetchSearchResult

use of org.opensearch.search.fetch.FetchSearchResult in project OpenSearch by opensearch-project.

the class FetchSearchPhaseTests method testCleanupIrrelevantContexts.

public void testCleanupIrrelevantContexts() {
    // contexts that are not fetched should be cleaned up
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    SearchPhaseController controller = new SearchPhaseController(writableRegistry(), s -> InternalAggregationTestCase.emptyReduceContextBuilder());
    QueryPhaseResultConsumer results = controller.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.getRequest(), 2, exc -> {
    });
    int resultSetSize = 1;
    final ShardSearchContextId ctx1 = new ShardSearchContextId(UUIDs.base64UUID(), 123);
    QuerySearchResult queryResult = new QuerySearchResult(ctx1, 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(resultSetSize);
    queryResult.setShardIndex(0);
    results.consumeResult(queryResult, () -> {
    });
    final ShardSearchContextId ctx2 = new ShardSearchContextId(UUIDs.base64UUID(), 321);
    queryResult = new QuerySearchResult(ctx2, new SearchShardTarget("node2", new ShardId("test", "na", 1), 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]);
    queryResult.size(resultSetSize);
    queryResult.setShardIndex(1);
    results.consumeResult(queryResult, () -> {
    });
    mockSearchPhaseContext.searchTransport = new SearchTransportService(null, null) {

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

        @Override
        public void run() {
            mockSearchPhaseContext.sendSearchResponse(searchResponse, null);
        }
    });
    assertEquals("fetch", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    SearchResponse searchResponse = mockSearchPhaseContext.searchResponse.get();
    assertNotNull(searchResponse);
    assertEquals(2, searchResponse.getHits().getTotalHits().value);
    assertEquals(1, searchResponse.getHits().getHits().length);
    assertEquals(84, searchResponse.getHits().getAt(0).docId());
    assertEquals(0, searchResponse.getFailedShards());
    assertEquals(2, searchResponse.getSuccessfulShards());
    assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(ctx1));
}