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

Example 26 with QuerySearchResult

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

the class IndicesService method loadIntoContext.

/**
 * Loads the cache result, computing it if needed by executing the query phase and otherwise deserializing the cached
 * value into the {@link SearchContext#queryResult() context's query result}. The combination of load + compute allows
 * to have a single load operation that will cause other requests with the same key to wait till its loaded an reuse
 * the same cache.
 */
public void loadIntoContext(ShardSearchRequest request, SearchContext context, QueryPhase queryPhase) throws Exception {
    assert canCache(request, context);
    final DirectoryReader directoryReader = context.searcher().getDirectoryReader();
    boolean[] loadedFromCache = new boolean[] { true };
    BytesReference bytesReference = cacheShardLevelResult(context.indexShard(), directoryReader, request.cacheKey(), out -> {
        queryPhase.execute(context);
        context.queryResult().writeToNoId(out);
        loadedFromCache[0] = false;
    });
    if (loadedFromCache[0]) {
        // restore the cached query result into the context
        final QuerySearchResult result = context.queryResult();
        StreamInput in = new NamedWriteableAwareStreamInput(bytesReference.streamInput(), namedWriteableRegistry);
        result.readFromWithId(context.id(), in);
        result.setSearchShardTarget(context.shardTarget());
    } else if (context.queryResult().searchTimedOut()) {
        // we have to invalidate the cache entry if we cached a query result form a request that timed out.
        // we can't really throw exceptions in the loading part to signal a timed out search to the outside world since if there are
        // multiple requests that wait for the cache entry to be calculated they'd fail all with the same exception.
        // instead we all caching such a result for the time being, return the timed out result for all other searches with that cache
        // key invalidate the result in the thread that caused the timeout. This will end up to be simpler and eventually correct since
        // running a search that times out concurrently will likely timeout again if it's run while we have this `stale` result in the
        // cache. One other option is to not cache requests with a timeout at all...
        indicesRequestCache.invalidate(new IndexShardCacheEntity(context.indexShard()), directoryReader, request.cacheKey());
        if (logger.isTraceEnabled()) {
            logger.trace("Query timed out, invalidating cache entry for request on shard [{}]:\n {}", request.shardId(), request.source());
        }
    }
}

Example 27 with QuerySearchResult

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

the class SearchPhaseControllerTests method testConsumerSortByField.

public void testConsumerSortByField() throws Exception {
    int expectedNumResults = randomIntBetween(1, 100);
    int bufferSize = randomIntBetween(2, 200);
    SearchRequest request = randomSearchRequest();
    int size = randomIntBetween(1, 10);
    request.setBatchedReduceSize(bufferSize);
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, request, expectedNumResults, exc -> {
    });
    AtomicInteger max = new AtomicInteger();
    SortField[] sortFields = { new SortField("field", SortField.Type.INT, true) };
    DocValueFormat[] docValueFormats = { DocValueFormat.RAW };
    CountDownLatch latch = new CountDownLatch(expectedNumResults);
    for (int i = 0; i < expectedNumResults; i++) {
        int number = randomIntBetween(1, 1000);
        max.updateAndGet(prev -> Math.max(prev, number));
        FieldDoc[] fieldDocs = { new FieldDoc(0, Float.NaN, new Object[] { number }) };
        TopDocs topDocs = new TopFieldDocs(new TotalHits(1, Relation.EQUAL_TO), fieldDocs, sortFields);
        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(Math.min(expectedNumResults, size), reduce.sortedTopDocs.scoreDocs.length);
    assertEquals(expectedNumResults, reduce.totalHits.value);
    assertEquals(max.get(), ((FieldDoc) reduce.sortedTopDocs.scoreDocs[0]).fields[0]);
    assertTrue(reduce.sortedTopDocs.isSortedByField);
    assertEquals(1, reduce.sortedTopDocs.sortFields.length);
    assertEquals("field", reduce.sortedTopDocs.sortFields[0].getField());
    assertEquals(SortField.Type.INT, reduce.sortedTopDocs.sortFields[0].getType());
    assertNull(reduce.sortedTopDocs.collapseField);
    assertNull(reduce.sortedTopDocs.collapseValues);
}

Example 28 with QuerySearchResult

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

the class SearchPhaseControllerTests method testConsumerSuggestions.

public void testConsumerSuggestions() throws Exception {
    int expectedNumResults = randomIntBetween(1, 100);
    int bufferSize = randomIntBetween(2, 200);
    SearchRequest request = randomSearchRequest();
    request.setBatchedReduceSize(bufferSize);
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, request, expectedNumResults, exc -> {
    });
    int maxScoreTerm = -1;
    int maxScorePhrase = -1;
    int maxScoreCompletion = -1;
    CountDownLatch latch = new CountDownLatch(expectedNumResults);
    for (int i = 0; i < expectedNumResults; i++) {
        QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId("", i), new SearchShardTarget("node", new ShardId("a", "b", i), null, OriginalIndices.NONE), null);
        List<Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>>> suggestions = new ArrayList<>();
        {
            TermSuggestion termSuggestion = new TermSuggestion("term", 1, SortBy.SCORE);
            TermSuggestion.Entry entry = new TermSuggestion.Entry(new Text("entry"), 0, 10);
            int numOptions = randomIntBetween(1, 10);
            for (int j = 0; j < numOptions; j++) {
                int score = numOptions - j;
                maxScoreTerm = Math.max(maxScoreTerm, score);
                entry.addOption(new TermSuggestion.Entry.Option(new Text("option"), randomInt(), score));
            }
            termSuggestion.addTerm(entry);
            suggestions.add(termSuggestion);
        }
        {
            PhraseSuggestion phraseSuggestion = new PhraseSuggestion("phrase", 1);
            PhraseSuggestion.Entry entry = new PhraseSuggestion.Entry(new Text("entry"), 0, 10);
            int numOptions = randomIntBetween(1, 10);
            for (int j = 0; j < numOptions; j++) {
                int score = numOptions - j;
                maxScorePhrase = Math.max(maxScorePhrase, score);
                entry.addOption(new PhraseSuggestion.Entry.Option(new Text("option"), new Text("option"), score));
            }
            phraseSuggestion.addTerm(entry);
            suggestions.add(phraseSuggestion);
        }
        {
            CompletionSuggestion completionSuggestion = new CompletionSuggestion("completion", 1, false);
            CompletionSuggestion.Entry entry = new CompletionSuggestion.Entry(new Text("entry"), 0, 10);
            int numOptions = randomIntBetween(1, 10);
            for (int j = 0; j < numOptions; j++) {
                int score = numOptions - j;
                maxScoreCompletion = Math.max(maxScoreCompletion, score);
                CompletionSuggestion.Entry.Option option = new CompletionSuggestion.Entry.Option(j, new Text("option"), score, Collections.emptyMap());
                entry.addOption(option);
            }
            completionSuggestion.addTerm(entry);
            suggestions.add(completionSuggestion);
        }
        result.suggest(new Suggest(suggestions));
        result.topDocs(new TopDocsAndMaxScore(Lucene.EMPTY_TOP_DOCS, Float.NaN), new DocValueFormat[0]);
        result.setShardIndex(i);
        result.size(0);
        consumer.consumeResult(result, latch::countDown);
    }
    latch.await();
    SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
    assertEquals(3, reduce.suggest.size());
    {
        TermSuggestion term = reduce.suggest.getSuggestion("term");
        assertEquals(1, term.getEntries().size());
        assertEquals(1, term.getEntries().get(0).getOptions().size());
        assertEquals(maxScoreTerm, term.getEntries().get(0).getOptions().get(0).getScore(), 0f);
    }
    {
        PhraseSuggestion phrase = reduce.suggest.getSuggestion("phrase");
        assertEquals(1, phrase.getEntries().size());
        assertEquals(1, phrase.getEntries().get(0).getOptions().size());
        assertEquals(maxScorePhrase, phrase.getEntries().get(0).getOptions().get(0).getScore(), 0f);
    }
    {
        CompletionSuggestion completion = reduce.suggest.getSuggestion("completion");
        assertEquals(1, completion.getSize());
        assertEquals(1, completion.getOptions().size());
        CompletionSuggestion.Entry.Option option = completion.getOptions().get(0);
        assertEquals(maxScoreCompletion, option.getScore(), 0f);
    }
    assertAggReduction(request);
    assertEquals(1, reduce.sortedTopDocs.scoreDocs.length);
    assertEquals(maxScoreCompletion, reduce.sortedTopDocs.scoreDocs[0].score, 0f);
    assertEquals(0, reduce.sortedTopDocs.scoreDocs[0].doc);
    assertNotEquals(-1, reduce.sortedTopDocs.scoreDocs[0].shardIndex);
    assertEquals(0, reduce.totalHits.value);
    assertFalse(reduce.sortedTopDocs.isSortedByField);
    assertNull(reduce.sortedTopDocs.sortFields);
    assertNull(reduce.sortedTopDocs.collapseField);
    assertNull(reduce.sortedTopDocs.collapseValues);
}

Example 29 with QuerySearchResult

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

the class SearchPhaseControllerTests method testReduceCase.

private void testReduceCase(boolean shouldFail) throws Exception {
    int expectedNumResults = randomIntBetween(20, 200);
    int bufferSize = randomIntBetween(2, expectedNumResults - 1);
    SearchRequest request = new SearchRequest();
    request.source(new SearchSourceBuilder().aggregation(AggregationBuilders.avg("foo")).size(0));
    request.setBatchedReduceSize(bufferSize);
    AtomicBoolean hasConsumedFailure = new AtomicBoolean();
    AssertingCircuitBreaker circuitBreaker = new AssertingCircuitBreaker(CircuitBreaker.REQUEST);
    boolean shouldFailPartial = shouldFail && randomBoolean();
    if (shouldFailPartial) {
        circuitBreaker.shouldBreak.set(true);
    }
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, circuitBreaker, SearchProgressListener.NOOP, request, expectedNumResults, exc -> hasConsumedFailure.set(true));
    CountDownLatch latch = new CountDownLatch(expectedNumResults);
    Thread[] threads = new Thread[expectedNumResults];
    for (int i = 0; i < expectedNumResults; i++) {
        final int index = i;
        threads[index] = new Thread(() -> {
            QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), index), new SearchShardTarget("node", new ShardId("a", "b", index), null, OriginalIndices.NONE), null);
            result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), Lucene.EMPTY_SCORE_DOCS), Float.NaN), new DocValueFormat[0]);
            InternalAggregations aggs = InternalAggregations.from(Collections.singletonList(new InternalMax("test", 0d, DocValueFormat.RAW, Collections.emptyMap())));
            result.aggregations(aggs);
            result.setShardIndex(index);
            result.size(1);
            consumer.consumeResult(result, latch::countDown);
        });
        threads[index].start();
    }
    for (int i = 0; i < expectedNumResults; i++) {
        threads[i].join();
    }
    latch.await();
    if (shouldFail) {
        if (shouldFailPartial == false) {
            circuitBreaker.shouldBreak.set(true);
        }
        CircuitBreakingException exc = expectThrows(CircuitBreakingException.class, () -> consumer.reduce());
        assertEquals(shouldFailPartial, hasConsumedFailure.get());
        assertThat(exc.getMessage(), containsString("<reduce_aggs>"));
        circuitBreaker.shouldBreak.set(false);
    } else {
        SearchPhaseController.ReducedQueryPhase phase = consumer.reduce();
    }
    consumer.close();
    assertThat(circuitBreaker.allocated, equalTo(0L));
}

Example 30 with QuerySearchResult

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

the class SearchPhaseControllerTests method testConsumerOnlyHits.

public void testConsumerOnlyHits() throws Exception {
    int expectedNumResults = randomIntBetween(1, 100);
    int bufferSize = randomIntBetween(2, 200);
    SearchRequest request = randomSearchRequest();
    if (randomBoolean()) {
        request.source(new SearchSourceBuilder().size(randomIntBetween(1, 10)));
    }
    request.setBatchedReduceSize(bufferSize);
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, request, expectedNumResults, exc -> {
    });
    AtomicInteger max = new AtomicInteger();
    CountDownLatch latch = new CountDownLatch(expectedNumResults);
    for (int i = 0; i < expectedNumResults; i++) {
        int number = randomIntBetween(1, 1000);
        max.updateAndGet(prev -> Math.max(prev, number));
        QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId("", i), new SearchShardTarget("node", new ShardId("a", "b", i), null, OriginalIndices.NONE), null);
        result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), new ScoreDoc[] { new ScoreDoc(0, number) }), number), new DocValueFormat[0]);
        result.setShardIndex(i);
        result.size(1);
        consumer.consumeResult(result, latch::countDown);
    }
    latch.await();
    SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
    assertAggReduction(request);
    assertEquals(1, reduce.sortedTopDocs.scoreDocs.length);
    assertEquals(max.get(), reduce.maxScore, 0.0f);
    assertEquals(expectedNumResults, reduce.totalHits.value);
    assertEquals(max.get(), reduce.sortedTopDocs.scoreDocs[0].score, 0.0f);
    assertFalse(reduce.sortedTopDocs.isSortedByField);
    assertNull(reduce.sortedTopDocs.sortFields);
    assertNull(reduce.sortedTopDocs.collapseField);
    assertNull(reduce.sortedTopDocs.collapseValues);
}