ESQL: Fix invariant test in TOP(bytes) (elastic#117049)

Fixes a self-test in the code for `TOP(bytes)`, specifically around the merging used for grouping by ordinals.
afoucret · Nov 20, 2024 · d5bc389 · d5bc389
1 parent 20166ae
commit d5bc389
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 1 deletion.
diff --git a/.../esql/compute/src/main/java/org/elasticsearch/compute/data/sort/BytesRefBucketedSort.java b/.../esql/compute/src/main/java/org/elasticsearch/compute/data/sort/BytesRefBucketedSort.java
@@ -147,7 +147,7 @@ public void merge(int bucket, BytesRefBucketedSort other, int otherBucket) {
             // The value was never collected.
             return;
         }
-        other.checkInvariant(bucket);
+        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

diff --git a/.../esql/compute/src/test/java/org/elasticsearch/compute/data/sort/BucketedSortTestCase.java b/.../esql/compute/src/test/java/org/elasticsearch/compute/data/sort/BucketedSortTestCase.java
@@ -363,6 +363,52 @@ public final void testMergeEmptyToEmpty() {
         }
     }
 
+    public final void testMergeOtherBigger() {
+        try (T sort = build(SortOrder.ASC, 3)) {
+            var values = threeSortedValues();
+
+            collect(sort, values.get(0), 0);
+            collect(sort, values.get(1), 0);
+            collect(sort, values.get(2), 0);
+
+            try (T other = build(SortOrder.ASC, 3)) {
+                collect(other, values.get(0), 0);
+                collect(other, values.get(1), 1);
+                collect(other, values.get(2), 2);
+
+                merge(sort, 0, other, 0);
+                merge(sort, 0, other, 1);
+                merge(sort, 0, other, 2);
+            }
+
+            assertBlock(sort, 0, List.of(values.get(0), values.get(0), values.get(1)));
+        }
+    }
+
+    public final void testMergeThisBigger() {
+        try (T sort = build(SortOrder.ASC, 3)) {
+            var values = threeSortedValues();
+
+            collect(sort, values.get(0), 0);
+            collect(sort, values.get(1), 1);
+            collect(sort, values.get(2), 2);
+
+            try (T other = build(SortOrder.ASC, 3)) {
+                collect(other, values.get(0), 0);
+                collect(other, values.get(1), 0);
+                collect(other, values.get(2), 0);
+
+                merge(sort, 0, other, 0);
+                merge(sort, 1, other, 0);
+                merge(sort, 2, other, 0);
+            }
+
+            assertBlock(sort, 0, List.of(values.get(0), values.get(0), values.get(1)));
+            assertBlock(sort, 1, List.of(values.get(0), values.get(1), values.get(1)));
+            assertBlock(sort, 2, values);
+        }
+    }
+
     protected void assertBlock(T sort, int groupId, List<V> values) {
         var blockFactory = TestBlockFactory.getNonBreakingInstance();