Examples with CollectorCreatorContext com.yelp.nrtsearch.server.luceneserver.search.collectors.CollectorCreatorContext used on opensource projects

Example 1 with CollectorCreatorContext

use of com.yelp.nrtsearch.server.luceneserver.search.collectors.CollectorCreatorContext in project nrtsearch by Yelp.

the class SearchRequestProcessor method buildContextForRequest.

/**
 * Create a {@link SearchContext} representing the given {@link SearchRequest} and index/searcher
 * state.
 *
 * @param searchRequest grpc request message
 * @param indexState index state
 * @param shardState shard state
 * @param searcherAndTaxonomy index searcher
 * @param profileResult container message for returned debug info
 * @return context info needed to execute the search query
 * @throws IOException if query rewrite fails
 */
public static SearchContext buildContextForRequest(SearchRequest searchRequest, IndexState indexState, ShardState shardState, SearcherTaxonomyManager.SearcherAndTaxonomy searcherAndTaxonomy, ProfileResult.Builder profileResult) throws IOException {
    SearchContext.Builder contextBuilder = SearchContext.newBuilder();
    SearchResponse.Builder responseBuilder = SearchResponse.newBuilder();
    contextBuilder.setIndexState(indexState).setShardState(shardState).setSearcherAndTaxonomy(searcherAndTaxonomy).setResponseBuilder(responseBuilder).setTimestampSec(System.currentTimeMillis() / 1000).setStartHit(searchRequest.getStartHit()).setTopHits(searchRequest.getTopHits());
    Map<String, FieldDef> queryVirtualFields = getVirtualFields(shardState, searchRequest);
    Map<String, FieldDef> queryFields = new HashMap<>(queryVirtualFields);
    addIndexFields(indexState, queryFields);
    contextBuilder.setQueryFields(Collections.unmodifiableMap(queryFields));
    Map<String, FieldDef> retrieveFields = getRetrieveFields(searchRequest, queryFields);
    contextBuilder.setRetrieveFields(Collections.unmodifiableMap(retrieveFields));
    Query query = extractQuery(indexState, searchRequest);
    if (profileResult != null) {
        profileResult.setParsedQuery(query.toString());
    }
    query = searcherAndTaxonomy.searcher.rewrite(query);
    if (profileResult != null) {
        profileResult.setRewrittenQuery(query.toString());
    }
    if (searchRequest.getFacetsCount() > 0) {
        query = addDrillDowns(indexState, query);
        if (profileResult != null) {
            profileResult.setDrillDownQuery(query.toString());
        }
    }
    contextBuilder.setFetchTasks(new FetchTasks(searchRequest.getFetchTasksList()));
    contextBuilder.setQuery(query);
    CollectorCreatorContext collectorCreatorContext = new CollectorCreatorContext(searchRequest, indexState, shardState, queryFields);
    DocCollector docCollector = buildDocCollector(collectorCreatorContext);
    contextBuilder.setCollector(docCollector);
    contextBuilder.setRescorers(getRescorers(indexState, searcherAndTaxonomy.searcher, searchRequest));
    contextBuilder.setSharedDocContext(new DefaultSharedDocContext());
    SearchContext searchContext = contextBuilder.build(true);
    // Give underlying collectors access to the search context
    docCollector.setSearchContext(searchContext);
    return searchContext;
}

Example 2 with CollectorCreatorContext

use of com.yelp.nrtsearch.server.luceneserver.search.collectors.CollectorCreatorContext in project nrtsearch by Yelp.

the class SearchRequestProcessor method buildDocCollector.

/**
 * Build {@link DocCollector} to provide the {@link org.apache.lucene.search.CollectorManager} for
 * collecting hits for this query.
 *
 * @return collector
 */
private static DocCollector buildDocCollector(CollectorCreatorContext collectorCreatorContext) {
    SearchRequest searchRequest = collectorCreatorContext.getRequest();
    List<AdditionalCollectorManager<? extends Collector, ? extends CollectorResult>> additionalCollectors = searchRequest.getCollectorsMap().entrySet().stream().map(e -> CollectorCreator.getInstance().createCollectorManager(collectorCreatorContext, e.getKey(), e.getValue())).collect(Collectors.toList());
    if (searchRequest.getQuerySort().getFields().getSortedFieldsList().isEmpty()) {
        if (hasLargeNumHits(searchRequest)) {
            return new LargeNumHitsCollector(collectorCreatorContext, additionalCollectors);
        } else {
            return new RelevanceCollector(collectorCreatorContext, additionalCollectors);
        }
    } else {
        return new SortFieldCollector(collectorCreatorContext, additionalCollectors);
    }
}