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

Example 16 with QuerySearchResult

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

the class QueryPhaseResultConsumerTests method testProgressListenerExceptionsAreCaught.

public void testProgressListenerExceptionsAreCaught() throws Exception {
    ThrowingSearchProgressListener searchProgressListener = new ThrowingSearchProgressListener();
    List<SearchShard> searchShards = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        searchShards.add(new SearchShard(null, new ShardId("index", "uuid", i)));
    }
    searchProgressListener.notifyListShards(searchShards, Collections.emptyList(), SearchResponse.Clusters.EMPTY, false);
    SearchRequest searchRequest = new SearchRequest("index");
    searchRequest.setBatchedReduceSize(2);
    AtomicReference<Exception> onPartialMergeFailure = new AtomicReference<>();
    QueryPhaseResultConsumer queryPhaseResultConsumer = new QueryPhaseResultConsumer(searchRequest, executor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), searchPhaseController, searchProgressListener, writableRegistry(), 10, e -> onPartialMergeFailure.accumulateAndGet(e, (prev, curr) -> {
        curr.addSuppressed(prev);
        return curr;
    }));
    CountDownLatch partialReduceLatch = new CountDownLatch(10);
    for (int i = 0; i < 10; i++) {
        SearchShardTarget searchShardTarget = new SearchShardTarget("node", new ShardId("index", "uuid", i), null, OriginalIndices.NONE);
        QuerySearchResult querySearchResult = new QuerySearchResult();
        TopDocs topDocs = new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]);
        querySearchResult.topDocs(new TopDocsAndMaxScore(topDocs, Float.NaN), new DocValueFormat[0]);
        querySearchResult.setSearchShardTarget(searchShardTarget);
        querySearchResult.setShardIndex(i);
        queryPhaseResultConsumer.consumeResult(querySearchResult, partialReduceLatch::countDown);
    }
    assertEquals(10, searchProgressListener.onQueryResult.get());
    assertTrue(partialReduceLatch.await(10, TimeUnit.SECONDS));
    assertNull(onPartialMergeFailure.get());
    assertEquals(8, searchProgressListener.onPartialReduce.get());
    queryPhaseResultConsumer.reduce();
    assertEquals(1, searchProgressListener.onFinalReduce.get());
}

Example 17 with QuerySearchResult

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

the class DlsFlsValveImpl method onQueryPhase.

@Override
public void onQueryPhase(QuerySearchResult queryResult) {
    InternalAggregations aggregations = queryResult.aggregations().expand();
    assert aggregations != null;
    queryResult.aggregations(InternalAggregations.from(StreamSupport.stream(aggregations.spliterator(), false).map(aggregation -> aggregateBuckets((InternalAggregation) aggregation)).collect(ImmutableList.toImmutableList())));
}

Example 18 with QuerySearchResult

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

the class TermsReduceBenchmark method reduceAggs.

@Benchmark
public SearchPhaseController.ReducedQueryPhase reduceAggs(TermsList candidateList) throws Exception {
    List<QuerySearchResult> shards = new ArrayList<>();
    for (int i = 0; i < candidateList.size(); i++) {
        QuerySearchResult result = new QuerySearchResult();
        result.setShardIndex(i);
        result.from(0);
        result.size(0);
        result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1000, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO), new ScoreDoc[0]), Float.NaN), new DocValueFormat[] { DocValueFormat.RAW });
        result.aggregations(candidateList.get(i));
        result.setSearchShardTarget(new SearchShardTarget("node", new ShardId(new Index("index", "index"), i), null, OriginalIndices.NONE));
        shards.add(result);
    }
    SearchRequest request = new SearchRequest();
    request.source(new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.terms("test")));
    request.setBatchedReduceSize(bufferSize);
    ExecutorService executor = Executors.newFixedThreadPool(1);
    QueryPhaseResultConsumer consumer = new QueryPhaseResultConsumer(request, executor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), controller, SearchProgressListener.NOOP, namedWriteableRegistry, shards.size(), exc -> {
    });
    CountDownLatch latch = new CountDownLatch(shards.size());
    for (int i = 0; i < shards.size(); i++) {
        consumer.consumeResult(shards.get(i), () -> latch.countDown());
    }
    latch.await();
    SearchPhaseController.ReducedQueryPhase phase = consumer.reduce();
    executor.shutdownNow();
    return phase;
}

Example 19 with QuerySearchResult

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

the class SearchScrollQueryThenFetchAsyncAction method moveToNextPhase.

@Override
protected SearchPhase moveToNextPhase(BiFunction<String, String, DiscoveryNode> clusterNodeLookup) {
    return new SearchPhase("fetch") {

        @Override
        public void run() {
            final SearchPhaseController.ReducedQueryPhase reducedQueryPhase = searchPhaseController.reducedScrollQueryPhase(queryResults.asList());
            ScoreDoc[] scoreDocs = reducedQueryPhase.sortedTopDocs.scoreDocs;
            if (scoreDocs.length == 0) {
                sendResponse(reducedQueryPhase, fetchResults);
                return;
            }
            final IntArrayList[] docIdsToLoad = searchPhaseController.fillDocIdsToLoad(queryResults.length(), scoreDocs);
            final ScoreDoc[] lastEmittedDocPerShard = searchPhaseController.getLastEmittedDocPerShard(reducedQueryPhase, queryResults.length());
            final CountDown counter = new CountDown(docIdsToLoad.length);
            for (int i = 0; i < docIdsToLoad.length; i++) {
                final int index = i;
                final IntArrayList docIds = docIdsToLoad[index];
                if (docIds != null) {
                    final QuerySearchResult querySearchResult = queryResults.get(index);
                    ScoreDoc lastEmittedDoc = lastEmittedDocPerShard[index];
                    ShardFetchRequest shardFetchRequest = new ShardFetchRequest(querySearchResult.getContextId(), docIds, lastEmittedDoc);
                    SearchShardTarget searchShardTarget = querySearchResult.getSearchShardTarget();
                    DiscoveryNode node = clusterNodeLookup.apply(searchShardTarget.getClusterAlias(), searchShardTarget.getNodeId());
                    assert node != null : "target node is null in secondary phase";
                    Transport.Connection connection = getConnection(searchShardTarget.getClusterAlias(), node);
                    searchTransportService.sendExecuteFetchScroll(connection, shardFetchRequest, task, new SearchActionListener<FetchSearchResult>(querySearchResult.getSearchShardTarget(), index) {

                        @Override
                        protected void innerOnResponse(FetchSearchResult response) {
                            fetchResults.setOnce(response.getShardIndex(), response);
                            if (counter.countDown()) {
                                sendResponse(reducedQueryPhase, fetchResults);
                            }
                        }

                        @Override
                        public void onFailure(Exception t) {
                            onShardFailure(getName(), counter, querySearchResult.getContextId(), t, querySearchResult.getSearchShardTarget(), () -> sendResponsePhase(reducedQueryPhase, fetchResults));
                        }
                    });
                } else {
                    // which can have null values so we have to count them down too
                    if (counter.countDown()) {
                        sendResponse(reducedQueryPhase, fetchResults);
                    }
                }
            }
        }
    };
}

Example 20 with QuerySearchResult

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

the class SearchPhaseController method reducedScrollQueryPhase.

/**
 * Reduces the given query results and consumes all aggregations and profile results.
 * @param queryResults a list of non-null query shard results
 */
ReducedQueryPhase reducedScrollQueryPhase(Collection<? extends SearchPhaseResult> queryResults) {
    InternalAggregation.ReduceContextBuilder aggReduceContextBuilder = new InternalAggregation.ReduceContextBuilder() {

        @Override
        public ReduceContext forPartialReduction() {
            throw new UnsupportedOperationException("Scroll requests don't have aggs");
        }

        @Override
        public ReduceContext forFinalReduction() {
            throw new UnsupportedOperationException("Scroll requests don't have aggs");
        }
    };
    final TopDocsStats topDocsStats = new TopDocsStats(SearchContext.TRACK_TOTAL_HITS_ACCURATE);
    final List<TopDocs> topDocs = new ArrayList<>();
    for (SearchPhaseResult sortedResult : queryResults) {
        QuerySearchResult queryResult = sortedResult.queryResult();
        final TopDocsAndMaxScore td = queryResult.consumeTopDocs();
        assert td != null;
        topDocsStats.add(td, queryResult.searchTimedOut(), queryResult.terminatedEarly());
        // make sure we set the shard index before we add it - the consumer didn't do that yet
        if (td.topDocs.scoreDocs.length > 0) {
            setShardIndex(td.topDocs, queryResult.getShardIndex());
            topDocs.add(td.topDocs);
        }
    }
    return reducedQueryPhase(queryResults, Collections.emptyList(), topDocs, topDocsStats, 0, true, aggReduceContextBuilder, true);
}