Search in sources :

Example 1 with SearchAfterSortedDocQuery

use of org.apache.lucene.queries.SearchAfterSortedDocQuery in project OpenSearch by opensearch-project.

the class CompositeAggregator method processLeafFromQuery.

private void processLeafFromQuery(LeafReaderContext ctx, Sort indexSortPrefix) throws IOException {
    DocValueFormat[] formats = new DocValueFormat[indexSortPrefix.getSort().length];
    for (int i = 0; i < formats.length; i++) {
        formats[i] = sources[i].format;
    }
    FieldDoc fieldDoc = SearchAfterBuilder.buildFieldDoc(new SortAndFormats(indexSortPrefix, formats), Arrays.copyOfRange(rawAfterKey.values(), 0, formats.length));
    if (indexSortPrefix.getSort().length < sources.length) {
        // include all docs that belong to the partial bucket
        fieldDoc.doc = -1;
    }
    BooleanQuery newQuery = new BooleanQuery.Builder().add(context.query(), BooleanClause.Occur.MUST).add(new SearchAfterSortedDocQuery(applySortFieldRounding(indexSortPrefix), fieldDoc), BooleanClause.Occur.FILTER).build();
    Weight weight = context.searcher().createWeight(context.searcher().rewrite(newQuery), ScoreMode.COMPLETE_NO_SCORES, 1f);
    Scorer scorer = weight.scorer(ctx);
    if (scorer != null) {
        DocIdSetIterator docIt = scorer.iterator();
        final LeafBucketCollector inner = queue.getLeafCollector(ctx, getFirstPassCollector(docIdSetBuilder, indexSortPrefix.getSort().length));
        inner.setScorer(scorer);
        final Bits liveDocs = ctx.reader().getLiveDocs();
        while (docIt.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
            if (liveDocs == null || liveDocs.get(docIt.docID())) {
                inner.collect(docIt.docID());
            }
        }
    }
}
Also used : BooleanQuery(org.apache.lucene.search.BooleanQuery) FieldDoc(org.apache.lucene.search.FieldDoc) DocValueFormat(org.opensearch.search.DocValueFormat) Scorer(org.apache.lucene.search.Scorer) SearchAfterSortedDocQuery(org.apache.lucene.queries.SearchAfterSortedDocQuery) SortAndFormats(org.opensearch.search.sort.SortAndFormats) Weight(org.apache.lucene.search.Weight) LeafBucketCollector(org.opensearch.search.aggregations.LeafBucketCollector) Bits(org.apache.lucene.util.Bits) DocIdSetIterator(org.apache.lucene.search.DocIdSetIterator)

Example 2 with SearchAfterSortedDocQuery

use of org.apache.lucene.queries.SearchAfterSortedDocQuery in project OpenSearch by opensearch-project.

the class QueryPhase method executeInternal.

/**
 * In a package-private method so that it can be tested without having to
 * wire everything (mapperService, etc.)
 * @return whether the rescoring phase should be executed
 */
static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExecutionException {
    final ContextIndexSearcher searcher = searchContext.searcher();
    final IndexReader reader = searcher.getIndexReader();
    QuerySearchResult queryResult = searchContext.queryResult();
    queryResult.searchTimedOut(false);
    try {
        queryResult.from(searchContext.from());
        queryResult.size(searchContext.size());
        Query query = searchContext.query();
        // already rewritten
        assert query == searcher.rewrite(query);
        final ScrollContext scrollContext = searchContext.scrollContext();
        if (scrollContext != null) {
            if (scrollContext.totalHits == null) {
                // first round
                assert scrollContext.lastEmittedDoc == null;
            // there is not much that we can optimize here since we want to collect all
            // documents in order to get the total number of hits
            } else {
                final ScoreDoc after = scrollContext.lastEmittedDoc;
                if (returnsDocsInOrder(query, searchContext.sort())) {
                    // skip to the desired doc
                    if (after != null) {
                        query = new BooleanQuery.Builder().add(query, BooleanClause.Occur.MUST).add(new MinDocQuery(after.doc + 1), BooleanClause.Occur.FILTER).build();
                    }
                    // ... and stop collecting after ${size} matches
                    searchContext.terminateAfter(searchContext.size());
                } else if (canEarlyTerminate(reader, searchContext.sort())) {
                    // skip to the desired doc
                    if (after != null) {
                        query = new BooleanQuery.Builder().add(query, BooleanClause.Occur.MUST).add(new SearchAfterSortedDocQuery(searchContext.sort().sort, (FieldDoc) after), BooleanClause.Occur.FILTER).build();
                    }
                }
            }
        }
        final LinkedList<QueryCollectorContext> collectors = new LinkedList<>();
        // whether the chain contains a collector that filters documents
        boolean hasFilterCollector = false;
        if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) {
            // add terminate_after before the filter collectors
            // it will only be applied on documents accepted by these filter collectors
            collectors.add(createEarlyTerminationCollectorContext(searchContext.terminateAfter()));
            // this collector can filter documents during the collection
            hasFilterCollector = true;
        }
        if (searchContext.parsedPostFilter() != null) {
            // add post filters before aggregations
            // it will only be applied to top hits
            collectors.add(createFilteredCollectorContext(searcher, searchContext.parsedPostFilter().query()));
            // this collector can filter documents during the collection
            hasFilterCollector = true;
        }
        if (searchContext.queryCollectors().isEmpty() == false) {
            // plug in additional collectors, like aggregations
            collectors.add(createMultiCollectorContext(searchContext.queryCollectors().values()));
        }
        if (searchContext.minimumScore() != null) {
            // apply the minimum score after multi collector so we filter aggs as well
            collectors.add(createMinScoreCollectorContext(searchContext.minimumScore()));
            // this collector can filter documents during the collection
            hasFilterCollector = true;
        }
        // optimizing sort on Numerics (long and date)
        if ((searchContext.sort() != null) && SYS_PROP_REWRITE_SORT) {
            enhanceSortOnNumeric(searchContext, searcher.getIndexReader());
        }
        boolean timeoutSet = scrollContext == null && searchContext.timeout() != null && searchContext.timeout().equals(SearchService.NO_TIMEOUT) == false;
        final Runnable timeoutRunnable;
        if (timeoutSet) {
            final long startTime = searchContext.getRelativeTimeInMillis();
            final long timeout = searchContext.timeout().millis();
            final long maxTime = startTime + timeout;
            timeoutRunnable = searcher.addQueryCancellation(() -> {
                final long time = searchContext.getRelativeTimeInMillis();
                if (time > maxTime) {
                    throw new TimeExceededException();
                }
            });
        } else {
            timeoutRunnable = null;
        }
        if (searchContext.lowLevelCancellation()) {
            searcher.addQueryCancellation(() -> {
                SearchShardTask task = searchContext.getTask();
                if (task != null && task.isCancelled()) {
                    throw new TaskCancelledException("cancelled task with reason: " + task.getReasonCancelled());
                }
            });
        }
        try {
            boolean shouldRescore = searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, timeoutSet);
            ExecutorService executor = searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH);
            if (executor instanceof QueueResizingOpenSearchThreadPoolExecutor) {
                QueueResizingOpenSearchThreadPoolExecutor rExecutor = (QueueResizingOpenSearchThreadPoolExecutor) executor;
                queryResult.nodeQueueSize(rExecutor.getCurrentQueueSize());
                queryResult.serviceTimeEWMA((long) rExecutor.getTaskExecutionEWMA());
            }
            return shouldRescore;
        } finally {
            // otherwise aggregation phase might get cancelled.
            if (timeoutRunnable != null) {
                searcher.removeQueryCancellation(timeoutRunnable);
            }
        }
    } catch (Exception e) {
        throw new QueryPhaseExecutionException(searchContext.shardTarget(), "Failed to execute main query", e);
    }
}
Also used : QueueResizingOpenSearchThreadPoolExecutor(org.opensearch.common.util.concurrent.QueueResizingOpenSearchThreadPoolExecutor) BooleanQuery(org.apache.lucene.search.BooleanQuery) Query(org.apache.lucene.search.Query) SearchAfterSortedDocQuery(org.apache.lucene.queries.SearchAfterSortedDocQuery) MatchAllDocsQuery(org.apache.lucene.search.MatchAllDocsQuery) MinDocQuery(org.apache.lucene.queries.MinDocQuery) ConstantScoreQuery(org.apache.lucene.search.ConstantScoreQuery) BooleanQuery(org.apache.lucene.search.BooleanQuery) ContextIndexSearcher(org.opensearch.search.internal.ContextIndexSearcher) ScrollContext(org.opensearch.search.internal.ScrollContext) SearchAfterSortedDocQuery(org.apache.lucene.queries.SearchAfterSortedDocQuery) LinkedList(java.util.LinkedList) TaskCancelledException(org.opensearch.tasks.TaskCancelledException) IOException(java.io.IOException) ScoreDoc(org.apache.lucene.search.ScoreDoc) MinDocQuery(org.apache.lucene.queries.MinDocQuery) IndexReader(org.apache.lucene.index.IndexReader) ExecutorService(java.util.concurrent.ExecutorService) TaskCancelledException(org.opensearch.tasks.TaskCancelledException) SearchShardTask(org.opensearch.action.search.SearchShardTask)

Aggregations

SearchAfterSortedDocQuery (org.apache.lucene.queries.SearchAfterSortedDocQuery)2 BooleanQuery (org.apache.lucene.search.BooleanQuery)2 IOException (java.io.IOException)1 LinkedList (java.util.LinkedList)1 ExecutorService (java.util.concurrent.ExecutorService)1 IndexReader (org.apache.lucene.index.IndexReader)1 MinDocQuery (org.apache.lucene.queries.MinDocQuery)1 ConstantScoreQuery (org.apache.lucene.search.ConstantScoreQuery)1 DocIdSetIterator (org.apache.lucene.search.DocIdSetIterator)1 FieldDoc (org.apache.lucene.search.FieldDoc)1 MatchAllDocsQuery (org.apache.lucene.search.MatchAllDocsQuery)1 Query (org.apache.lucene.search.Query)1 ScoreDoc (org.apache.lucene.search.ScoreDoc)1 Scorer (org.apache.lucene.search.Scorer)1 Weight (org.apache.lucene.search.Weight)1 Bits (org.apache.lucene.util.Bits)1 SearchShardTask (org.opensearch.action.search.SearchShardTask)1 QueueResizingOpenSearchThreadPoolExecutor (org.opensearch.common.util.concurrent.QueueResizingOpenSearchThreadPoolExecutor)1 DocValueFormat (org.opensearch.search.DocValueFormat)1 LeafBucketCollector (org.opensearch.search.aggregations.LeafBucketCollector)1