Examples with InternalAggregations org.opensearch.search.aggregations.InternalAggregations used on opensource projects

Example 1 with InternalAggregations

use of org.opensearch.search.aggregations.InternalAggregations in project OpenSearch by opensearch-project.

the class SearchResponseMerger method getMergedResponse.

/**
 * Returns the merged response. To be called once all responses have been added through {@link #add(SearchResponse)}
 * so that all responses are merged into a single one.
 */
SearchResponse getMergedResponse(SearchResponse.Clusters clusters) {
    // we end up calling merge without anything to merge, we just return an empty search response
    if (searchResponses.size() == 0) {
        return SearchResponse.empty(searchTimeProvider::buildTookInMillis, clusters);
    }
    int totalShards = 0;
    int skippedShards = 0;
    int successfulShards = 0;
    // the current reduce phase counts as one
    int numReducePhases = 1;
    List<ShardSearchFailure> failures = new ArrayList<>();
    Map<String, ProfileShardResult> profileResults = new HashMap<>();
    List<InternalAggregations> aggs = new ArrayList<>();
    Map<ShardIdAndClusterAlias, Integer> shards = new TreeMap<>();
    List<TopDocs> topDocsList = new ArrayList<>(searchResponses.size());
    Map<String, List<Suggest.Suggestion>> groupedSuggestions = new HashMap<>();
    Boolean trackTotalHits = null;
    SearchPhaseController.TopDocsStats topDocsStats = new SearchPhaseController.TopDocsStats(trackTotalHitsUpTo);
    for (SearchResponse searchResponse : searchResponses) {
        totalShards += searchResponse.getTotalShards();
        skippedShards += searchResponse.getSkippedShards();
        successfulShards += searchResponse.getSuccessfulShards();
        numReducePhases += searchResponse.getNumReducePhases();
        Collections.addAll(failures, searchResponse.getShardFailures());
        profileResults.putAll(searchResponse.getProfileResults());
        if (searchResponse.getAggregations() != null) {
            InternalAggregations internalAggs = (InternalAggregations) searchResponse.getAggregations();
            aggs.add(internalAggs);
        }
        Suggest suggest = searchResponse.getSuggest();
        if (suggest != null) {
            for (Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> entries : suggest) {
                List<Suggest.Suggestion> suggestionList = groupedSuggestions.computeIfAbsent(entries.getName(), s -> new ArrayList<>());
                suggestionList.add(entries);
            }
            List<CompletionSuggestion> completionSuggestions = suggest.filter(CompletionSuggestion.class);
            for (CompletionSuggestion completionSuggestion : completionSuggestions) {
                for (CompletionSuggestion.Entry options : completionSuggestion) {
                    for (CompletionSuggestion.Entry.Option option : options) {
                        SearchShardTarget shard = option.getHit().getShard();
                        ShardIdAndClusterAlias shardId = new ShardIdAndClusterAlias(shard.getShardId(), shard.getClusterAlias());
                        shards.putIfAbsent(shardId, null);
                    }
                }
            }
        }
        SearchHits searchHits = searchResponse.getHits();
        final TotalHits totalHits;
        if (searchHits.getTotalHits() == null) {
            // in case we didn't track total hits, we get null from each cluster, but we need to set 0 eq to the TopDocs
            totalHits = new TotalHits(0, TotalHits.Relation.EQUAL_TO);
            assert trackTotalHits == null || trackTotalHits == false;
            trackTotalHits = false;
        } else {
            totalHits = searchHits.getTotalHits();
            assert trackTotalHits == null || trackTotalHits;
            trackTotalHits = true;
        }
        TopDocs topDocs = searchHitsToTopDocs(searchHits, totalHits, shards);
        topDocsStats.add(new TopDocsAndMaxScore(topDocs, searchHits.getMaxScore()), searchResponse.isTimedOut(), searchResponse.isTerminatedEarly());
        if (searchHits.getHits().length > 0) {
            // there is no point in adding empty search hits and merging them with the others. Also, empty search hits always come
            // without sort fields and collapse info, despite sort by field and/or field collapsing was requested, which causes
            // issues reconstructing the proper TopDocs instance and breaks mergeTopDocs which expects the same type for each result.
            topDocsList.add(topDocs);
        }
    }
    // after going through all the hits and collecting all their distinct shards, we assign shardIndex and set it to the ScoreDocs
    setTopDocsShardIndex(shards, topDocsList);
    TopDocs topDocs = SearchPhaseController.mergeTopDocs(topDocsList, size, from);
    SearchHits mergedSearchHits = topDocsToSearchHits(topDocs, topDocsStats);
    setSuggestShardIndex(shards, groupedSuggestions);
    Suggest suggest = groupedSuggestions.isEmpty() ? null : new Suggest(Suggest.reduce(groupedSuggestions));
    InternalAggregations reducedAggs = InternalAggregations.topLevelReduce(aggs, aggReduceContextBuilder.forFinalReduction());
    ShardSearchFailure[] shardFailures = failures.toArray(ShardSearchFailure.EMPTY_ARRAY);
    SearchProfileShardResults profileShardResults = profileResults.isEmpty() ? null : new SearchProfileShardResults(profileResults);
    // make failures ordering consistent between ordinary search and CCS by looking at the shard they come from
    Arrays.sort(shardFailures, FAILURES_COMPARATOR);
    InternalSearchResponse response = new InternalSearchResponse(mergedSearchHits, reducedAggs, suggest, profileShardResults, topDocsStats.timedOut, topDocsStats.terminatedEarly, numReducePhases);
    long tookInMillis = searchTimeProvider.buildTookInMillis();
    return new SearchResponse(response, null, totalShards, successfulShards, skippedShards, tookInMillis, shardFailures, clusters, null);
}

Example 2 with InternalAggregations

use of org.opensearch.search.aggregations.InternalAggregations in project OpenSearch by opensearch-project.

the class SearchProgressActionListenerIT method testCase.

private void testCase(NodeClient client, SearchRequest request, List<SearchShard> expectedShards, boolean hasFetchPhase) throws InterruptedException {
    AtomicInteger numQueryResults = new AtomicInteger();
    AtomicInteger numQueryFailures = new AtomicInteger();
    AtomicInteger numFetchResults = new AtomicInteger();
    AtomicInteger numFetchFailures = new AtomicInteger();
    AtomicInteger numReduces = new AtomicInteger();
    AtomicReference<SearchResponse> searchResponse = new AtomicReference<>();
    AtomicReference<List<SearchShard>> shardsListener = new AtomicReference<>();
    CountDownLatch latch = new CountDownLatch(1);
    SearchProgressActionListener listener = new SearchProgressActionListener() {

        @Override
        public void onListShards(List<SearchShard> shards, List<SearchShard> skippedShards, SearchResponse.Clusters clusters, boolean fetchPhase) {
            shardsListener.set(shards);
            assertEquals(fetchPhase, hasFetchPhase);
        }

        @Override
        public void onQueryResult(int shardIndex) {
            assertThat(shardIndex, lessThan(shardsListener.get().size()));
            numQueryResults.incrementAndGet();
        }

        @Override
        public void onQueryFailure(int shardIndex, SearchShardTarget shardTarget, Exception exc) {
            assertThat(shardIndex, lessThan(shardsListener.get().size()));
            numQueryFailures.incrementAndGet();
        }

        @Override
        public void onFetchResult(int shardIndex) {
            assertThat(shardIndex, lessThan(shardsListener.get().size()));
            numFetchResults.incrementAndGet();
        }

        @Override
        public void onFetchFailure(int shardIndex, SearchShardTarget shardTarget, Exception exc) {
            assertThat(shardIndex, lessThan(shardsListener.get().size()));
            numFetchFailures.incrementAndGet();
        }

        @Override
        public void onPartialReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {
            numReduces.incrementAndGet();
        }

        @Override
        public void onFinalReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {
            numReduces.incrementAndGet();
        }

        @Override
        public void onResponse(SearchResponse response) {
            searchResponse.set(response);
            latch.countDown();
        }

        @Override
        public void onFailure(Exception e) {
            throw new AssertionError();
        }
    };
    client.executeLocally(SearchAction.INSTANCE, new SearchRequest(request) {

        @Override
        public SearchTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
            SearchTask task = super.createTask(id, type, action, parentTaskId, headers);
            task.setProgressListener(listener);
            return task;
        }
    }, listener);
    latch.await();
    assertThat(shardsListener.get(), equalTo(expectedShards));
    assertThat(numQueryResults.get(), equalTo(searchResponse.get().getSuccessfulShards()));
    assertThat(numQueryFailures.get(), equalTo(searchResponse.get().getFailedShards()));
    if (hasFetchPhase) {
        assertThat(numFetchResults.get(), equalTo(searchResponse.get().getSuccessfulShards()));
        assertThat(numFetchFailures.get(), equalTo(0));
    } else {
        assertThat(numFetchResults.get(), equalTo(0));
        assertThat(numFetchFailures.get(), equalTo(0));
    }
    assertThat(numReduces.get(), equalTo(searchResponse.get().getNumReducePhases()));
}

Example 3 with InternalAggregations

use of org.opensearch.search.aggregations.InternalAggregations in project OpenSearch by opensearch-project.

the class QueryPhaseResultConsumer method partialReduce.

private MergeResult partialReduce(QuerySearchResult[] toConsume, List<SearchShard> emptyResults, SearchPhaseController.TopDocsStats topDocsStats, MergeResult lastMerge, int numReducePhases) {
    // ensure consistent ordering
    Arrays.sort(toConsume, Comparator.comparingInt(QuerySearchResult::getShardIndex));
    for (QuerySearchResult result : toConsume) {
        topDocsStats.add(result.topDocs(), result.searchTimedOut(), result.terminatedEarly());
    }
    final TopDocs newTopDocs;
    if (hasTopDocs) {
        List<TopDocs> topDocsList = new ArrayList<>();
        if (lastMerge != null) {
            topDocsList.add(lastMerge.reducedTopDocs);
        }
        for (QuerySearchResult result : toConsume) {
            TopDocsAndMaxScore topDocs = result.consumeTopDocs();
            SearchPhaseController.setShardIndex(topDocs.topDocs, result.getShardIndex());
            topDocsList.add(topDocs.topDocs);
        }
        newTopDocs = SearchPhaseController.mergeTopDocs(topDocsList, // we have to merge here in the same way we collect on a shard
        topNSize, 0);
    } else {
        newTopDocs = null;
    }
    final InternalAggregations newAggs;
    if (hasAggs) {
        List<InternalAggregations> aggsList = new ArrayList<>();
        if (lastMerge != null) {
            aggsList.add(lastMerge.reducedAggs);
        }
        for (QuerySearchResult result : toConsume) {
            aggsList.add(result.consumeAggs().expand());
        }
        newAggs = InternalAggregations.topLevelReduce(aggsList, aggReduceContextBuilder.forPartialReduction());
    } else {
        newAggs = null;
    }
    List<SearchShard> processedShards = new ArrayList<>(emptyResults);
    if (lastMerge != null) {
        processedShards.addAll(lastMerge.processedShards);
    }
    for (QuerySearchResult result : toConsume) {
        SearchShardTarget target = result.getSearchShardTarget();
        processedShards.add(new SearchShard(target.getClusterAlias(), target.getShardId()));
    }
    progressListener.notifyPartialReduce(processedShards, topDocsStats.getTotalHits(), newAggs, numReducePhases);
    // we leave the results un-serialized because serializing is slow but we compute the serialized
    // size as an estimate of the memory used by the newly reduced aggregations.
    long serializedSize = hasAggs ? newAggs.getSerializedSize() : 0;
    return new MergeResult(processedShards, newTopDocs, newAggs, hasAggs ? serializedSize : 0);
}

Example 4 with InternalAggregations

use of org.opensearch.search.aggregations.InternalAggregations in project OpenSearch by opensearch-project.

the class InternalSingleBucketAggregation method reduce.

@Override
public InternalAggregation reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
    long docCount = 0L;
    List<InternalAggregations> subAggregationsList = new ArrayList<>(aggregations.size());
    for (InternalAggregation aggregation : aggregations) {
        assert aggregation.getName().equals(getName());
        docCount += ((InternalSingleBucketAggregation) aggregation).docCount;
        subAggregationsList.add(((InternalSingleBucketAggregation) aggregation).aggregations);
    }
    final InternalAggregations aggs = InternalAggregations.reduce(subAggregationsList, reduceContext);
    return newAggregation(getName(), docCount, aggs);
}

Example 5 with InternalAggregations

use of org.opensearch.search.aggregations.InternalAggregations in project OpenSearch by opensearch-project.

the class InternalSingleBucketAggregation method reducePipelines.

/**
 * Amulti-bucket agg needs to first reduce the buckets and *their* pipelines
 * before allowing sibling pipelines to materialize.
 */
@Override
public final InternalAggregation reducePipelines(InternalAggregation reducedAggs, ReduceContext reduceContext, PipelineTree pipelineTree) {
    assert reduceContext.isFinalReduce();
    InternalAggregation reduced = this;
    if (pipelineTree.hasSubTrees()) {
        List<InternalAggregation> aggs = new ArrayList<>();
        for (Aggregation agg : getAggregations().asList()) {
            PipelineTree subTree = pipelineTree.subTree(agg.getName());
            aggs.add(((InternalAggregation) agg).reducePipelines((InternalAggregation) agg, reduceContext, subTree));
        }
        InternalAggregations reducedSubAggs = InternalAggregations.from(aggs);
        reduced = create(reducedSubAggs);
    }
    return super.reducePipelines(reduced, reduceContext, pipelineTree);
}