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

Example 1 with RescoreContext

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

the class ExampleRescoreBuilderTests method testRescore.

public void testRescore() throws IOException {
    // Always use a factor > 1 so rescored fields are sorted in front of the unrescored fields.
    float factor = (float) randomDoubleBetween(1.0d, Float.MAX_VALUE, false);
    // Skipping factorField because it is much harder to mock. We'll catch it in an integration test.
    String fieldFactor = null;
    ExampleRescoreBuilder builder = new ExampleRescoreBuilder(factor, fieldFactor).windowSize(2);
    RescoreContext context = builder.buildContext(null);
    TopDocs docs = new TopDocs(new TotalHits(10, TotalHits.Relation.EQUAL_TO), new ScoreDoc[3]);
    docs.scoreDocs[0] = new ScoreDoc(0, 1.0f);
    docs.scoreDocs[1] = new ScoreDoc(1, 1.0f);
    docs.scoreDocs[2] = new ScoreDoc(2, 1.0f);
    context.rescorer().rescore(docs, null, context);
    assertEquals(factor, docs.scoreDocs[0].score, 0.0f);
    assertEquals(factor, docs.scoreDocs[1].score, 0.0f);
    assertEquals(1.0f, docs.scoreDocs[2].score, 0.0f);
}

Example 2 with RescoreContext

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

the class ExplainPhase method getProcessor.

@Override
public FetchSubPhaseProcessor getProcessor(FetchContext context) {
    if (context.explain() == false) {
        return null;
    }
    return new FetchSubPhaseProcessor() {

        @Override
        public void setNextReader(LeafReaderContext readerContext) {
        }

        @Override
        public void process(HitContext hitContext) throws IOException {
            final int topLevelDocId = hitContext.hit().docId();
            Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
            for (RescoreContext rescore : context.rescore()) {
                explanation = rescore.rescorer().explain(topLevelDocId, context.searcher(), rescore, explanation);
            }
            // we use the top level doc id, since we work with the top level searcher
            hitContext.hit().explanation(explanation);
        }
    };
}

Example 3 with RescoreContext

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

the class DefaultSearchContext method preProcess.

/**
 * Should be called before executing the main query and after all other parameters have been set.
 */
@Override
public void preProcess(boolean rewrite) {
    if (hasOnlySuggest()) {
        return;
    }
    long from = from() == -1 ? 0 : from();
    long size = size() == -1 ? 10 : size();
    long resultWindow = from + size;
    int maxResultWindow = indexService.getIndexSettings().getMaxResultWindow();
    if (resultWindow > maxResultWindow) {
        if (scrollContext() == null) {
            throw new IllegalArgumentException("Result window is too large, from + size must be less than or equal to: [" + maxResultWindow + "] but was [" + resultWindow + "]. 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.");
        }
        throw new IllegalArgumentException("Batch size is too large, size must be less than or equal to: [" + maxResultWindow + "] but was [" + resultWindow + "]. 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.");
    }
    if (rescore != null) {
        if (sort != null) {
            throw new IllegalArgumentException("Cannot use [sort] option in conjunction with [rescore].");
        }
        int maxWindow = indexService.getIndexSettings().getMaxRescoreWindow();
        for (RescoreContext rescoreContext : rescore()) {
            if (rescoreContext.getWindowSize() > maxWindow) {
                throw new IllegalArgumentException("Rescore window [" + rescoreContext.getWindowSize() + "] is too large. " + "It must be less than [" + maxWindow + "]. 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.");
            }
        }
    }
    if (sliceBuilder != null) {
        int sliceLimit = indexService.getIndexSettings().getMaxSlicesPerScroll();
        int numSlices = sliceBuilder.getMax();
        if (numSlices > sliceLimit) {
            throw new IllegalArgumentException("The number of slices [" + numSlices + "] is too large. It must " + "be less than [" + sliceLimit + "]. This limit can be set by changing the [" + IndexSettings.MAX_SLICES_PER_SCROLL.getKey() + "] index level setting.");
        }
    }
    // initialize the filtering alias based on the provided filters
    try {
        final QueryBuilder queryBuilder = request.getAliasFilter().getQueryBuilder();
        aliasFilter = queryBuilder == null ? null : queryBuilder.toQuery(queryShardContext);
    } catch (IOException e) {
        throw new UncheckedIOException(e);
    }
    if (query() == null) {
        parsedQuery(ParsedQuery.parsedMatchAllQuery());
    }
    if (queryBoost() != AbstractQueryBuilder.DEFAULT_BOOST) {
        parsedQuery(new ParsedQuery(new BoostQuery(query(), queryBoost), parsedQuery()));
    }
    this.query = buildFilteredQuery(query);
    if (rewrite) {
        try {
            this.query = searcher.rewrite(query);
        } catch (IOException e) {
            throw new QueryPhaseExecutionException(shardTarget, "Failed to rewrite main query", e);
        }
    }
}

Example 4 with RescoreContext

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

the class SearchContext method rescoreDocIds.

public final RescoreDocIds rescoreDocIds() {
    final List<RescoreContext> rescore = rescore();
    if (rescore == null) {
        return RescoreDocIds.EMPTY;
    }
    Map<Integer, Set<Integer>> rescoreDocIds = null;
    for (int i = 0; i < rescore.size(); i++) {
        final Set<Integer> docIds = rescore.get(i).getRescoredDocs();
        if (docIds != null && docIds.isEmpty() == false) {
            if (rescoreDocIds == null) {
                rescoreDocIds = new HashMap<>();
            }
            rescoreDocIds.put(i, docIds);
        }
    }
    return rescoreDocIds == null ? RescoreDocIds.EMPTY : new RescoreDocIds(rescoreDocIds);
}

Example 5 with RescoreContext

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

the class TopDocsCollectorContext method createTopDocsCollectorContext.

/**
 * Creates a {@link TopDocsCollectorContext} from the provided <code>searchContext</code>.
 * @param hasFilterCollector True if the collector chain contains at least one collector that can filters document.
 */
static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searchContext, boolean hasFilterCollector) throws IOException {
    final IndexReader reader = searchContext.searcher().getIndexReader();
    final Query query = searchContext.query();
    // top collectors don't like a size of 0
    final int totalNumDocs = Math.max(1, reader.numDocs());
    if (searchContext.size() == 0) {
        // no matter what the value of from is
        return new EmptyTopDocsCollectorContext(reader, query, searchContext.sort(), searchContext.trackTotalHitsUpTo(), hasFilterCollector);
    } else if (searchContext.scrollContext() != null) {
        // we can disable the tracking of total hits after the initial scroll query
        // since the total hits is preserved in the scroll context.
        int trackTotalHitsUpTo = searchContext.scrollContext().totalHits != null ? SearchContext.TRACK_TOTAL_HITS_DISABLED : SearchContext.TRACK_TOTAL_HITS_ACCURATE;
        // no matter what the value of from is
        int numDocs = Math.min(searchContext.size(), totalNumDocs);
        return new ScrollingTopDocsCollectorContext(reader, query, searchContext.scrollContext(), searchContext.sort(), numDocs, searchContext.trackScores(), searchContext.numberOfShards(), trackTotalHitsUpTo, hasFilterCollector);
    } else if (searchContext.collapse() != null) {
        boolean trackScores = searchContext.sort() == null ? true : searchContext.trackScores();
        int numDocs = Math.min(searchContext.from() + searchContext.size(), totalNumDocs);
        return new CollapsingTopDocsCollectorContext(searchContext.collapse(), searchContext.sort(), numDocs, trackScores);
    } else {
        int numDocs = Math.min(searchContext.from() + searchContext.size(), totalNumDocs);
        final boolean rescore = searchContext.rescore().isEmpty() == false;
        if (rescore) {
            assert searchContext.sort() == null;
            for (RescoreContext rescoreContext : searchContext.rescore()) {
                numDocs = Math.max(numDocs, rescoreContext.getWindowSize());
            }
        }
        return new SimpleTopDocsCollectorContext(reader, query, searchContext.sort(), searchContext.searchAfter(), numDocs, searchContext.trackScores(), searchContext.trackTotalHitsUpTo(), hasFilterCollector) {

            @Override
            boolean shouldRescore() {
                return rescore;
            }
        };
    }
}