Implement parallel execution of sub-queries for hybrid search

[VijayanB]

Description

1. Add new thread pool to schedule tasks that are related to hybrid query execution
2. Register executor builders with Plugin
3. Use Lucene's Task Executor to execute and collect results
4. Parallelize Query re-write
5. Parallelize score supplier creation
6. Parallelize build hybrid scores

Issues Resolved

#279

Check List

• New functionality includes testing.
  • All tests pass
• New functionality has been documented.
  • New functionality has javadoc added
• Commits are signed as per the DCO using --signoff

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[codecov]

Codecov Report

Attention: Patch coverage is 83.78378% with 12 lines in your changes are missing coverage. Please review.

Please upload report for BASE (feature/parallelize-hybrid-search@806042c). Learn more about missing BASE report.

❗ Current head 2750ad4 differs from pull request most recent head bc6b885

Please upload reports for the commit bc6b885 to get more accurate results.

Files	Patch %	Lines
...ch/neuralsearch/executors/HybridQueryExecutor.java	44.44%	4 Missing and 1 partial ⚠️
...org/opensearch/neuralsearch/query/HybridQuery.java	89.47%	2 Missing ⚠️
...ensearch/neuralsearch/query/HybridQueryScorer.java	86.66%	2 Missing ⚠️
...ensearch/neuralsearch/query/HybridQueryWeight.java	85.71%	2 Missing ⚠️
...g/opensearch/neuralsearch/plugin/NeuralSearch.java	50.00%	1 Missing ⚠️

Additional details and impacted files

@@                         Coverage Diff                          @@
##             feature/parallelize-hybrid-search     #749   +/-   ##
====================================================================
  Coverage                                     ?   84.89%           
  Complexity                                   ?      812           
====================================================================
  Files                                        ?       65           
  Lines                                        ?     2490           
  Branches                                     ?      410           
====================================================================
  Hits                                         ?     2114           
  Misses                                       ?      213           
  Partials                                     ?      163           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[navneet1v]

Looked PR at high level and added comments. Please resolve them after that I will do 1 more review of the PR.

[navneet1v]

can we avoid storing queries in a map? and who is going to store this simple map entry?

[VijayanB]

HybridQueryExecutorCollector.collect() will store this entry. If we don't store map, we can't assemble back to caller once the execution is completed.

[shatejas]

You can see if you can use optional instead. That will simplify the data structure and make it more readable

return rewrittenQuery != query ? Optional.of(rewrittenQuery) : Optional.empty() 

[martin-gaievski]

do we have data or can we run benchmark to see how much overhead (if any) adds the Optional approach? I agree this way code looks simpler, but if we're aiming for performance improvement and this is critical section then we may want to save CPU cycles of wrapping/unwrapping result to Optional.

[martin-gaievski]

and +1 to Navneet's point regarding map and specifically using Query as a key. I was using similar approach in initial implementation of hybrid query, problem is that for some queries hashCode can be slow.
How much time this will be called - one time per sub-query x num of shards?

[martin-gaievski]

understood, I mismatch the Map.Entry<Query, Boolean> with Map<Query, Boolean>, I guess the Map.Entry is used to operate on pair of objects. I'm good now, thanks Vijay

[martin-gaievski]

Used optional to avoid null condition, but we need query irrespective of whether query is re-written or not.

We can explicitly check for null, my point was performance. Doing first Optional.of(rewrittenQuery), and later optional.get() both will add some cycles of wrapping query object into Optional and then unwrapping it. If we have done benchmarks with Optional and without it then we can use it as datapoint to make a decision.

[VijayanB]

Sounds good. I will run benchmark with and without before merging into main. In the meantime, can we keep Optional? What do you think?

[martin-gaievski]

Yes, I'm not sure what is the delta in performance if any, let's keep it with Optional, run benchmark and act per per data

[shatejas]

the intent was never to have an Optional as a member variable, the intent was to only return optional to indicate the value is not rewritten. It got misunderstood.

I realize that the query is needed irrespective of if its rewritten so Optional cannot be used. An Entry is a good option here.

In terms of performance, if Optional does turn out to be an issue then we might have to deep dive and see if wrapping results in an object is an issue overall

[shatejas]

Just for my understanding, why is a collector chosen here for collecting results? you could have plugged in TaskExecutor here and then simply use the list of results returned by invokeAll. The entire collector seems to be a round about way to get the results. Manager would have been useful if you wanted to return a different collector based on certain conditions IMO

[VijayanB]

Good question. The main idea is to move logic of collecting and reducing/merging into its own manager. By introducing TaskExecutor inside Scorer/Weight we are coupling how to parallelize with what needs to be parallelized. With introduction of Collector and Collector Manager, we decoupled parallelization from logic. IMO, this abstracts parallelization implementation from all of its caller.

[shatejas]

Lets dissect this a little more. There are 3 classes here:

1. HybridQueryExecutor: responsible for initializing and holding a common threadpool for query, scorer and weight,
2. HybridQueryExecutorCollectorManager: Responsible for giving a new collector
3. HybridQueryExecutorCollector : Holds the query results

With introduction of Collector and Collector Manager, we decoupled parallelization from logic. IMO, this abstracts parallelization implementation from all of its caller.

HybridQueryExecutor seems to be a good wrapper to make sure the same task executor is used. I am more curious about the additional value of HybridQueryExecutorCollectorManager and HybridQueryExecutorCollector

Why can't it be

final List<Callable<Entry<Query, Boolean>>> queryRewriteTasks = new ArrayList<>();
// for each subquery
queryRewriteTasks.add(() -> rewriteQuery(subQuery); //rewrite needs to return entry<query, boolean>;
List<Entry<Query, Boolean>> rewrittenQueries = HybridQueryExecutor.getExecutor().invokeAll(queryRewriteTasks);

// rest of the logic

This way you don't have to worry about the threadsafety of the collectors. There are some utility methods in collector manager which arent related to managing a collector but more of related to rewritten queries itself.

Overall while this works fine, I want to make sure these interfaces have a clear definition and its clear on how to use them. Currently it seems like we can simplify this by not having it unless I am misunderstanding the value it brings

[chishui]

Is HybridQueryExecutorCollector thread safe?

[VijayanB]

There are no setters and only way to update result is by calling collect method. Do you have any particular concerns?

[chishui]

Since collect() could set result, if collect() and getResults() are called asynchronously, technically, there could be race condition, right?

[VijayanB]

You are right. Currently getResult() is called from collector manager, where the contract is that it can be called only after all collectors are finished collection. This was taken care in the implementation. I can add this note in collector as well. i am reluctant to add synchronize to result variable, since it will add additional latency for scenario that is not possible at this moment. Do you have any other suggestions?

[shatejas]

Adding a comment is a good intention. you won't be able to enforce it. You are relying on the hope that the code doesn't change in future to hit that race condition so there will always be a risk.

If you really want to enforce; your options here would be to synchronize or move away from result collection to returning results as pointed out in this thread

[chishui]

i am reluctant to add synchronize to result variable, since it will add additional latency for scenario that is not possible at this moment

Actually, the overhead of synchronization is way less than we thought, here are some synchronization benchmark articles I found,

• https://baptiste-wicht.com/posts/2010/09/java-synchronization-mutual-exclusion-benchmark.html
• https://isuru-perera.blogspot.com/2016/05/benchmarking-java-locks-with-counters.html
  One operation of "synchronized" is < 1 micro second, considering the time spending on actual query logic with intensive I/O, this is negligible and you have simple implementation with concurrent situation covered.

[VijayanB]

@shatejas @chishui Fair enough. I am changing to synchronized, however while benchmarking i am planning to compare with and without synchronization to see that it is not adding any latency to this block.

[vibrantvarun]

+1 to @VijayanB comment. I think the purpose of this whole PR is to reduce latency. Therefore, I think @VijayanB it would good if you compare the benchmark results, if you see there is little degradation also in the latency, then we should take a call on this tradeoff whether to keep synchronized or not.

[martin-gaievski]

do we have data or can we run benchmark to see how much overhead (if any) adds the Optional approach? I agree this way code looks simpler, but if we're aiming for performance improvement and this is critical section then we may want to save CPU cycles of wrapping/unwrapping result to Optional.

[martin-gaievski]

and +1 to Navneet's point regarding map and specifically using Query as a key. I was using similar approach in initial implementation of hybrid query, problem is that for some queries hashCode can be slow.
How much time this will be called - one time per sub-query x num of shards?

[martin-gaievski]

Looks good to me, thank you Vijay.
Please plan to add unit tests for the new classes from the executor package. We can do it in a separate PR.