perf: optimize bucket counting and tracking in LongKeyedBucketOrds

[rajat315315]

Description

This PR optimizes LongKeyedBucketOrds.FromMany to replace $O(N)$ linear scans in bucketsInOrd and maxOwningBucketOrd with $O(1)$ operations:

1. Tracks maxOwningBucketOrd dynamically as a class field updated on insertions.
2. Caches bucket counts per owningBucketOrd in a LongArray managed by BigArrays, incrementing it incrementally.

Tests and Verification

• Unit Tests: Running ./gradlew :server:test --tests "org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrdsTests" passed successfully.
• Spotless Formatting: Verified with ./gradlew spotlessJavaCheck.
• Microbenchmarking: Microbenchmarks show a ~4x speedup in lookup operations (see the corresponding issue for details).

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

Related Issues

Resolves #22449

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

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.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit ccec60b)

Here are some key observations to aid the review process:

🧪 No relevant tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Resource Leak on Growth Failure In add(), bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1) is called after ords.add() has already inserted the entry. If grow throws (e.g., CircuitBreakingException), the ord has been added to ords but the count is not incremented, and maxOwningBucketOrd has already been updated. The caller may not see the returned ord, leading to inconsistent state between ords and bucketOrdsCounts. Consider growing the array before calling ords.add(), or reassigning bucketOrdsCounts only after the grow succeeds (the current assignment pattern does that, but the partially-updated maxOwningBucketOrd and the successful ords.add still leave state inconsistent on failure). long ord = ords.add(owningBucketOrd, value); if (ord >= 0) { maxOwningBucketOrd = Math.max(maxOwningBucketOrd, owningBucketOrd); bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); bucketOrdsCounts.set(owningBucketOrd, bucketOrdsCounts.get(owningBucketOrd) + 1); } return ord;

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to ccec60b
Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Track max owning bucket on all adds The LongLongHash.add returns a negative value (-1-ord) when the key already exists, meaning the bucket count is only incremented on the first insertion of each unique (owningBucketOrd, value) pair, which is correct. However, maxOwningBucketOrd should be updated regardless of whether the entry is new, since the same owningBucketOrd with a duplicate value is still a valid owning bucket. Move the maxOwningBucketOrd update outside the if (ord >= 0) block to ensure correctness. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [247-253] long ord = ords.add(owningBucketOrd, value); +maxOwningBucketOrd = Math.max(maxOwningBucketOrd, owningBucketOrd); if (ord >= 0) { - maxOwningBucketOrd = Math.max(maxOwningBucketOrd, owningBucketOrd); bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); bucketOrdsCounts.set(owningBucketOrd, bucketOrdsCounts.get(owningBucketOrd) + 1); } return ord; Suggestion importance[1-10]: 8 __ Why: Valid correctness concern: if an owningBucketOrd is only added with duplicate values (ord < 0), it would not be reflected in maxOwningBucketOrd, potentially causing incorrect results. Moving the update outside the if block ensures correctness.	Medium

---

Previous suggestions

Suggestions up to commit aba4ca9

Category	Suggestion	Impact
Possible issue	Prevent resource leak on constructor failure If LongLongHash construction succeeds but newLongArray throws (e.g., CircuitBreakingException), the just-allocated ords will leak. Wrap allocations to release ords on failure to prevent memory leaks under memory pressure. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [238-242] public FromMany(BigArrays bigArrays) { this.bigArrays = bigArrays; - ords = new LongLongHash(2, bigArrays); - bucketOrdsCounts = bigArrays.newLongArray(1, true); + LongLongHash hash = new LongLongHash(2, bigArrays); + boolean success = false; + try { + bucketOrdsCounts = bigArrays.newLongArray(1, true); + success = true; + } finally { + if (!success) { + hash.close(); + } + } + ords = hash; } Suggestion importance[1-10]: 6 __ Why: Valid concern: if newLongArray throws after LongLongHash allocation, ords would leak. This is a legitimate resource-safety improvement, though the failure scenario is uncommon.	Low
General	Handle negative input in bounds check Guard against negative owningBucketOrd values to avoid an ArrayIndexOutOfBoundsException from LongArray.get. The previous implementation iterated the hash and would naturally return 0 for negative inputs. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [267-272] @Override public long bucketsInOrd(long owningBucketOrd) { - if (owningBucketOrd >= bucketOrdsCounts.size()) { + if (owningBucketOrd < 0 || owningBucketOrd >= bucketOrdsCounts.size()) { return 0; } return bucketOrdsCounts.get(owningBucketOrd); } Suggestion importance[1-10]: 5 __ Why: Guarding against negative input preserves previous behavior and prevents potential exceptions, though it's unclear if negative values are ever passed in practice.	Low
	Use atomic increment for counter update The condition ord >= 0 skips updating counts and maxOwningBucketOrd when the entry already exists (negative ord). However, maxOwningBucketOrd should still be updated for existing entries too, because a caller may add the same (owningBucketOrd, value) pair with a larger owningBucketOrd never seen before as a new entry — but actually the issue is inverse: LongLongHash.add returns negative when the key already exists, meaning it was previously added and maxOwningBucketOrd was already updated then. That part is correct. But bucketOrdsCounts intentionally only increments on new insertions — that's also correct. However, this means maxOwningBucketOrd tracks max owning bucket even when only duplicates are added — verify semantics match previous behavior which counted only distinct entries in ords. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [247-253] long ord = ords.add(owningBucketOrd, value); if (ord >= 0) { maxOwningBucketOrd = Math.max(maxOwningBucketOrd, owningBucketOrd); bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); - bucketOrdsCounts.set(owningBucketOrd, bucketOrdsCounts.get(owningBucketOrd) + 1); + bucketOrdsCounts.increment(owningBucketOrd, 1); } return ord; Suggestion importance[1-10]: 3 __ Why: The suggestion is mostly a verification/discussion of semantics with a minor style change (using increment instead of get+set). The improvement is marginal and LongArray.increment may not exist with that exact signature.	Low

[github-actions]

Persistent review updated to latest commit ccec60b

[github-actions]

✅ Gradle check result for ccec60b: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 93.75000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 73.46%. Comparing base (27e5d58) to head (ccec60b).
⚠️ Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
...aggregations/bucket/terms/LongKeyedBucketOrds.java	93.75%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22450      +/-   ##
============================================
+ Coverage     73.39%   73.46%   +0.06%     
- Complexity    76397    76495      +98     
============================================
  Files          6105     6104       -1     
  Lines        346613   346566      -47     
  Branches      49888    49880       -8     
============================================
+ Hits         254403   254594     +191     
+ Misses        71958    71722     -236     
+ Partials      20252    20250       -2     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[msfroh]

Thanks @rajat315315 !

This looks good overall. I left some pretty minor cleanup comments.

[msfroh]

You can use bucketOrdsCounts.increment(owningBucketOrd, 1). It might help performance just a little bit, but also IMO it's more readable.

[msfroh]

I feel like we can slightly reduce the number of comparisons here (which normally I wouldn't worry about, but it is in the critical path, so it might actually be worth it).

What if we make this:

if (owningBucketOrd > maxOwningBucketOrd) {
  maxOwningBucketOrd = owningBucketOrd;
  bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1);
}

Basically, we get two comparisons for the price of one. What do you think?

[msfroh]

A couple of things here:

1. ords.close() should never throw an exception, so I don't believe the try is necessary.
2. bucketOrdsCounts should never be null (since it's allocated in the constructor). So I don't think we need the null check.