Implement Indexed Priority Queue to improve grouping performance (#251)

[Dilaxn]

Is your feature request related to a problem?
Yes. The current grouping logic uses a standard priority queue, forcing removal operations to run in O(n)O(n). This creates a significant performance bottleneck—especially in scenarios with frequent updates or large data sets.

What solution would you like?
Introduce an Indexed Priority Queue in the query grouping logic. By allowing direct lookup of an element’s position in the heap, arbitrary removals and updates can be done in O(log⁡n) time instead of O(n). This significantly boosts performance for frequent insertions and removals.

What alternatives have you considered?

1.Binary Heap without indexing: Simpler to implement, but removal of an arbitrary element still requires an O(n) scan.

2.Fibonacci Heap or Pairing Heap: While they can provide even better theoretical performance for certain operations, they are more complex and less practical for our specific use case.

Additional context
I have reviewed the existing code and tested how the new Indexed Priority Queue integrates with the grouping logic. Early benchmarks show a noticeable reduction in latency for frequent queue updates. I plan to add thorough test coverage to validate correctness and performance improvements before merging.

[kkhatua]

@Dilaxn Thank you for taking interest in improving performance. Unfortunately we don't have any performance tests as yet to validate this. Would you be able to help share some numbers so as to showcase the performance gains? (cc: @deshsidd )

[ansjcy]

@Dilaxn Thanks for your work on this! I've added several comments regarding the generic type handling and thread safety.

@deshsidd, could you please also take a look to verify the functional correctness of the Indexed Priority Queue logic? Thanks!

[ansjcy]

We don't do deletion for data stored in the PriorityQueue in TopQueriesService, I think the performance of PriorityBlockingQueue and IndexedPriorityQueue would be the same here.

[Dilaxn]

We’re already handing the very same topQueriesStore into the MinMaxHeapQueryGrouper, so splitting it into two separate queues just adds needless complexity. By switching both service and grouper to use a single IndexedPriorityQueue, we get one shared heap—no extra merge logic, no change in performance characteristics—and all of our inserts, polls, and snapshots happen against the exact same data structure. For those reasons, I’d stick with the IPQ here.

[ansjcy]

In here, we use T as the map key, this will require T to have equals / hashCode implemented. Also if T is mutable (which is true for our case), the equals and hashCode will change after the data is inserted, in which case the indexMap will become inconsistent (e.g. contains and remove might fail). Also the order will depend on the comparator T implemented (this should be fine in our case though).

I think we should probably change IndexedPriorityQueue to be IndexedPriorityQueue<K, V>, and make index map: Map<K, Integer> indexMap; and heap be List<Entry<K, V>> heap . We should also pass a comparator in the constructor IMO.

[Dilaxn]

Keys are now of type K, separate from your mutable V values, so you only ever hash or compare on that stable identifier.

[ansjcy]

As I think about it more, using synchronized on each function is a good start to ensure thread safety, but it might be a bottleneck under high contention. I think it makes more sense to use more fine-grained locking mechanisms instead of doing synchronized for all public functions - the main challenge here should be ensuring data consistency between the heap and the indexMap (e.g. in siftUp, updating on the heap and indexMap should be atomic).

[Dilaxn]

Implemented separate locks for read and write operations.

[ansjcy]

Also the unit test testAddRecordToLimitAndDrain is failing, could you double check?
https://github.com/opensearch-project/query-insights/actions/runs/14431905665/job/40541189714?pr=299

You can run ./gradlew test on your local env to validate your changes as well.

[Dilaxn]

Thank you @ansjcy and @kkhatua for your feedback. I do check the review comments and try to solve those. Thank you.

[Dilaxn]

Hi @ansjcy I’ve updated the code per your review:
-Replaced T with Entry<K, V>.
-Refined the locking mechanism.
-Applied fixes to the unit tests.
Please verify the changes, Thank you for ur time and consideration.

[ansjcy]

As .kkhatua mentioned:

Would you be able to help share some numbers so as to showcase the performance gains?

We might need to do some benchmarking tests and CPU profiling to understand the performance gains. We need to:

• run an OpenSearch cluster
• use https://docs.opensearch.org/docs/latest/benchmark/ to ingest some workload.
• Then run some search requests with different shapes.
• use https://github.com/async-profiler/async-profiler to see the CPU usage impact numbers.

[deshsidd]

@Dilaxn Hope you are able to make progress on this PR. Please let us know if you need any guidance :)

[Dilaxn]

@deshsidd @ansjcy
Thank you for your patience. I’ve been spending additional time testing both functionality and performance to ensure the solution is solid. I’ll continue working on it over the weekend and will post an update soon.
Appreciate the support and guidance!

[Dilaxn]

@ansjcy @kkhatua
PQ.json
IPQ.json
Attached the benchmark results comparing the PriorityBlockingQueue (PQ) version of the Query Insights plugin with the updated IndexedPriorityQueue (IPQ) version.

Benchmark Setup:
Benchmark tool: OpenSearch Benchmark
Workload: nyc_taxis
Pipeline: benchmark-only
Environment: Local, external provisioned OpenSearch cluster (3.0.0-SNAPSHOT)
Metric source: JSON results from benchmark-results.json captured after each run
Plugin versions: plugin_version: pq vs. plugin_version: ipq (tagged via --user-tag)

Here's a summary of key performance metrics:

Metric	PQ	IPQ	Δ
Index throughput (mean)	50,856 docs/s	61,300 docs/s	↑ 20.5%
p90 latency	1,203 ms	804 ms	↓ 33%
p99 latency	5,929 ms	2,043 ms	↓ 65%
Indexing duration	3,532 s	2,825 s	↓ 20%
Distance agg service time (mean)	6,250 ms	11,173 ms	↑ 78%
Merge time	6,422 s	5,886 s	↓ 8%
Segment count	27	21	↓ 22%
Young GC time	167 s	145 s	↓ 13%

Key Takeaways:

The IPQ-based version significantly improves indexing throughput and reduces latency, especially at p99.

A slight increase in distance aggregation service time was observed 
Overall, the switch to IPQ results in lower GC pressure, faster merging, and fewer segments.

Please feel free to correct me if there’s any mistake. Thank you!

[Dilaxn]

Subject: ✅ Functional Parity & Concurrency Benchmark Validation for IPQ

Hi @deshsidd ,

I've completed a comprehensive validation of our IndexedPriorityQueue (IPQ) against the suggested IndexMinPQ implementation.
🔍 Functional Parity

✅ Verified 1:1 behavioral parity using:

    1,000 randomized functional test scenarios

    Property-based tests (250+ diverse operation sequences)

    Edge case coverage (insert, remove, peek, poll correctness)

🎯 Both queues produced identical sorted outputs and consistent state transitions.

⚙️ Concurrency Comparison

Benchmarked concurrent insertions (8 threads, 100,000 entries) against sequential IndexMinPQ.

Observation: IPQ performs slightly slower than the non-thread-safe version in microbenchmarks, primarily due to the overhead introduced by ReentrantReadWriteLock — which is expected for thread-safe implementations.

Advantage: IPQ is concurrency-safe by design and scales better under real-world multithreaded conditions (read/write workloads).

For this, I created a separate Gradle project with the help of GPT. The project is : https://github.com/Dilaxn/indexed-priority-queue-compat

Thanks again for the guidance!

Best regards,
Dilakshan

[ansjcy]

@Dilaxn This is very impressive, thanks for doing the performance benchmarking! Let me take a final check on the PR.
Also, since this change mainly affects search path, I'm wondering why are we seeing an improvement on indexing. Do you have more performance numbers on the search path?

[Dilaxn]

@Dilaxn This is very impressive, thanks for doing the performance benchmarking! Let me take a final check on the PR. Also, since this change mainly affects search path, I'm wondering why are we seeing an improvement on indexing. Do you have more performance numbers on the search path?

Hi @ansjcy, Thanks for the feedback!
I’ve attached the benchmark reports for the nyc_taxis workload using both IndexedPriorityQueue (the new change) and PriorityBlockingQueue (master).
IndexedPriorityQueue.json
master benchmark.txt
PriorityBlockingQueue.json

[deshsidd]

@Dilaxn Thanks for the parity checks and the benchmarks. This is pretty impressive! Let me take a final look today.

[deshsidd]

LGTM overall. Thanks again for the change. Some minor comments @Dilaxn
Once addressed we can aim to rebase approve and merge.

[ansjcy]

hmm, looking at the full benchmark run results I see latency increases on indexed PQ on almost all search workloads. I want us to do some extra benchmarks and cpu profilings. Let's hold on this PR for now until we fully understand the performance impact.