ESQL: Fix a bug in TOP

docs/changelog/121552.yaml

@@ -0,0 +1,5 @@
+pr: 121552
+summary: Fix a bug in TOP
+area: ES|QL
+type: bug
+issues: []

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/sort/BytesRefBucketedSort.java

@@ -143,13 +143,13 @@ public void collect(BytesRef value, int bucket) {
      */
     public void merge(int bucket, BytesRefBucketedSort other, int otherBucket) {
         long otherRootIndex = other.common.rootIndex(otherBucket);
-        if (otherRootIndex >= other.values.size()) {
+        long otherEnd = other.common.endIndex(otherRootIndex);
+        if (otherEnd >= other.values.size()) {
             // The value was never collected.
             return;
         }
         other.checkInvariant(otherBucket);
         long otherStart = other.startIndex(otherBucket, otherRootIndex);
-        long otherEnd = other.common.endIndex(otherRootIndex);
         // TODO: This can be improved for heapified buckets by making use of the heap structures
         for (long i = otherStart; i < otherEnd; i++) {
             collect(other.values.get(i).bytesRefView(), bucket);

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/sort/BucketedSortTestCase.java

@@ -409,6 +409,42 @@ public final void testMergeThisBigger() {
         }
     }
 
+    public final void testMergePastEnd() {
+        int buckets = 10000;
+        int bucketSize = between(1, 1000);
+        int target = between(0, buckets);
+        List<V> values = randomList(buckets, buckets, this::randomValue);
+        Collections.sort(values);
+        try (T sort = build(SortOrder.ASC, bucketSize)) {
+            // Add a single value to the main sort.
+            for (int b = 0; b < buckets; b++) {
+                collect(sort, values.get(b), b);
+            }
+
+            try (T other = build(SortOrder.ASC, bucketSize)) {
+                // Add *all* values to the target bucket of the secondary sort.
+                for (int i = 0; i < values.size(); i++) {
+                    if (i != target) {
+                        collect(other, values.get(i), target);
+                    }
+                }
+
+                // Merge all buckets pairwise. Most of the secondary ones are empty.
+                for (int b = 0; b < buckets; b++) {
+                    merge(sort, b, other, b);
+                }
+            }
+
+            for (int b = 0; b < buckets; b++) {
+                if (b == target) {
+                    assertBlock(sort, b, values.subList(0, bucketSize));
+                } else {
+                    assertBlock(sort, b, List.of(values.get(b)));
+                }
+            }
+        }
+    }
+
     protected void assertBlock(T sort, int groupId, List<V> values) {
         var blockFactory = TestBlockFactory.getNonBreakingInstance();