Examples with InternalProfileCollector org.opensearch.search.profile.query.InternalProfileCollector used on opensource projects

Example 1 with InternalProfileCollector

use of org.opensearch.search.profile.query.InternalProfileCollector in project OpenSearch by opensearch-project.

the class QueryCollectorContext method createMultiCollectorContext.

/**
 * Creates a multi collector from the provided <code>subs</code>
 */
static QueryCollectorContext createMultiCollectorContext(Collection<Collector> subs) {
    return new QueryCollectorContext(REASON_SEARCH_MULTI) {

        @Override
        Collector create(Collector in) {
            List<Collector> subCollectors = new ArrayList<>();
            subCollectors.add(in);
            subCollectors.addAll(subs);
            return MultiCollector.wrap(subCollectors);
        }

        @Override
        protected InternalProfileCollector createWithProfiler(InternalProfileCollector in) {
            final List<InternalProfileCollector> subCollectors = new ArrayList<>();
            subCollectors.add(in);
            if (subs.stream().anyMatch((col) -> col instanceof InternalProfileCollector == false)) {
                throw new IllegalArgumentException("non-profiling collector");
            }
            for (Collector collector : subs) {
                subCollectors.add((InternalProfileCollector) collector);
            }
            final Collector collector = MultiCollector.wrap(subCollectors);
            return new InternalProfileCollector(collector, REASON_SEARCH_MULTI, subCollectors);
        }
    };
}

Example 2 with InternalProfileCollector

use of org.opensearch.search.profile.query.InternalProfileCollector in project OpenSearch by opensearch-project.

the class AggregationPhase method execute.

public void execute(SearchContext context) {
    if (context.aggregations() == null) {
        context.queryResult().aggregations(null);
        return;
    }
    if (context.queryResult().hasAggs()) {
        // no need to compute the aggs twice, they should be computed on a per context basis
        return;
    }
    Aggregator[] aggregators = context.aggregations().aggregators();
    List<Aggregator> globals = new ArrayList<>();
    for (int i = 0; i < aggregators.length; i++) {
        if (aggregators[i] instanceof GlobalAggregator) {
            globals.add(aggregators[i]);
        }
    }
    // optimize the global collector based execution
    if (!globals.isEmpty()) {
        BucketCollector globalsCollector = MultiBucketCollector.wrap(globals);
        Query query = context.buildFilteredQuery(Queries.newMatchAllQuery());
        try {
            final Collector collector;
            if (context.getProfilers() == null) {
                collector = globalsCollector;
            } else {
                InternalProfileCollector profileCollector = new InternalProfileCollector(globalsCollector, CollectorResult.REASON_AGGREGATION_GLOBAL, // TODO: report on sub collectors
                Collections.emptyList());
                collector = profileCollector;
                // start a new profile with this collector
                context.getProfilers().addQueryProfiler().setCollector(profileCollector);
            }
            globalsCollector.preCollection();
            context.searcher().search(query, collector);
        } catch (Exception e) {
            throw new QueryPhaseExecutionException(context.shardTarget(), "Failed to execute global aggregators", e);
        }
    }
    List<InternalAggregation> aggregations = new ArrayList<>(aggregators.length);
    context.aggregations().resetBucketMultiConsumer();
    for (Aggregator aggregator : context.aggregations().aggregators()) {
        try {
            aggregator.postCollection();
            aggregations.add(aggregator.buildTopLevel());
        } catch (IOException e) {
            throw new AggregationExecutionException("Failed to build aggregation [" + aggregator.name() + "]", e);
        }
    }
    context.queryResult().aggregations(new InternalAggregations(aggregations, context.request().source().aggregations()::buildPipelineTree));
    // disable aggregations so that they don't run on next pages in case of scrolling
    context.aggregations(null);
    context.queryCollectors().remove(AggregationPhase.class);
}

Example 3 with InternalProfileCollector

use of org.opensearch.search.profile.query.InternalProfileCollector in project OpenSearch by opensearch-project.

the class AggregationPhase method preProcess.

public void preProcess(SearchContext context) {
    if (context.aggregations() != null) {
        List<Aggregator> collectors = new ArrayList<>();
        Aggregator[] aggregators;
        try {
            AggregatorFactories factories = context.aggregations().factories();
            aggregators = factories.createTopLevelAggregators(context);
            for (int i = 0; i < aggregators.length; i++) {
                if (aggregators[i] instanceof GlobalAggregator == false) {
                    collectors.add(aggregators[i]);
                }
            }
            context.aggregations().aggregators(aggregators);
            if (!collectors.isEmpty()) {
                Collector collector = MultiBucketCollector.wrap(collectors);
                ((BucketCollector) collector).preCollection();
                if (context.getProfilers() != null) {
                    collector = new InternalProfileCollector(collector, CollectorResult.REASON_AGGREGATION, // TODO: report on child aggs as well
                    Collections.emptyList());
                }
                context.queryCollectors().put(AggregationPhase.class, collector);
            }
        } catch (IOException e) {
            throw new AggregationInitializationException("Could not initialize aggregators", e);
        }
    }
}

Example 4 with InternalProfileCollector

use of org.opensearch.search.profile.query.InternalProfileCollector in project OpenSearch by opensearch-project.

the class QueryPhase method searchWithCollector.

private static boolean searchWithCollector(SearchContext searchContext, ContextIndexSearcher searcher, Query query, LinkedList<QueryCollectorContext> collectors, boolean hasFilterCollector, boolean timeoutSet) throws IOException {
    // create the top docs collector last when the other collectors are known
    final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
    // add the top docs collector, the first collector context in the chain
    collectors.addFirst(topDocsFactory);
    final Collector queryCollector;
    if (searchContext.getProfilers() != null) {
        InternalProfileCollector profileCollector = QueryCollectorContext.createQueryCollectorWithProfiler(collectors);
        searchContext.getProfilers().getCurrentQueryProfiler().setCollector(profileCollector);
        queryCollector = profileCollector;
    } else {
        queryCollector = QueryCollectorContext.createQueryCollector(collectors);
    }
    QuerySearchResult queryResult = searchContext.queryResult();
    try {
        searcher.search(query, queryCollector);
    } catch (EarlyTerminatingCollector.EarlyTerminationException e) {
        queryResult.terminatedEarly(true);
    } catch (TimeExceededException e) {
        assert timeoutSet : "TimeExceededException thrown even though timeout wasn't set";
        if (searchContext.request().allowPartialSearchResults() == false) {
            // Can't rethrow TimeExceededException because not serializable
            throw new QueryPhaseExecutionException(searchContext.shardTarget(), "Time exceeded");
        }
        queryResult.searchTimedOut(true);
    }
    if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER && queryResult.terminatedEarly() == null) {
        queryResult.terminatedEarly(false);
    }
    for (QueryCollectorContext ctx : collectors) {
        ctx.postProcess(queryResult);
    }
    return topDocsFactory.shouldRescore();
}