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

Example 1 with SearchPhaseResult

use of org.opensearch.search.SearchPhaseResult 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 2 with SearchPhaseResult

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

the class SearchPhaseController method getHits.

private SearchHits getHits(ReducedQueryPhase reducedQueryPhase, boolean ignoreFrom, Collection<? extends SearchPhaseResult> fetchResults, IntFunction<SearchPhaseResult> resultsLookup) {
    SortedTopDocs sortedTopDocs = reducedQueryPhase.sortedTopDocs;
    int sortScoreIndex = -1;
    if (sortedTopDocs.isSortedByField) {
        SortField[] sortFields = sortedTopDocs.sortFields;
        for (int i = 0; i < sortFields.length; i++) {
            if (sortFields[i].getType() == SortField.Type.SCORE) {
                sortScoreIndex = i;
            }
        }
    }
    // clean the fetch counter
    for (SearchPhaseResult entry : fetchResults) {
        entry.fetchResult().initCounter();
    }
    int from = ignoreFrom ? 0 : reducedQueryPhase.from;
    int numSearchHits = (int) Math.min(reducedQueryPhase.fetchHits - from, reducedQueryPhase.size);
    // with collapsing we can have more fetch hits than sorted docs
    numSearchHits = Math.min(sortedTopDocs.scoreDocs.length, numSearchHits);
    // merge hits
    List<SearchHit> hits = new ArrayList<>();
    if (!fetchResults.isEmpty()) {
        for (int i = 0; i < numSearchHits; i++) {
            ScoreDoc shardDoc = sortedTopDocs.scoreDocs[i];
            SearchPhaseResult fetchResultProvider = resultsLookup.apply(shardDoc.shardIndex);
            if (fetchResultProvider == null) {
                // TODO it would be nice to assert this in the future
                continue;
            }
            FetchSearchResult fetchResult = fetchResultProvider.fetchResult();
            final int index = fetchResult.counterGetAndIncrement();
            assert index < fetchResult.hits().getHits().length : "not enough hits fetched. index [" + index + "] length: " + fetchResult.hits().getHits().length;
            SearchHit searchHit = fetchResult.hits().getHits()[index];
            searchHit.shard(fetchResult.getSearchShardTarget());
            if (sortedTopDocs.isSortedByField) {
                FieldDoc fieldDoc = (FieldDoc) shardDoc;
                searchHit.sortValues(fieldDoc.fields, reducedQueryPhase.sortValueFormats);
                if (sortScoreIndex != -1) {
                    searchHit.score(((Number) fieldDoc.fields[sortScoreIndex]).floatValue());
                }
            } else {
                searchHit.score(shardDoc.score);
            }
            hits.add(searchHit);
        }
    }
    return new SearchHits(hits.toArray(new SearchHit[0]), reducedQueryPhase.totalHits, reducedQueryPhase.maxScore, sortedTopDocs.sortFields, sortedTopDocs.collapseField, sortedTopDocs.collapseValues);
}

Example 3 with SearchPhaseResult

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

the class DfsQueryPhase method run.

@Override
public void run() throws IOException {
    // TODO we can potentially also consume the actual per shard results from the initial phase here in the aggregateDfs
    // to free up memory early
    final CountedCollector<SearchPhaseResult> counter = new CountedCollector<>(queryResult, searchResults.size(), () -> context.executeNextPhase(this, nextPhaseFactory.apply(queryResult)), context);
    for (final DfsSearchResult dfsResult : searchResults) {
        final SearchShardTarget searchShardTarget = dfsResult.getSearchShardTarget();
        Transport.Connection connection = context.getConnection(searchShardTarget.getClusterAlias(), searchShardTarget.getNodeId());
        QuerySearchRequest querySearchRequest = new QuerySearchRequest(searchShardTarget.getOriginalIndices(), dfsResult.getContextId(), dfsResult.getShardSearchRequest(), dfs);
        final int shardIndex = dfsResult.getShardIndex();
        searchTransportService.sendExecuteQuery(connection, querySearchRequest, context.getTask(), new SearchActionListener<QuerySearchResult>(searchShardTarget, shardIndex) {

            @Override
            protected void innerOnResponse(QuerySearchResult response) {
                try {
                    counter.onResult(response);
                } catch (Exception e) {
                    context.onPhaseFailure(DfsQueryPhase.this, "", e);
                }
            }

            @Override
            public void onFailure(Exception exception) {
                try {
                    context.getLogger().debug(() -> new ParameterizedMessage("[{}] Failed to execute query phase", querySearchRequest.contextId()), exception);
                    progressListener.notifyQueryFailure(shardIndex, searchShardTarget, exception);
                    counter.onFailure(shardIndex, searchShardTarget, exception);
                } finally {
                    if (context.getRequest().pointInTimeBuilder() == null) {
                        // the query might not have been executed at all (for example because thread pool rejected
                        // execution) and the search context that was created in dfs phase might not be released.
                        // release it again to be in the safe side
                        context.sendReleaseSearchContext(querySearchRequest.contextId(), connection, searchShardTarget.getOriginalIndices());
                    }
                }
            }
        });
    }
}

Example 4 with SearchPhaseResult

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

the class TransportSearchHelperTests method testParseScrollId.

public void testParseScrollId() {
    final Version version = VersionUtils.randomVersion(random());
    boolean includeUUID = version.onOrAfter(LegacyESVersion.V_7_7_0);
    final AtomicArray<SearchPhaseResult> queryResults = generateQueryResults();
    String scrollId = TransportSearchHelper.buildScrollId(queryResults, version);
    ParsedScrollId parseScrollId = TransportSearchHelper.parseScrollId(scrollId);
    assertEquals(3, parseScrollId.getContext().length);
    assertEquals("node_1", parseScrollId.getContext()[0].getNode());
    assertEquals("cluster_x", parseScrollId.getContext()[0].getClusterAlias());
    assertEquals(1, parseScrollId.getContext()[0].getSearchContextId().getId());
    if (includeUUID) {
        assertThat(parseScrollId.getContext()[0].getSearchContextId().getSessionId(), equalTo("a"));
    } else {
        assertThat(parseScrollId.getContext()[0].getSearchContextId().getSessionId(), equalTo(""));
    }
    assertEquals("node_2", parseScrollId.getContext()[1].getNode());
    assertEquals("cluster_y", parseScrollId.getContext()[1].getClusterAlias());
    assertEquals(12, parseScrollId.getContext()[1].getSearchContextId().getId());
    if (includeUUID) {
        assertThat(parseScrollId.getContext()[1].getSearchContextId().getSessionId(), equalTo("b"));
    } else {
        assertThat(parseScrollId.getContext()[1].getSearchContextId().getSessionId(), equalTo(""));
    }
    assertEquals("node_3", parseScrollId.getContext()[2].getNode());
    assertNull(parseScrollId.getContext()[2].getClusterAlias());
    assertEquals(42, parseScrollId.getContext()[2].getSearchContextId().getId());
    if (includeUUID) {
        assertThat(parseScrollId.getContext()[2].getSearchContextId().getSessionId(), equalTo("c"));
    } else {
        assertThat(parseScrollId.getContext()[2].getSearchContextId().getSessionId(), equalTo(""));
    }
}

Example 5 with SearchPhaseResult

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

the class ClearScrollControllerTests method testClearScrollIdsWithFailure.

public void testClearScrollIdsWithFailure() throws IOException, InterruptedException {
    DiscoveryNode node1 = new DiscoveryNode("node_1", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode node2 = new DiscoveryNode("node_2", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode node3 = new DiscoveryNode("node_3", buildNewFakeTransportAddress(), Version.CURRENT);
    AtomicArray<SearchPhaseResult> array = new AtomicArray<>(3);
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult1 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), 1), node1);
    testSearchPhaseResult1.setSearchShardTarget(new SearchShardTarget("node_1", new ShardId("idx", "uuid1", 2), null, null));
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult2 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), 12), node2);
    testSearchPhaseResult2.setSearchShardTarget(new SearchShardTarget("node_2", new ShardId("idy", "uuid2", 42), null, null));
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult3 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), 42), node3);
    testSearchPhaseResult3.setSearchShardTarget(new SearchShardTarget("node_3", new ShardId("idy", "uuid2", 43), null, null));
    array.setOnce(0, testSearchPhaseResult1);
    array.setOnce(1, testSearchPhaseResult2);
    array.setOnce(2, testSearchPhaseResult3);
    AtomicInteger numFreed = new AtomicInteger(0);
    AtomicInteger numFailures = new AtomicInteger(0);
    AtomicInteger numConnectionFailures = new AtomicInteger(0);
    String scrollId = TransportSearchHelper.buildScrollId(array, VersionUtils.randomVersion(random()));
    DiscoveryNodes nodes = DiscoveryNodes.builder().add(node1).add(node2).add(node3).build();
    CountDownLatch latch = new CountDownLatch(1);
    ActionListener<ClearScrollResponse> listener = new LatchedActionListener<>(new ActionListener<ClearScrollResponse>() {

        @Override
        public void onResponse(ClearScrollResponse clearScrollResponse) {
            assertEquals(numFreed.get(), clearScrollResponse.getNumFreed());
            if (numFailures.get() > 0) {
                assertFalse(clearScrollResponse.isSucceeded());
            } else {
                assertTrue(clearScrollResponse.isSucceeded());
            }
        }

        @Override
        public void onFailure(Exception e) {
            throw new AssertionError(e);
        }
    }, latch);
    List<DiscoveryNode> nodesInvoked = new CopyOnWriteArrayList<>();
    SearchTransportService searchTransportService = new SearchTransportService(null, null) {

        @Override
        public void sendFreeContext(Transport.Connection connection, ShardSearchContextId contextId, ActionListener<SearchFreeContextResponse> listener) {
            nodesInvoked.add(connection.getNode());
            boolean freed = randomBoolean();
            boolean fail = randomBoolean();
            Thread t = new Thread(() -> {
                if (fail) {
                    numFailures.incrementAndGet();
                    listener.onFailure(new IllegalArgumentException("boom"));
                } else {
                    if (freed) {
                        numFreed.incrementAndGet();
                    }
                    listener.onResponse(new SearchFreeContextResponse(freed));
                }
            });
            t.start();
        }

        @Override
        public Transport.Connection getConnection(String clusterAlias, DiscoveryNode node) {
            if (randomBoolean()) {
                numFailures.incrementAndGet();
                numConnectionFailures.incrementAndGet();
                throw new NodeNotConnectedException(node, "boom");
            }
            return new SearchAsyncActionTests.MockConnection(node);
        }
    };
    ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
    clearScrollRequest.scrollIds(Arrays.asList(scrollId));
    ClearScrollController controller = new ClearScrollController(clearScrollRequest, listener, nodes, logger, searchTransportService);
    controller.run();
    latch.await();
    assertEquals(3 - numConnectionFailures.get(), nodesInvoked.size());
}