Emit ordinals for BytesRefLongBlockHash (#136704)

This change tries to emit ordinals in BytesRefLongBlockHash when possible.

Author: dnhatn

server/src/main/java/org/elasticsearch/common/util/BytesRefArray.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.core.Releasables;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 
 /**
  * Compact serializable container for ByteRefs
@@ -136,6 +137,40 @@ public static BytesRefArray takeOwnershipOf(BytesRefArray other) {
         return b;
     }
 
+    /**
+     * Creates a deep copy of the given BytesRefArray.
+     */
+    public static BytesRefArray deepCopy(BytesRefArray other) {
+        LongArray startOffsets = null;
+        ByteArray bytes = null;
+        BytesRefArray result = null;
+        try {
+            startOffsets = other.bigArrays.newLongArray(other.startOffsets.size());
+            for (long i = 0; i < other.startOffsets.size(); i++) {
+                startOffsets.set(i, other.startOffsets.get(i));
+            }
+            bytes = other.bigArrays.newByteArray(other.bytes.size());
+            BytesRefIterator it = other.bytes.iterator();
+            BytesRef ref;
+            long pos = 0;
+            try {
+                while ((ref = it.next()) != null) {
+                    bytes.set(pos, ref.bytes, ref.offset, ref.length);
+                    pos += ref.length;
+                }
+            } catch (IOException e) {
+                assert false : new AssertionError("BytesRefIterator should not throw IOException", e);
+                throw new UncheckedIOException(e);
+            }
+            result = new BytesRefArray(startOffsets, bytes, other.size, other.bigArrays);
+            return result;
+        } finally {
+            if (result == null) {
+                Releasables.closeExpectNoException(startOffsets, bytes);
+            }
+        }
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         assert startOffsets != null : "using BytesRefArray after ownership taken";

server/src/test/java/org/elasticsearch/common/util/BytesRefArrayTests.java

@@ -164,6 +164,17 @@ public void testReadWrittenRepeated(boolean halfEmpty, int listSize) {
         testReadWritten(IntStream.range(0, listSize).mapToObj(i -> values).flatMap(List::stream).toList(), 10);
     }
 
+    public void testDeepCopy() {
+        try (
+            BytesRefArray bytes1 = randomArray(between(1, 50_000), between(10, 50_000), mockBigArrays());
+            BytesRefArray bytes2 = BytesRefArray.deepCopy(bytes1);
+            BytesRefArray bytes3 = BytesRefArray.deepCopy(bytes2)
+        ) {
+            assertEquality(bytes1, bytes2);
+            assertEquality(bytes1, bytes3);
+        }
+    }
+
     private void testReadWritten(List<BytesRef> values, int initialCapacity) {
         try (BytesRefArray array = new BytesRefArray(initialCapacity, mockBigArrays())) {
             for (BytesRef v : values) {
@@ -191,7 +202,7 @@ private void assertEquality(BytesRefArray original, BytesRefArray copy) {
         for (int i = 0; i < original.size(); ++i) {
             original.get(i, scratch);
             copy.get(i, scratch2);
-            assertEquals(scratch, scratch2);
+            assertEquals(Integer.toString(i), scratch, scratch2);
         }
     }
 }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/BytesRefLongBlockHash.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BitArray;
+import org.elasticsearch.common.util.BytesRefArray;
 import org.elasticsearch.common.util.LongLongHash;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.aggregation.SeenGroupIds;
@@ -22,6 +23,7 @@
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.compute.data.LongVector;
+import org.elasticsearch.compute.data.OrdinalBytesRefBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.mvdedupe.IntLongBlockAdd;
 import org.elasticsearch.core.ReleasableIterator;
@@ -143,28 +145,56 @@ public ReleasableIterator<IntBlock> lookup(Page page, ByteSizeValue targetBlockS
 
     @Override
     public Block[] getKeys() {
-        int positions = (int) finalHash.size();
         BytesRefBlock k1 = null;
         LongVector k2 = null;
-        try (
-            BytesRefBlock.Builder keys1 = blockFactory.newBytesRefBlockBuilder(positions);
-            LongVector.Builder keys2 = blockFactory.newLongVectorBuilder(positions)
-        ) {
-            BytesRef scratch = new BytesRef();
-            for (long i = 0; i < positions; i++) {
-                keys2.appendLong(finalHash.getKey2(i));
-                long h1 = finalHash.getKey1(i);
-                if (h1 == 0) {
-                    keys1.appendNull();
-                } else {
-                    keys1.appendBytesRef(bytesHash.hash.get(h1 - 1, scratch));
+        int positions = (int) finalHash.size();
+        if (OrdinalBytesRefBlock.isDense(positions, bytesHash.hash.size())) {
+            try (var ordinals = blockFactory.newIntBlockBuilder(positions); var longs = blockFactory.newLongVectorBuilder(positions)) {
+                for (long p = 0; p < positions; p++) {
+                    long h1 = finalHash.getKey1(p);
+                    if (h1 == 0) {
+                        ordinals.appendNull();
+                    } else {
+                        ordinals.appendInt(Math.toIntExact(h1 - 1));
+                    }
+                    longs.appendLong(finalHash.getKey2(p));
+                }
+                // TODO: make takeOwnershipOf work?
+                BytesRefArray bytes = BytesRefArray.deepCopy(bytesHash.hash.getBytesRefs());
+                try {
+                    var dict = blockFactory.newBytesRefArrayVector(bytes, Math.toIntExact(bytes.size()));
+                    bytes = null; // transfer ownership to dict
+                    k1 = new OrdinalBytesRefBlock(ordinals.build(), dict);
+                } finally {
+                    Releasables.closeExpectNoException(bytes);
+                }
+                k2 = longs.build();
+            } finally {
+                if (k2 == null) {
+                    Releasables.closeExpectNoException(k1);
                 }
             }
-            k1 = keys1.build();
-            k2 = keys2.build();
-        } finally {
-            if (k2 == null) {
-                Releasables.closeExpectNoException(k1);
+        } else {
+            try (
+                BytesRefBlock.Builder keys1 = blockFactory.newBytesRefBlockBuilder(positions);
+                LongVector.Builder keys2 = blockFactory.newLongVectorBuilder(positions)
+            ) {
+                BytesRef scratch = new BytesRef();
+                for (long i = 0; i < positions; i++) {
+                    long h1 = finalHash.getKey1(i);
+                    if (h1 == 0) {
+                        keys1.appendNull();
+                    } else {
+                        keys1.appendBytesRef(bytesHash.hash.get(h1 - 1, scratch));
+                    }
+                    keys2.appendLong(finalHash.getKey2(i));
+                }
+                k1 = keys1.build();
+                k2 = keys2.build();
+            } finally {
+                if (k2 == null) {
+                    Releasables.closeExpectNoException(k1);
+                }
             }
         }
         if (reverseOutput) {