Examples with RescoreContext org.opensearch.search.rescore.RescoreContext used on opensource projects

Example 6 with RescoreContext

use of org.opensearch.search.rescore.RescoreContext in project OpenSearch by opensearch-project.

the class TransportExplainAction method shardOperation.

@Override
protected ExplainResponse shardOperation(ExplainRequest request, ShardId shardId) throws IOException {
    ShardSearchRequest shardSearchLocalRequest = new ShardSearchRequest(shardId, request.nowInMillis, request.filteringAlias());
    SearchContext context = searchService.createSearchContext(shardSearchLocalRequest, SearchService.NO_TIMEOUT);
    Engine.GetResult result = null;
    try {
        // No need to check the type, IndexShard#get does it for us
        Term uidTerm = new Term(IdFieldMapper.NAME, Uid.encodeId(request.id()));
        result = context.indexShard().get(new Engine.Get(false, false, request.id(), uidTerm));
        if (!result.exists()) {
            return new ExplainResponse(shardId.getIndexName(), request.id(), false);
        }
        context.parsedQuery(context.getQueryShardContext().toQuery(request.query()));
        context.preProcess(true);
        int topLevelDocId = result.docIdAndVersion().docId + result.docIdAndVersion().docBase;
        Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
        for (RescoreContext ctx : context.rescore()) {
            Rescorer rescorer = ctx.rescorer();
            explanation = rescorer.explain(topLevelDocId, context.searcher(), ctx, explanation);
        }
        if (request.storedFields() != null || (request.fetchSourceContext() != null && request.fetchSourceContext().fetchSource())) {
            // Advantage is that we're not opening a second searcher to retrieve the _source. Also
            // because we are working in the same searcher in engineGetResult we can be sure that a
            // doc isn't deleted between the initial get and this call.
            GetResult getResult = context.indexShard().getService().get(result, request.id(), request.storedFields(), request.fetchSourceContext());
            return new ExplainResponse(shardId.getIndexName(), request.id(), true, explanation, getResult);
        } else {
            return new ExplainResponse(shardId.getIndexName(), request.id(), true, explanation);
        }
    } catch (IOException e) {
        throw new OpenSearchException("Could not explain", e);
    } finally {
        Releasables.close(result, context);
    }
}

Example 7 with RescoreContext

use of org.opensearch.search.rescore.RescoreContext in project OpenSearch by opensearch-project.

the class DfsPhase method execute.

public void execute(SearchContext context) {
    try {
        ObjectObjectHashMap<String, CollectionStatistics> fieldStatistics = HppcMaps.newNoNullKeysMap();
        Map<Term, TermStatistics> stats = new HashMap<>();
        IndexSearcher searcher = new IndexSearcher(context.searcher().getIndexReader()) {

            @Override
            public TermStatistics termStatistics(Term term, int docFreq, long totalTermFreq) throws IOException {
                if (context.isCancelled()) {
                    throw new TaskCancelledException("cancelled task with reason: " + context.getTask().getReasonCancelled());
                }
                TermStatistics ts = super.termStatistics(term, docFreq, totalTermFreq);
                if (ts != null) {
                    stats.put(term, ts);
                }
                return ts;
            }

            @Override
            public CollectionStatistics collectionStatistics(String field) throws IOException {
                if (context.isCancelled()) {
                    throw new TaskCancelledException("cancelled task with reason: " + context.getTask().getReasonCancelled());
                }
                CollectionStatistics cs = super.collectionStatistics(field);
                if (cs != null) {
                    fieldStatistics.put(field, cs);
                }
                return cs;
            }
        };
        searcher.createWeight(context.searcher().rewrite(context.query()), ScoreMode.COMPLETE, 1);
        for (RescoreContext rescoreContext : context.rescore()) {
            for (Query query : rescoreContext.getQueries()) {
                searcher.createWeight(context.searcher().rewrite(query), ScoreMode.COMPLETE, 1);
            }
        }
        Term[] terms = stats.keySet().toArray(new Term[0]);
        TermStatistics[] termStatistics = new TermStatistics[terms.length];
        for (int i = 0; i < terms.length; i++) {
            termStatistics[i] = stats.get(terms[i]);
        }
        context.dfsResult().termsStatistics(terms, termStatistics).fieldStatistics(fieldStatistics).maxDoc(context.searcher().getIndexReader().maxDoc());
    } catch (Exception e) {
        throw new DfsPhaseExecutionException(context.shardTarget(), "Exception during dfs phase", e);
    }
}

Example 8 with RescoreContext

use of org.opensearch.search.rescore.RescoreContext in project OpenSearch by opensearch-project.

the class TopHitsAggregator method buildAggregation.

@Override
public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOException {
    Collectors collectors = topDocsCollectors.get(owningBucketOrdinal);
    if (collectors == null) {
        return buildEmptyAggregation();
    }
    TopDocsCollector<?> topDocsCollector = collectors.topDocsCollector;
    TopDocs topDocs = topDocsCollector.topDocs();
    float maxScore = Float.NaN;
    if (subSearchContext.sort() == null) {
        for (RescoreContext ctx : context().rescore()) {
            try {
                topDocs = ctx.rescorer().rescore(topDocs, context.searcher(), ctx);
            } catch (IOException e) {
                throw new OpenSearchException("Rescore TopHits Failed", e);
            }
        }
        if (topDocs.scoreDocs.length > 0) {
            maxScore = topDocs.scoreDocs[0].score;
        }
    } else if (subSearchContext.trackScores()) {
        TopFieldCollector.populateScores(topDocs.scoreDocs, subSearchContext.searcher(), subSearchContext.query());
        maxScore = collectors.maxScoreCollector.getMaxScore();
    }
    final TopDocsAndMaxScore topDocsAndMaxScore = new TopDocsAndMaxScore(topDocs, maxScore);
    subSearchContext.queryResult().topDocs(topDocsAndMaxScore, subSearchContext.sort() == null ? null : subSearchContext.sort().formats);
    int[] docIdsToLoad = new int[topDocs.scoreDocs.length];
    for (int i = 0; i < topDocs.scoreDocs.length; i++) {
        docIdsToLoad[i] = topDocs.scoreDocs[i].doc;
    }
    subSearchContext.docIdsToLoad(docIdsToLoad, 0, docIdsToLoad.length);
    fetchPhase.execute(subSearchContext);
    FetchSearchResult fetchResult = subSearchContext.fetchResult();
    SearchHit[] internalHits = fetchResult.fetchResult().hits().getHits();
    for (int i = 0; i < internalHits.length; i++) {
        ScoreDoc scoreDoc = topDocs.scoreDocs[i];
        SearchHit searchHitFields = internalHits[i];
        searchHitFields.shard(subSearchContext.shardTarget());
        searchHitFields.score(scoreDoc.score);
        if (scoreDoc instanceof FieldDoc) {
            FieldDoc fieldDoc = (FieldDoc) scoreDoc;
            searchHitFields.sortValues(fieldDoc.fields, subSearchContext.sort().formats);
        }
    }
    return new InternalTopHits(name, subSearchContext.from(), subSearchContext.size(), topDocsAndMaxScore, fetchResult.hits(), metadata());
}

Example 9 with RescoreContext

use of org.opensearch.search.rescore.RescoreContext in project OpenSearch by opensearch-project.

the class TopHitsAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
    // Create leaf collectors here instead of at the aggregator level. Otherwise in case this collector get invoked
    // when post collecting then we have already replaced the leaf readers on the aggregator level have already been
    // replaced with the next leaf readers and then post collection pushes docids of the previous segment, which
    // then causes assertions to trip or incorrect top docs to be computed.
    final LongObjectHashMap<LeafCollector> leafCollectors = new LongObjectHashMap<>(1);
    return new LeafBucketCollectorBase(sub, null) {

        Scorable scorer;

        @Override
        public void setScorer(Scorable scorer) throws IOException {
            this.scorer = scorer;
            super.setScorer(scorer);
            for (ObjectCursor<LeafCollector> cursor : leafCollectors.values()) {
                cursor.value.setScorer(scorer);
            }
        }

        @Override
        public void collect(int docId, long bucket) throws IOException {
            Collectors collectors = topDocsCollectors.get(bucket);
            if (collectors == null) {
                SortAndFormats sort = subSearchContext.sort();
                int topN = subSearchContext.from() + subSearchContext.size();
                if (sort == null) {
                    for (RescoreContext rescoreContext : context.rescore()) {
                        topN = Math.max(rescoreContext.getWindowSize(), topN);
                    }
                }
                // In the QueryPhase we don't need this protection, because it is build into the IndexSearcher,
                // but here we create collectors ourselves and we need prevent OOM because of crazy an offset and size.
                topN = Math.min(topN, subSearchContext.searcher().getIndexReader().maxDoc());
                if (sort == null) {
                    collectors = new Collectors(TopScoreDocCollector.create(topN, Integer.MAX_VALUE), null);
                } else {
                    // TODO: can we pass trackTotalHits=subSearchContext.trackTotalHits(){
                    // Note that this would require to catch CollectionTerminatedException
                    collectors = new Collectors(TopFieldCollector.create(sort.sort, topN, Integer.MAX_VALUE), subSearchContext.trackScores() ? new MaxScoreCollector() : null);
                }
                topDocsCollectors.put(bucket, collectors);
            }
            final LeafCollector leafCollector;
            final int key = leafCollectors.indexOf(bucket);
            if (key < 0) {
                leafCollector = collectors.collector.getLeafCollector(ctx);
                if (scorer != null) {
                    leafCollector.setScorer(scorer);
                }
                leafCollectors.indexInsert(key, bucket, leafCollector);
            } else {
                leafCollector = leafCollectors.indexGet(key);
            }
            leafCollector.collect(docId);
        }
    };
}

Example 10 with RescoreContext

use of org.opensearch.search.rescore.RescoreContext in project OpenSearch by opensearch-project.

the class DefaultSearchContextTests method testPreProcess.

public void testPreProcess() throws Exception {
    TimeValue timeout = new TimeValue(randomIntBetween(1, 100));
    ShardSearchRequest shardSearchRequest = mock(ShardSearchRequest.class);
    when(shardSearchRequest.searchType()).thenReturn(SearchType.DEFAULT);
    ShardId shardId = new ShardId("index", UUID.randomUUID().toString(), 1);
    when(shardSearchRequest.shardId()).thenReturn(shardId);
    ThreadPool threadPool = new TestThreadPool(this.getClass().getName());
    IndexShard indexShard = mock(IndexShard.class);
    QueryCachingPolicy queryCachingPolicy = mock(QueryCachingPolicy.class);
    when(indexShard.getQueryCachingPolicy()).thenReturn(queryCachingPolicy);
    when(indexShard.getThreadPool()).thenReturn(threadPool);
    int maxResultWindow = randomIntBetween(50, 100);
    int maxRescoreWindow = randomIntBetween(50, 100);
    int maxSlicesPerScroll = randomIntBetween(50, 100);
    Settings settings = Settings.builder().put("index.max_result_window", maxResultWindow).put("index.max_slices_per_scroll", maxSlicesPerScroll).put("index.max_rescore_window", maxRescoreWindow).put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1).put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 2).build();
    IndexService indexService = mock(IndexService.class);
    IndexCache indexCache = mock(IndexCache.class);
    QueryCache queryCache = mock(QueryCache.class);
    when(indexCache.query()).thenReturn(queryCache);
    when(indexService.cache()).thenReturn(indexCache);
    QueryShardContext queryShardContext = mock(QueryShardContext.class);
    when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class))).thenReturn(queryShardContext);
    MapperService mapperService = mock(MapperService.class);
    when(mapperService.hasNested()).thenReturn(randomBoolean());
    when(indexService.mapperService()).thenReturn(mapperService);
    IndexMetadata indexMetadata = IndexMetadata.builder("index").settings(settings).build();
    IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY);
    when(indexService.getIndexSettings()).thenReturn(indexSettings);
    when(mapperService.getIndexSettings()).thenReturn(indexSettings);
    BigArrays bigArrays = new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), new NoneCircuitBreakerService());
    try (Directory dir = newDirectory();
        RandomIndexWriter w = new RandomIndexWriter(random(), dir)) {
        final Supplier<Engine.SearcherSupplier> searcherSupplier = () -> new Engine.SearcherSupplier(Function.identity()) {

            @Override
            protected void doClose() {
            }

            @Override
            protected Engine.Searcher acquireSearcherInternal(String source) {
                try {
                    IndexReader reader = w.getReader();
                    return new Engine.Searcher("test", reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), reader);
                } catch (IOException exc) {
                    throw new AssertionError(exc);
                }
            }
        };
        SearchShardTarget target = new SearchShardTarget("node", shardId, null, OriginalIndices.NONE);
        ReaderContext readerWithoutScroll = new ReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), randomNonNegativeLong(), false);
        DefaultSearchContext contextWithoutScroll = new DefaultSearchContext(readerWithoutScroll, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        contextWithoutScroll.from(300);
        contextWithoutScroll.close();
        // resultWindow greater than maxResultWindow and scrollContext is null
        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> contextWithoutScroll.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Result window is too large, from + size must be less than or equal to:" + " [" + maxResultWindow + "] but was [310]. See the scroll api for a more efficient way to request large data sets. " + "This limit can be set by changing the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey() + "] index level setting."));
        // resultWindow greater than maxResultWindow and scrollContext isn't null
        when(shardSearchRequest.scroll()).thenReturn(new Scroll(TimeValue.timeValueMillis(randomInt(1000))));
        ReaderContext readerContext = new LegacyReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), shardSearchRequest, randomNonNegativeLong());
        DefaultSearchContext context1 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        context1.from(300);
        exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Batch size is too large, size must be less than or equal to: [" + maxResultWindow + "] but was [310]. Scroll batch sizes cost as much memory as result windows so they are " + "controlled by the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey() + "] index level setting."));
        // resultWindow not greater than maxResultWindow and both rescore and sort are not null
        context1.from(0);
        DocValueFormat docValueFormat = mock(DocValueFormat.class);
        SortAndFormats sortAndFormats = new SortAndFormats(new Sort(), new DocValueFormat[] { docValueFormat });
        context1.sort(sortAndFormats);
        RescoreContext rescoreContext = mock(RescoreContext.class);
        when(rescoreContext.getWindowSize()).thenReturn(500);
        context1.addRescore(rescoreContext);
        exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Cannot use [sort] option in conjunction with [rescore]."));
        // rescore is null but sort is not null and rescoreContext.getWindowSize() exceeds maxResultWindow
        context1.sort(null);
        exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Rescore window [" + rescoreContext.getWindowSize() + "] is too large. " + "It must be less than [" + maxRescoreWindow + "]. This prevents allocating massive heaps for storing the results " + "to be rescored. This limit can be set by changing the [" + IndexSettings.MAX_RESCORE_WINDOW_SETTING.getKey() + "] index level setting."));
        readerContext.close();
        readerContext = new ReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), randomNonNegativeLong(), false);
        // rescore is null but sliceBuilder is not null
        DefaultSearchContext context2 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        SliceBuilder sliceBuilder = mock(SliceBuilder.class);
        int numSlices = maxSlicesPerScroll + randomIntBetween(1, 100);
        when(sliceBuilder.getMax()).thenReturn(numSlices);
        context2.sliceBuilder(sliceBuilder);
        exception = expectThrows(IllegalArgumentException.class, () -> context2.preProcess(false));
        assertThat(exception.getMessage(), equalTo("The number of slices [" + numSlices + "] is too large. It must " + "be less than [" + maxSlicesPerScroll + "]. This limit can be set by changing the [" + IndexSettings.MAX_SLICES_PER_SCROLL.getKey() + "] index level setting."));
        // No exceptions should be thrown
        when(shardSearchRequest.getAliasFilter()).thenReturn(AliasFilter.EMPTY);
        when(shardSearchRequest.indexBoost()).thenReturn(AbstractQueryBuilder.DEFAULT_BOOST);
        DefaultSearchContext context3 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        ParsedQuery parsedQuery = ParsedQuery.parsedMatchAllQuery();
        context3.sliceBuilder(null).parsedQuery(parsedQuery).preProcess(false);
        assertEquals(context3.query(), context3.buildFilteredQuery(parsedQuery.query()));
        when(queryShardContext.getIndexSettings()).thenReturn(indexSettings);
        when(queryShardContext.fieldMapper(anyString())).thenReturn(mock(MappedFieldType.class));
        when(shardSearchRequest.indexRoutings()).thenReturn(new String[0]);
        readerContext.close();
        readerContext = new ReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), randomNonNegativeLong(), false);
        DefaultSearchContext context4 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        context4.sliceBuilder(new SliceBuilder(1, 2)).parsedQuery(parsedQuery).preProcess(false);
        Query query1 = context4.query();
        context4.sliceBuilder(new SliceBuilder(0, 2)).parsedQuery(parsedQuery).preProcess(false);
        Query query2 = context4.query();
        assertTrue(query1 instanceof MatchNoDocsQuery || query2 instanceof MatchNoDocsQuery);
        readerContext.close();
        threadPool.shutdown();
    }
}