perf: optimize bucket counting and tracking in LongKeyedBucketOrds by maintaining per-bucket counts and max ordinal metadata.

Author: rajat315315

benchmarks/src/main/java/org/opensearch/benchmark/search/aggregations/bucket/terms/LongKeyedBucketOrdsBenchmark.java

@@ -182,4 +182,17 @@ public void multiBucket(Blackhole bh) {
             bh.consume(ords);
         }
     }
+
+    @Benchmark
+    public void benchmarkBucketsInOrdAndMaxOwning(Blackhole bh) {
+        try (LongKeyedBucketOrds ords = LongKeyedBucketOrds.build(bigArrays, CardinalityUpperBound.MANY)) {
+            for (long i = 0; i < 50_000; i++) {
+                ords.add(i % 100, i % DISTINCT_VALUES);
+            }
+            for (long j = 0; j < 10_000; j++) {
+                bh.consume(ords.bucketsInOrd(j % 100));
+                bh.consume(ords.maxOwningBucketOrd());
+            }
+        }
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongKeyedBucketOrds.java

@@ -34,6 +34,7 @@
 
 import org.opensearch.common.lease.Releasable;
 import org.opensearch.common.util.BigArrays;
+import org.opensearch.common.util.LongArray;
 import org.opensearch.common.util.LongLongHash;
 import org.opensearch.common.util.ReorganizingLongHash;
 import org.opensearch.search.aggregations.CardinalityUpperBound;
@@ -229,16 +230,27 @@ public void close() {
      * @opensearch.internal
      */
     public static class FromMany extends LongKeyedBucketOrds {
+        private final BigArrays bigArrays;
         private final LongLongHash ords;
+        private long maxOwningBucketOrd = -1;
+        private LongArray bucketOrdsCounts;
 
         public FromMany(BigArrays bigArrays) {
+            this.bigArrays = bigArrays;
             ords = new LongLongHash(2, bigArrays);
+            bucketOrdsCounts = bigArrays.newLongArray(1, true);
         }
 
         @Override
         public long add(long owningBucketOrd, long value) {
             // This is in the critical path for collecting most aggs. Be careful of performance.
-            return ords.add(owningBucketOrd, value);
+            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;
         }
 
         @Override
@@ -253,14 +265,10 @@ public long get(long ordinal) {
 
         @Override
         public long bucketsInOrd(long owningBucketOrd) {
-            // TODO it'd be faster to count the number of buckets in a list of these ords rather than one at a time
-            long count = 0;
-            for (long i = 0; i < ords.size(); i++) {
-                if (ords.getKey1(i) == owningBucketOrd) {
-                    count++;
-                }
+            if (owningBucketOrd >= bucketOrdsCounts.size()) {
+                return 0;
             }
-            return count;
+            return bucketOrdsCounts.get(owningBucketOrd);
         }
 
         @Override
@@ -270,12 +278,7 @@ public long size() {
 
         @Override
         public long maxOwningBucketOrd() {
-            // TODO this is fairly expensive to compute. Can we avoid needing it?
-            long max = -1;
-            for (long i = 0; i < ords.size(); i++) {
-                max = Math.max(max, ords.getKey1(i));
-            }
-            return max;
+            return maxOwningBucketOrd;
         }
 
         @Override
@@ -313,7 +316,13 @@ public long ord() {
 
         @Override
         public void close() {
-            ords.close();
+            try {
+                ords.close();
+            } finally {
+                if (bucketOrdsCounts != null) {
+                    bucketOrdsCounts.close();
+                }
+            }
         }
     }
 }