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

Example 31 with QuerySearchResult

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

the class SearchPhaseControllerTests method testSortDocsIsIdempotent.

public void testSortDocsIsIdempotent() throws Exception {
    int nShards = randomIntBetween(1, 20);
    int queryResultSize = randomBoolean() ? 0 : randomIntBetween(1, nShards * 2);
    long randomSeed = randomLong();
    boolean useConstantScore = randomBoolean();
    AtomicArray<SearchPhaseResult> results = generateSeededQueryResults(randomSeed, nShards, Collections.emptyList(), queryResultSize, useConstantScore);
    boolean ignoreFrom = randomBoolean();
    Optional<SearchPhaseResult> first = results.asList().stream().findFirst();
    int from = 0, size = 0;
    if (first.isPresent()) {
        from = first.get().queryResult().from();
        size = first.get().queryResult().size();
    }
    List<TopDocs> topDocsList = new ArrayList<>();
    for (SearchPhaseResult result : results.asList()) {
        QuerySearchResult queryResult = result.queryResult();
        TopDocs topDocs = queryResult.consumeTopDocs().topDocs;
        topDocsList.add(topDocs);
        SearchPhaseController.setShardIndex(topDocs, result.getShardIndex());
    }
    ScoreDoc[] sortedDocs = SearchPhaseController.sortDocs(ignoreFrom, topDocsList, from, size, Collections.emptyList()).scoreDocs;
    results = generateSeededQueryResults(randomSeed, nShards, Collections.emptyList(), queryResultSize, useConstantScore);
    topDocsList = new ArrayList<>();
    for (SearchPhaseResult result : results.asList()) {
        QuerySearchResult queryResult = result.queryResult();
        TopDocs topDocs = queryResult.consumeTopDocs().topDocs;
        topDocsList.add(topDocs);
        SearchPhaseController.setShardIndex(topDocs, result.getShardIndex());
    }
    ScoreDoc[] sortedDocs2 = SearchPhaseController.sortDocs(ignoreFrom, topDocsList, from, size, Collections.emptyList()).scoreDocs;
    assertEquals(sortedDocs.length, sortedDocs2.length);
    for (int i = 0; i < sortedDocs.length; i++) {
        assertEquals(sortedDocs[i].doc, sortedDocs2[i].doc);
        assertEquals(sortedDocs[i].shardIndex, sortedDocs2[i].shardIndex);
        assertEquals(sortedDocs[i].score, sortedDocs2[i].score, 0.0f);
    }
}

Example 32 with QuerySearchResult

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

the class SearchPhaseControllerTests method consumerTestCase.

private void consumerTestCase(int numEmptyResponses) throws Exception {
    long beforeCompletedTasks = fixedExecutor.getCompletedTaskCount();
    int numShards = 3 + numEmptyResponses;
    int bufferSize = randomIntBetween(2, 3);
    CountDownLatch latch = new CountDownLatch(numShards);
    SearchRequest request = randomSearchRequest();
    request.source(new SearchSourceBuilder().aggregation(AggregationBuilders.avg("foo")));
    request.setBatchedReduceSize(bufferSize);
    ArraySearchPhaseResults<SearchPhaseResult> consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, request, 3 + numEmptyResponses, exc -> {
    });
    if (numEmptyResponses == 0) {
        assertEquals(0, reductions.size());
    }
    if (numEmptyResponses > 0) {
        QuerySearchResult empty = QuerySearchResult.nullInstance();
        int shardId = 2 + numEmptyResponses;
        empty.setShardIndex(2 + numEmptyResponses);
        empty.setSearchShardTarget(new SearchShardTarget("node", new ShardId("a", "b", shardId), null, OriginalIndices.NONE));
        consumer.consumeResult(empty, latch::countDown);
        numEmptyResponses--;
    }
    QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId("", 0), new SearchShardTarget("node", new ShardId("a", "b", 0), null, OriginalIndices.NONE), null);
    result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]), Float.NaN), new DocValueFormat[0]);
    InternalAggregations aggs = InternalAggregations.from(singletonList(new InternalMax("test", 1.0D, DocValueFormat.RAW, emptyMap())));
    result.aggregations(aggs);
    result.setShardIndex(0);
    consumer.consumeResult(result, latch::countDown);
    result = new QuerySearchResult(new ShardSearchContextId("", 1), new SearchShardTarget("node", new ShardId("a", "b", 0), null, OriginalIndices.NONE), null);
    result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]), Float.NaN), new DocValueFormat[0]);
    aggs = InternalAggregations.from(singletonList(new InternalMax("test", 3.0D, DocValueFormat.RAW, emptyMap())));
    result.aggregations(aggs);
    result.setShardIndex(2);
    consumer.consumeResult(result, latch::countDown);
    result = new QuerySearchResult(new ShardSearchContextId("", 1), new SearchShardTarget("node", new ShardId("a", "b", 0), null, OriginalIndices.NONE), null);
    result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]), Float.NaN), new DocValueFormat[0]);
    aggs = InternalAggregations.from(singletonList(new InternalMax("test", 2.0D, DocValueFormat.RAW, emptyMap())));
    result.aggregations(aggs);
    result.setShardIndex(1);
    consumer.consumeResult(result, latch::countDown);
    while (numEmptyResponses > 0) {
        result = QuerySearchResult.nullInstance();
        int shardId = 2 + numEmptyResponses;
        result.setShardIndex(shardId);
        result.setSearchShardTarget(new SearchShardTarget("node", new ShardId("a", "b", shardId), null, OriginalIndices.NONE));
        consumer.consumeResult(result, latch::countDown);
        numEmptyResponses--;
    }
    latch.await();
    final int numTotalReducePhases;
    if (numShards > bufferSize) {
        if (bufferSize == 2) {
            assertEquals(1, ((QueryPhaseResultConsumer) consumer).getNumReducePhases());
            assertEquals(1, reductions.size());
            assertEquals(false, reductions.get(0));
            numTotalReducePhases = 2;
        } else {
            assertEquals(0, ((QueryPhaseResultConsumer) consumer).getNumReducePhases());
            assertEquals(0, reductions.size());
            numTotalReducePhases = 1;
        }
    } else {
        assertEquals(0, reductions.size());
        numTotalReducePhases = 1;
    }
    SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
    assertEquals(numTotalReducePhases, reduce.numReducePhases);
    assertEquals(numTotalReducePhases, reductions.size());
    assertAggReduction(request);
    InternalMax max = (InternalMax) reduce.aggregations.asList().get(0);
    assertEquals(3.0D, max.getValue(), 0.0D);
    assertFalse(reduce.sortedTopDocs.isSortedByField);
    assertNull(reduce.sortedTopDocs.sortFields);
    assertNull(reduce.sortedTopDocs.collapseField);
    assertNull(reduce.sortedTopDocs.collapseValues);
}

Example 33 with QuerySearchResult

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

the class SearchPhaseControllerTests method testSortDocs.

public void testSortDocs() {
    List<CompletionSuggestion> suggestions = new ArrayList<>();
    for (int i = 0; i < randomIntBetween(1, 5); i++) {
        suggestions.add(new CompletionSuggestion(randomAlphaOfLength(randomIntBetween(1, 5)), randomIntBetween(1, 20), false));
    }
    int nShards = randomIntBetween(1, 20);
    int queryResultSize = randomBoolean() ? 0 : randomIntBetween(1, nShards * 2);
    AtomicArray<SearchPhaseResult> results = generateQueryResults(nShards, suggestions, queryResultSize, false);
    Optional<SearchPhaseResult> first = results.asList().stream().findFirst();
    int from = 0, size = 0;
    if (first.isPresent()) {
        from = first.get().queryResult().from();
        size = first.get().queryResult().size();
    }
    int accumulatedLength = Math.min(queryResultSize, getTotalQueryHits(results));
    List<CompletionSuggestion> reducedCompletionSuggestions = reducedSuggest(results);
    for (Suggest.Suggestion<?> suggestion : reducedCompletionSuggestions) {
        int suggestionSize = suggestion.getEntries().get(0).getOptions().size();
        accumulatedLength += suggestionSize;
    }
    List<TopDocs> topDocsList = new ArrayList<>();
    for (SearchPhaseResult result : results.asList()) {
        QuerySearchResult queryResult = result.queryResult();
        TopDocs topDocs = queryResult.consumeTopDocs().topDocs;
        SearchPhaseController.setShardIndex(topDocs, result.getShardIndex());
        topDocsList.add(topDocs);
    }
    ScoreDoc[] sortedDocs = SearchPhaseController.sortDocs(true, topDocsList, from, size, reducedCompletionSuggestions).scoreDocs;
    assertThat(sortedDocs.length, equalTo(accumulatedLength));
}

Example 34 with QuerySearchResult

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

the class SearchPhaseControllerTests method testConsumerFieldCollapsing.

public void testConsumerFieldCollapsing() throws Exception {
    int expectedNumResults = randomIntBetween(30, 100);
    int bufferSize = randomIntBetween(2, 200);
    SearchRequest request = randomSearchRequest();
    int size = randomIntBetween(5, 10);
    request.setBatchedReduceSize(bufferSize);
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, request, expectedNumResults, exc -> {
    });
    SortField[] sortFields = { new SortField("field", SortField.Type.STRING) };
    BytesRef a = new BytesRef("a");
    BytesRef b = new BytesRef("b");
    BytesRef c = new BytesRef("c");
    Object[] collapseValues = new Object[] { a, b, c };
    DocValueFormat[] docValueFormats = { DocValueFormat.RAW };
    CountDownLatch latch = new CountDownLatch(expectedNumResults);
    for (int i = 0; i < expectedNumResults; i++) {
        Object[] values = { randomFrom(collapseValues) };
        FieldDoc[] fieldDocs = { new FieldDoc(0, Float.NaN, values) };
        TopDocs topDocs = new CollapseTopFieldDocs("field", new TotalHits(1, Relation.EQUAL_TO), fieldDocs, sortFields, values);
        QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId("", i), new SearchShardTarget("node", new ShardId("a", "b", i), null, OriginalIndices.NONE), null);
        result.topDocs(new TopDocsAndMaxScore(topDocs, Float.NaN), docValueFormats);
        result.setShardIndex(i);
        result.size(size);
        consumer.consumeResult(result, latch::countDown);
    }
    latch.await();
    SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
    assertAggReduction(request);
    assertEquals(3, reduce.sortedTopDocs.scoreDocs.length);
    assertEquals(expectedNumResults, reduce.totalHits.value);
    assertEquals(a, ((FieldDoc) reduce.sortedTopDocs.scoreDocs[0]).fields[0]);
    assertEquals(b, ((FieldDoc) reduce.sortedTopDocs.scoreDocs[1]).fields[0]);
    assertEquals(c, ((FieldDoc) reduce.sortedTopDocs.scoreDocs[2]).fields[0]);
    assertTrue(reduce.sortedTopDocs.isSortedByField);
    assertEquals(1, reduce.sortedTopDocs.sortFields.length);
    assertEquals("field", reduce.sortedTopDocs.sortFields[0].getField());
    assertEquals(SortField.Type.STRING, reduce.sortedTopDocs.sortFields[0].getType());
    assertEquals("field", reduce.sortedTopDocs.collapseField);
    assertArrayEquals(collapseValues, reduce.sortedTopDocs.collapseValues);
}

Example 35 with QuerySearchResult

use of org.opensearch.search.query.QuerySearchResult 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));
}