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 d0f6b04)

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 Incorrect grow condition bucketOrdsCounts is grown only when owningBucketOrd > maxOwningBucketOrd. However, the initial maxOwningBucketOrd is -1 while bucketOrdsCounts has size 1. If a caller inserts with owningBucketOrd = 0 first (common case), then later inserts with a larger ord that does not exceed the current array capacity but does exceed size 1, this works. But consider ords inserted in decreasing order: the first insert with owningBucketOrd = 5 grows the array to size 6, maxOwningBucketOrd = 5. A subsequent insert with owningBucketOrd = 3 will still call increment(3, 1) — this is fine. However, if owningBucketOrd equals the current maxOwningBucketOrd but exceeds bucketOrdsCounts.size() (which shouldn't happen given the invariant, but is worth verifying), no growth occurs. The invariant appears to hold, but any code path that inserts without going through add (or bypasses growth) could cause an out-of-bounds. Confirm the invariant bucketOrdsCounts.size() >= maxOwningBucketOrd + 1 is always maintained. if (ord >= 0) { if (owningBucketOrd > maxOwningBucketOrd) { maxOwningBucketOrd = owningBucketOrd; bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); } bucketOrdsCounts.increment(owningBucketOrd, 1); } Resource leak on constructor failure In the FromMany constructor, if bigArrays.newLongArray(1, true) throws after ords is successfully constructed, the ords LongLongHash will leak because there is no try/catch to release it. Consider wrapping the allocation to release ords on failure. public FromMany(BigArrays bigArrays) { this.bigArrays = bigArrays; ords = new LongLongHash(2, bigArrays); bucketOrdsCounts = bigArrays.newLongArray(1, true); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to d0f6b04

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Guard array growth by size, not max The counter should be incremented only on newly inserted ords, but LongLongHash.add returns -1-existingOrd for duplicates (negative), so ord >= 0 is correct for new entries. However, the array growth check only triggers when owningBucketOrd > maxOwningBucketOrd, which misses the case where a new (non-duplicate) key/value pair is inserted for a previously-seen owningBucketOrd but the array was never sized for it (e.g., if the first insertion updated maxOwningBucketOrd but not sufficiently). Ensure bucketOrdsCounts is grown whenever owningBucketOrd >= bucketOrdsCounts.size() to prevent IndexOutOfBoundsException on increment. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [250-258] long ord = ords.add(owningBucketOrd, value); if (ord >= 0) { if (owningBucketOrd > maxOwningBucketOrd) { maxOwningBucketOrd = owningBucketOrd; + } + if (owningBucketOrd >= bucketOrdsCounts.size()) { bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); } bucketOrdsCounts.increment(owningBucketOrd, 1); } return ord; Suggestion importance[1-10]: 3 __ Why: The suggestion's premise is largely incorrect: bigArrays.grow sizes the array to at least owningBucketOrd + 1, and this is triggered whenever owningBucketOrd > maxOwningBucketOrd, which correctly ensures the array is large enough. The suggested change is functionally equivalent in practice, offering minimal benefit.	Low

---

Previous suggestions

Suggestions up to commit 1e3b18d

Category	Suggestion	Impact
General	Guarantee both resources are closed If ords.close() throws, bucketOrdsCounts.close() will not be invoked, leaking memory. Use Releasables.close(ords, bucketOrdsCounts) to guarantee both are released. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [323-326] @Override public void close() { - ords.close(); - bucketOrdsCounts.close(); + org.opensearch.common.lease.Releasables.close(ords, bucketOrdsCounts); } Suggestion importance[1-10]: 5 __ Why: Using Releasables.close guards against a leak if ords.close() throws. It's a valid minor robustness improvement for resource cleanup.	Low
Possible issue	Ensure array grows for existing keys too LongLongHash.add returns a negative value (-1-ord) when the key already exists, meaning the current logic only increments the count for newly-added buckets and undercounts duplicates. The count should be incremented on every successful add call regardless of whether the key is new, since bucketsInOrd needs the total number of buckets (unique (owningBucketOrd, value) pairs) — but the array must still be grown to accommodate owningBucketOrd. Restructure so growth happens unconditionally on new ords and increment only occurs for new ords (ord >= 0), which matches the original semantics of counting distinct entries. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [250-258] long ord = ords.add(owningBucketOrd, value); +if (owningBucketOrd > maxOwningBucketOrd) { + maxOwningBucketOrd = owningBucketOrd; + bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); +} if (ord >= 0) { - if (owningBucketOrd > maxOwningBucketOrd) { - maxOwningBucketOrd = owningBucketOrd; - bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); - } bucketOrdsCounts.increment(owningBucketOrd, 1); } return ord; Suggestion importance[1-10]: 4 __ Why: The suggestion is partially valid: when owningBucketOrd is new but the value already exists in the hash, the array growth still happens because a new (owningBucketOrd, value) pair means ord >= 0. However, if the same owningBucketOrd reappears with an already-seen value, growth is not needed since the array was already sized. The current logic appears correct for counting distinct entries, and the suggestion's rewrite is functionally equivalent in behavior for the existing semantics. Impact is low.	Low

Suggestions up to commit 7f3a278

Category	Suggestion	Impact
Possible issue	Decouple array growth from max ord tracking The LongLongHash.add returns a negative value (-1-ord) when the key already exists, meaning the bucket count is only incremented for newly-inserted entries. However, bucketsInOrd is expected to return the number of distinct buckets (values) for an owningBucketOrd, which equals the count of insertions — so this is correct only if duplicates are never counted. Verify this matches the semantics of the previous implementation which counted all entries where getKey1 == owningBucketOrd (each entry is unique per key1+key2 pair, so semantics match). Also, ensure bucketOrdsCounts is grown even when owningBucketOrd equals maxOwningBucketOrd on the very first insert (initial size 1 covers ord 0, so OK), but consider growing whenever owningBucketOrd >= bucketOrdsCounts.size() rather than comparing to maxOwningBucketOrd for clarity and safety. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [250-258] long ord = ords.add(owningBucketOrd, value); if (ord >= 0) { + if (owningBucketOrd >= bucketOrdsCounts.size()) { + bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); + } if (owningBucketOrd > maxOwningBucketOrd) { maxOwningBucketOrd = owningBucketOrd; - bucketOrdsCounts = bigArrays.grow(bucketOrdsCounts, owningBucketOrd + 1); } bucketOrdsCounts.increment(owningBucketOrd, 1); } return ord; Suggestion importance[1-10]: 6 __ Why: Valid observation: coupling array growth to maxOwningBucketOrd is fragile. While the current code works because owningBucketOrd > maxOwningBucketOrd implies growing is needed, decoupling improves clarity and safety.	Low
General	Prevent resource leak on close failure If ords.close() throws, bucketOrdsCounts.close() will not be invoked and the LongArray will leak. Use Releasables.close(...) to ensure both are released even on exceptions. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [323-326] @Override public void close() { - ords.close(); - bucketOrdsCounts.close(); + org.opensearch.common.lease.Releasables.close(ords, bucketOrdsCounts); } Suggestion importance[1-10]: 6 __ Why: Correctly identifies that if ords.close() throws, bucketOrdsCounts will leak. Using Releasables.close() is a legitimate improvement for resource safety, though such exceptions are rare.	Low

Suggestions up to commit 43c8f88

Category	Suggestion	Impact
Possible issue	Track max owning ord for duplicates too The counter is only incremented when ord >= 0 (new insertion), but maxOwningBucketOrd should also be updated for duplicate insertions since an owning bucket may only ever be seen with values that already exist for other buckets. Move the maxOwningBucketOrd update outside the if block to correctly track the maximum owning bucket ordinal even when values are duplicates. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [251-257] 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.increment(owningBucketOrd, 1); } return ord; Suggestion importance[1-10]: 7 __ Why: Valid concern - maxOwningBucketOrd may not be updated if a value has already been added for a different owning bucket, since ords.add would return a negative ord. Moving the update outside the if block ensures correctness.	Medium
General	Ensure both resources close on exception If ords.close() throws, bucketOrdsCounts will leak. Use Releasables.close(...) (or a try/finally) to guarantee both resources are released even if one throws. server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java [322-325] @Override public void close() { - ords.close(); - bucketOrdsCounts.close(); + org.opensearch.common.lease.Releasables.close(ords, bucketOrdsCounts); } Suggestion importance[1-10]: 5 __ Why: Valid resource management improvement - using Releasables.close(...) prevents potential resource leak if ords.close() throws, though this is a minor concern in practice.	Low

Suggestions up to commit ccec60b

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

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

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.39%. Comparing base (27e5d58) to head (1e3b18d).
⚠️ Report is 14 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22450      +/-   ##
============================================
- Coverage     73.39%   73.39%   -0.01%     
- Complexity    76397    76431      +34     
============================================
  Files          6105     6104       -1     
  Lines        346613   346566      -47     
  Branches      49888    49880       -8     
============================================
- Hits         254403   254353      -50     
- Misses        71958    71987      +29     
+ Partials      20252    20226      -26     

☔ 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.

[github-actions]

Persistent review updated to latest commit 43c8f88

[github-actions]

❌ Gradle check result for 43c8f88: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit 7f3a278

[github-actions]

❌ Gradle check result for 7f3a278: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit 1e3b18d

[github-actions]

✅ Gradle check result for 1e3b18d: SUCCESS

[msfroh]

@rajat315315 -- one last thing. I noticed that your commits are signed off with:

Signed-off-by: Your Name <rajatjain.ix@gmail.com>

Unless I'm mistaken, our DCO rules want each commit signed using a real name. Can you please run git config --global user.name "Rajat Jain", then do a git commit --amend -s to update the signature on the last commit?

[rajat315315]

Done.

[github-actions]

Persistent review updated to latest commit d0f6b04

[github-actions]

✅ Gradle check result for d0f6b04: SUCCESS