Examples with SearchShardTarget org.opensearch.search.SearchShardTarget used on opensource projects

Example 1 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class MeanReciprocalRankTests method createSearchHits.

/**
 * Create SearchHits for testing, starting from dociId 'from' up to docId 'to'.
 * The search hits index also need to be provided
 */
private static SearchHit[] createSearchHits(int from, int to, String index) {
    SearchHit[] hits = new SearchHit[to + 1 - from];
    for (int i = from; i <= to; i++) {
        hits[i] = new SearchHit(i, i + "", Collections.emptyMap(), Collections.emptyMap());
        hits[i].shard(new SearchShardTarget("testnode", new ShardId(index, "uuid", 0), null, OriginalIndices.NONE));
    }
    return hits;
}

Example 2 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget 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 3 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class SearchContextId method encode.

public static String encode(List<SearchPhaseResult> searchPhaseResults, Map<String, AliasFilter> aliasFilter, Version version) {
    final Map<ShardId, SearchContextIdForNode> shards = new HashMap<>();
    for (SearchPhaseResult searchPhaseResult : searchPhaseResults) {
        final SearchShardTarget target = searchPhaseResult.getSearchShardTarget();
        shards.put(target.getShardId(), new SearchContextIdForNode(target.getClusterAlias(), target.getNodeId(), searchPhaseResult.getContextId()));
    }
    try (BytesStreamOutput out = new BytesStreamOutput()) {
        out.setVersion(version);
        Version.writeVersion(version, out);
        out.writeMap(shards, (o, k) -> k.writeTo(o), (o, v) -> v.writeTo(o));
        out.writeMap(aliasFilter, StreamOutput::writeString, (o, v) -> v.writeTo(o));
        return Base64.getUrlEncoder().encodeToString(BytesReference.toBytes(out.bytes()));
    } catch (IOException e) {
        throw new IllegalArgumentException(e);
    }
}

Example 4 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class SearchScrollAsyncAction method run.

private void run(BiFunction<String, String, DiscoveryNode> clusterNodeLookup, final SearchContextIdForNode[] context) {
    final CountDown counter = new CountDown(scrollId.getContext().length);
    for (int i = 0; i < context.length; i++) {
        SearchContextIdForNode target = context[i];
        final int shardIndex = i;
        final Transport.Connection connection;
        try {
            DiscoveryNode node = clusterNodeLookup.apply(target.getClusterAlias(), target.getNode());
            if (node == null) {
                throw new IllegalStateException("node [" + target.getNode() + "] is not available");
            }
            connection = getConnection(target.getClusterAlias(), node);
        } catch (Exception ex) {
            onShardFailure("query", counter, target.getSearchContextId(), ex, null, () -> SearchScrollAsyncAction.this.moveToNextPhase(clusterNodeLookup));
            continue;
        }
        final InternalScrollSearchRequest internalRequest = TransportSearchHelper.internalScrollSearchRequest(target.getSearchContextId(), request);
        // we can't create a SearchShardTarget here since we don't know the index and shard ID we are talking to
        // we only know the node and the search context ID. Yet, the response will contain the SearchShardTarget
        // from the target node instead...that's why we pass null here
        SearchActionListener<T> searchActionListener = new SearchActionListener<T>(null, shardIndex) {

            @Override
            protected void setSearchShardTarget(T response) {
                // don't do this - it's part of the response...
                assert response.getSearchShardTarget() != null : "search shard target must not be null";
                if (target.getClusterAlias() != null) {
                    // re-create the search target and add the cluster alias if there is any,
                    // we need this down the road for subseq. phases
                    SearchShardTarget searchShardTarget = response.getSearchShardTarget();
                    response.setSearchShardTarget(new SearchShardTarget(searchShardTarget.getNodeId(), searchShardTarget.getShardId(), target.getClusterAlias(), null));
                }
            }

            @Override
            protected void innerOnResponse(T result) {
                assert shardIndex == result.getShardIndex() : "shard index mismatch: " + shardIndex + " but got: " + result.getShardIndex();
                onFirstPhaseResult(shardIndex, result);
                if (counter.countDown()) {
                    SearchPhase phase = moveToNextPhase(clusterNodeLookup);
                    try {
                        phase.run();
                    } catch (Exception e) {
                        // we need to fail the entire request here - the entire phase just blew up
                        // don't call onShardFailure or onFailure here since otherwise we'd countDown the counter
                        // again which would result in an exception
                        listener.onFailure(new SearchPhaseExecutionException(phase.getName(), "Phase failed", e, ShardSearchFailure.EMPTY_ARRAY));
                    }
                }
            }

            @Override
            public void onFailure(Exception t) {
                onShardFailure("query", counter, target.getSearchContextId(), t, null, () -> SearchScrollAsyncAction.this.moveToNextPhase(clusterNodeLookup));
            }
        };
        executeInitialPhase(connection, internalRequest, searchActionListener);
    }
}

Example 5 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class PrecisionAtKTests method testIgnoreUnlabeled.

public void testIgnoreUnlabeled() {
    List<RatedDocument> rated = new ArrayList<>();
    rated.add(createRatedDoc("test", "0", RELEVANT_RATING));
    rated.add(createRatedDoc("test", "1", RELEVANT_RATING));
    // add an unlabeled search hit
    SearchHit[] searchHits = Arrays.copyOf(toSearchHits(rated, "test"), 3);
    searchHits[2] = new SearchHit(2, "2", Collections.emptyMap(), Collections.emptyMap());
    searchHits[2].shard(new SearchShardTarget("testnode", new ShardId("index", "uuid", 0), null, OriginalIndices.NONE));
    EvalQueryQuality evaluated = (new PrecisionAtK()).evaluate("id", searchHits, rated);
    assertEquals((double) 2 / 3, evaluated.metricScore(), 0.00001);
    assertEquals(2, ((PrecisionAtK.Detail) evaluated.getMetricDetails()).getRelevantRetrieved());
    assertEquals(3, ((PrecisionAtK.Detail) evaluated.getMetricDetails()).getRetrieved());
    // also try with setting `ignore_unlabeled`
    PrecisionAtK prec = new PrecisionAtK(true);
    evaluated = prec.evaluate("id", searchHits, rated);
    assertEquals((double) 2 / 2, evaluated.metricScore(), 0.00001);
    assertEquals(2, ((PrecisionAtK.Detail) evaluated.getMetricDetails()).getRelevantRetrieved());
    assertEquals(2, ((PrecisionAtK.Detail) evaluated.getMetricDetails()).getRetrieved());
}