[Feature Request] Enable injecting collector context from the plugin in the QueryPhase directly

[vibrantvarun]

Is your feature request related to a problem? Please describe

The current workflow of the search in OpenSearch is divided into two main phases: The Query Phase (QP) and The Fetch Phase (FP). At a high level, the QP first loads the Aggregation Processor and then performs its pre-processing. Following that, it loads the collectors and initializes the collector manager to execute the search on the shard. Later, the aggregation processor performs the post-processing of the results.

For traditional searches like bool, match, and term, the top docs collector context is created. The collector context internally initializes the TopDocsCollectorManager and loads the TopDocsCollector into it. The initialization and usage of TopDocsCollector is hardcoded in the search process. This creates a limitation where plugins cannot inject a custom collector context during the search in an ideal place when the TopDocsCollector gets instantiated.

A classic example of a query that has custom logic is the hybrid query, which resides in the neural search plugin. Due to the limitation mentioned earlier, the neural search plugin has to inject the HybridCollectorManager and HybridTopScoreDocCollector during the aggregation pre-process phase. It also has to provide a custom aggregation processor called HybridQueryAggregationProcessor, which is essentially a wrapper around DefaultAggregationProcessor. Moreover, in order to skip TopDocsCollectorContext initialization during HQ execution, there is a parody of the searchWithCollector method which injects an empty collector context in the search.

Recently, the team has performed the POC of moving Hybrid Search to OpenSearch core. There were multiple phases in which the POC was done and at each phase the benchmarking was performed. The baseline of these benchmarks is the current hybrid query implementation in the neural search plugin.

Dataset: noaa-semantic-search

Phase 1: Move Hybrid Query logic and all of its related classes to OpenSearch core. This also includes Normalization processor. Here the assumption is doing so will reduce the network calls.

Hybrid query with 3 subqueries: Term, Range and Date

Latency	3.0-beta (Phase 1) (Min distribution)	3.0-beta (GA )
p50	306.41	280.16
p90	348.01	299.51
p99	405.09	326.92
p100	434.87	334.57

No improvement was observed. The reason is that the plugin and core run in the same JVM, so it does not provide a performance boost. The degradation in latency is observed in the min distribution because it is not as stable as the GA version. Also, since the code in the min distribution is an MVP, we are just looking for any small improvement.

But, from the above experiment it can be said that a custom query can either lie in the OS core or a plugin it does not have an impact on network calls.

Phase 2: Create HybridQueryCollectorContext and inject it in the same way how TopDocsCollectorContext is injected. Also remove the EmptyCollectorContext initialization and HybridQueryAggregationProcessor and switch to DefaultAggregationProcessor

Latency	3.0-beta (Phase 2) (Min distribution)	3.0-beta (GA )	Improvement
p50	259.81	280.16	----> 7.26% improvement
p90	278.41	299.51	----> 7.04% improvement
p99	298.77	326.92	----> 8.61% improvement
p100	313.3	334.57	----> 6.35% improvement

We clearly see an improvement. To further bolster the improvement, we did a performance benchmarking of complete distribution tarball with the GA one.

Latency	3.0-beta (POC Complete distribution)	3.0-beta (GA )	Improvement
p50	246.99	280.16	--> 11.83% improvement
p90	250.9	299.51	--> 16.22% improvement
p99	289.33	326.92	--> 11.49% improvement
p100	324.19	334.57	--> 3.10% improvement

Therefore, if we make the process of injecting collector context extensible through plugins, it will help custom query types improve their performance.

Essentially, in the searchWithCollector method, the QueryCollectorContext can be injected by the plugin like how we injected in the POC for hybrid query.

Describe the solution you'd like

We can make the createQueryCollectorContext method extensible and the plugins can provide there implementation.

Related component

Search:Performance

Describe alternatives you've considered

No response

Additional context

No response

[martin-gaievski]

Good suggestion, thanks for opening this issue Varun.

I think it can be done in parallel with request for a new extension point that has been created some time ago: #17593. It should help to address the workaround of using AggregationProcessor to reduce collected results, you need to do this anyways, doesn't matter how score collector is called.

[sandeshkr419]

[Search Triage] @vibrantvarun Thanks for doing the POC on this. Do you also have CPU/JVM or any other regression impact with these changes as well?

[vibrantvarun]

CPU Utilization	3.0-beta (Optimized Complete Distribution	3.0-beta GA	Improvement/Decrement
Max Cpu Utilization	16.45%	18.59%	2.14% improvement
Avg Cpu Utilization	4%	3.75%	0.25% decrement

cc: @sandeshkr419 @martin-gaievski

[navneet1v]

@vibrantvarun could you share some diagram on the collector injection will work?

[vibrantvarun]

JVM heap utilization

3.0.0-beta GA

Average JVM heap Utilization

Max JVM heap Utilization

3.0.0-beta GA (Complete Optimization)

Average JVM heap Utilization

Max JVM heap Utilization

JVM Utilization	3.0-beta (Optimized Complete Distribution	3.0-beta GA	Improvement/Decrement
Max JVM Utilization	61%	61%	no improvement
Avg JVM Utilization	38.41%	51.5%	13.09% improvement

Conclusion: There is significant improvement (~13%) in the average JVM utilization of the optimization version.

cc: @sandeshkr419 @martin-gaievski @navneet1v

[vibrantvarun]

HLD

Following HLD proposes to open access QueryCollectorContext class to be extensible from plugin and then use registry and factory design patterns in registering query collector context and extending it respectively.

Conclusion

In short, if we open up the QueryCollectorContext class then we see a 13% improvement in average cpu utilization in heap for hybrid query.

cc: @reta @msfroh @navneet1v @martin-gaievski

[vibrantvarun]

In Plugin, we can extend following method

 @Override
    public List<FactorySpec<?>> getCollectorContextFactory(){
        return Arrays.asList(
                new FactorySpec<>(HybridQueryBuilder.NAME, HybridQuery.class, new HybridQueryCollectorContextFactory())
        );
    }

If the plugin does not provide its own implementation of collectorContext then we use default TopDocsCollectorContext.

The searchWithCollector method of QueryPhase can be modified as

protected boolean searchWithCollector(
            SearchContext searchContext,
            ContextIndexSearcher searcher,
            Query query,
            LinkedList<QueryCollectorContext> collectors,
            boolean hasFilterCollector,
            boolean hasTimeout
        ) throws IOException {
            // create the top docs collector last when the other collectors are known
            QueryCollectorContextFactory queryCollectorContextFactory = QueryCollectorContextFactoryRegistery.getFactory(query);
            CollectorContextParams collectorContextParams = new CollectorContextParams.Builder().withSearchContext(searchContext)
                .withHasFilterCollector(hasFilterCollector)
                .build();
            if (queryCollectorContextFactory == null) {
                queryCollectorContextFactory = new TopDocsCollectorContextFactory();
            }
            final QueryCollectorContext queryCollectorContext = queryCollectorContextFactory.createCollectorContext(collectorContextParams);
            return searchWithCollector(searchContext, searcher, query, collectors, queryCollectorContext, hasFilterCollector, hasTimeout);
        }

This way if in future we introduce any other query type which does not need TopDocsCollector to be loaded and has its own collectorContext implementation then we can easily make it possible by earlier proposed HLD.

[reta]

Thanks @vibrantvarun for the proposal

Following HLD proposes to open access QueryCollectorContext class to be extensible from plugin and then use registry and factory design patterns in registering query collector context and extending it respectively.

The QueryCollectorContext is public api already, but QueryPhaseSearcher implementation do not offer the option to instantiate the anything but TopDocsCollectorContext by default. What if we add newQueryCollectorContext to existing QueryPhaseSearcher so the plugins / extensions could supply own QueryCollectorContext while concisely using the same (single) APIs?

public interface QueryPhaseSearcher {
        @Override
        public default QueryCollectorContext newQueryCollectorContext(SearchContext searchContext, boolean hasFilterCollector) {
            return createTopDocsCollectorContext(searchContext, hasFilterCollector);
        } 
}

[vibrantvarun]

@reta Earlier when I mentioned to open the access it meant there are few methods in QueryCollectorContext class which are package-private like postProcess, createCollector, createManager etc to make them public.

I liked the idea you proposed, I will evaluate the feasibility of the proposed solution once. But this solution would also need the earlier called out methods to become public because then only we will able to create HybridQueryCollectorContex in NS plugin in the following way

public class HybridCollectorContext extends QueryCollectorContext {
    protected final @Nullable SortAndFormats sortAndFormats;
    private final int numHits;
    private final int trackTotalHitsUpTo;
    private final HitsThresholdChecker hitsThresholdChecker;
    private final @Nullable Weight filterWeight;
    private final TopDocsMerger topDocsMerger;
    private final @Nullable FieldDoc after;
    private final SearchContext searchContext;
    private final Collector collector;
    private HybridCollectorManager manager;

    public static HybridCollectorContext createHybridCollectorContext(final SearchContext searchContext) throws IOException {
        if (searchContext.scrollContext() != null) {
            throw new IllegalArgumentException(String.format(Locale.ROOT, "Scroll operation is not supported in hybrid query"));
        }
        final IndexReader reader = searchContext.searcher().getIndexReader();
        final int totalNumDocs = Math.max(0, reader.numDocs());
        int numDocs = Math.min(getSubqueryResultsRetrievalSize(searchContext), totalNumDocs);
        Weight filteringWeight = null;
        // Check for post filter to create weight for filter query and later use that weight in the search workflow
        if (Objects.nonNull(searchContext.parsedPostFilter()) && Objects.nonNull(searchContext.parsedPostFilter().query())) {
            Query filterQuery = searchContext.parsedPostFilter().query();
            ContextIndexSearcher searcher = searchContext.searcher();
            // ScoreMode COMPLETE_NO_SCORES will be passed as post_filter does not contribute in scoring. COMPLETE_NO_SCORES means it is not
            // a scoring clause
            // Boost factor 1f is taken because if boost is multiplicative of 1 then it means "no boost"
            // Previously this code in OpenSearch looked like
            // https://github.com/opensearch-project/OpenSearch/commit/36a5cf8f35e5cbaa1ff857b5a5db8c02edc1a187
            filteringWeight = searcher.createWeight(searcher.rewrite(filterQuery), ScoreMode.COMPLETE_NO_SCORES, 1f);
        }
        SortAndFormats sort = searchContext.sort();
        if (sort != null) {
            validateSortCriteria(searchContext, searchContext.trackScores());
        }

        boolean isSingleShard = searchContext.numberOfShards() == 1;
        // In case of single shard, it can happen that fetch phase might execute before normalization phase. Moreover, The pagination logic
        // lies in the fetch phase.
        // If the fetch phase gets executed before the normalization phase, then the result will be not paginated as per normalized score.
        // Therefore, to avoid it we will update from value in search context to 0. This will stop fetch phase to trim results prematurely.
        // Later in the normalization phase we will update QuerySearchResult object with the right from value, to handle the effective
        // trimming of results.
        if (isSingleShard && searchContext.from() > 0) {
            searchContext.from(0);
        }

        HitsThresholdChecker hitsThreshold = new HitsThresholdChecker(Math.max(numDocs, searchContext.trackTotalHitsUpTo()));
        // HybridTopScoreDocCollector hybridCollector = new HybridTopScoreDocCollector(numDocs, hitsThreshold);
        // Collector hcollector = Objects.nonNull(filteringWeight) ? new FilteredCollector(hybridCollector, filteringWeight) :
        // hybridCollector;
        return new HybridCollectorContext(
            REASON_SEARCH_TOP_HITS,
            sort,
            numDocs,
            searchContext.trackTotalHitsUpTo(),
            hitsThreshold,
            filteringWeight,
            new TopDocsMerger(searchContext.sort()),
            searchContext.searchAfter(),
            searchContext
        );
    }

    private static int getSubqueryResultsRetrievalSize(final SearchContext searchContext) {
        HybridQuery hybridQuery = unwrapHybridQuery(searchContext);
        Integer paginationDepth = hybridQuery.getQueryContext().getPaginationDepth();

        // Pagination is expected to work only when pagination_depth is provided in the search request.
        if (Objects.isNull(paginationDepth) && searchContext.from() > 0) {
            throw new IllegalArgumentException(String.format(Locale.ROOT, "pagination_depth param is missing in the search request"));
        }

        if (Objects.nonNull(paginationDepth)) {
            return paginationDepth;
        }

        // Switch to from+size retrieval size during standard hybrid query execution where from is 0.
        return searchContext.size();
    }

    private static HybridQuery unwrapHybridQuery(final SearchContext searchContext) {
        HybridQuery hybridQuery;
        Query query = searchContext.query();
        // In case of nested fields and alias filter, hybrid query is wrapped under bool query and lies in the first clause.
        if (isHybridQueryWrappedInBooleanQuery(searchContext, searchContext.query())) {
            BooleanQuery booleanQuery = (BooleanQuery) query;
            hybridQuery = (HybridQuery) booleanQuery.clauses().get(0).query();
        } else {
            hybridQuery = (HybridQuery) query;
        }
        return hybridQuery;
    }

    public static boolean isHybridQueryWrappedInBooleanQuery(final SearchContext searchContext, final Query query) {
        return ((hasAliasFilter(query, searchContext) || hasNestedFieldOrNestedDocs(query, searchContext))
            && isWrappedHybridQuery(query)
            && !((BooleanQuery) query).clauses().isEmpty());
    }

    private static boolean hasNestedFieldOrNestedDocs(final Query query, final SearchContext searchContext) {
        return searchContext.mapperService().hasNested() && new NestedHelper(searchContext.mapperService()).mightMatchNestedDocs(query);
    }

    private static boolean isWrappedHybridQuery(final Query query) {
        return query instanceof BooleanQuery
            && ((BooleanQuery) query).clauses().stream().anyMatch(clauseQuery -> clauseQuery.query() instanceof HybridQuery);
    }

    private static boolean hasAliasFilter(final Query query, final SearchContext searchContext) {
        return Objects.nonNull(searchContext.aliasFilter());
    }

    HybridCollectorContext(
        String profileName,
        @Nullable SortAndFormats sortAndFormats,
        int numHits,
        int trackTotalHitsUpTo,
        HitsThresholdChecker hitsThresholdChecker,
        @Nullable Weight filterWeight,
        TopDocsMerger topDocsMerger,
        @Nullable FieldDoc after,
        SearchContext searchContext
    ) {
        super(profileName);
        this.sortAndFormats = sortAndFormats;
        this.numHits = numHits;
        this.trackTotalHitsUpTo = trackTotalHitsUpTo;
        this.hitsThresholdChecker = hitsThresholdChecker;
        this.filterWeight = filterWeight;
        this.topDocsMerger = topDocsMerger;
        this.after = after;
        this.searchContext = searchContext;
        HybridCollectorManager hybridCollectorManager = new HybridCollectorManager(
            numHits,
            hitsThresholdChecker,
            trackTotalHitsUpTo,
            sortAndFormats,
            filterWeight,
            topDocsMerger,
            after,
            searchContext
        );
        searchContext.queryCollectorManagers().put(HybridCollectorManager.class, hybridCollectorManager);

        this.manager = hybridCollectorManager;
        this.collector = manager.newCollector();
    }

    @Override
    CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
        this.manager = new HybridCollectorManager(
            numHits,
            hitsThresholdChecker,
            trackTotalHitsUpTo,
            sortAndFormats,
            filterWeight,
            topDocsMerger,
            after,
            searchContext
        );
        return manager;
    }
   
    private static void validateSortCriteria(SearchContext searchContext, boolean trackScores) {
        SortField[] sortFields = searchContext.sort().sort.getSort();
        boolean hasFieldSort = false;
        boolean hasScoreSort = false;
        for (SortField sortField : sortFields) {
            SortField.Type type = sortField.getType();
            if (type.equals(SortField.Type.SCORE)) {
                hasScoreSort = true;
            } else {
                hasFieldSort = true;
            }
            if (hasScoreSort && hasFieldSort) {
                break;
            }
        }
        if (hasScoreSort && hasFieldSort) {
            throw new IllegalArgumentException(
                "_score sort criteria cannot be applied with any other criteria. Please select one sort criteria out of them."
            );
        }
        if (trackScores && hasFieldSort) {
            throw new IllegalArgumentException(
                "Hybrid search results when sorted by any field, docId or _id, track_scores must be set to false."
            );
        }
        if (trackScores && hasScoreSort) {
            throw new IllegalArgumentException("Hybrid search results are by default sorted by _score, track_scores must be set to false.");
        }
    }

    public static boolean isHybridQuery(final Query query, final SearchContext searchContext) {
        if (query instanceof HybridQuery
            || (Objects.nonNull(searchContext.parsedQuery()) && searchContext.parsedQuery().query() instanceof HybridQuery)) {
            return true;
        }
        return false;
    }

    @Override
    Collector create(Collector in) throws IOException {
        return collector;
    }

    @Override
    void postProcess(QuerySearchResult result) throws IOException {
        CollectorManager<Collector, ReduceableSearchResult> collectorManager = (CollectorManager<
            Collector,
            ReduceableSearchResult>) searchContext.queryCollectorManagers().get(HybridCollectorManager.class);
        final Collection<Collector> subCollectors = new ArrayList<>();
        subCollectors.add(collector);
        try {
            collectorManager.reduce(subCollectors).reduce(result);
        } catch (IOException e) {
            throw new QueryPhaseExecutionException(searchContext.shardTarget(), "failed to execute hybrid query aggregation processor", e);
        }
    }

}