Don't use a BytesStreamOutput to copy keys in BytesRefBlockHash (elas…

…tic#114819) (elastic#115288) Removes he size limit on BytesStreamOutput when copying keys.
jimczi · Oct 22, 2024 · 186a15b · 186a15b
1 parent 2203e4c
commit 186a15b
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 32 deletions.
diff --git a/docs/changelog/114819.yaml b/docs/changelog/114819.yaml
@@ -0,0 +1,6 @@
+pr: 114819
+summary: Don't use a `BytesStreamOutput` to copy keys in `BytesRefBlockHash`
+area: EQL
+type: bug
+issues:
+ - 114599
diff --git a/...main/generated-src/org/elasticsearch/compute/aggregation/blockhash/BytesRefBlockHash.java b/...main/generated-src/org/elasticsearch/compute/aggregation/blockhash/BytesRefBlockHash.java
diff --git a/...compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/X-BlockHash.java.st b/...compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/X-BlockHash.java.st
@@ -9,15 +9,10 @@ package org.elasticsearch.compute.aggregation.blockhash;
 
 $if(BytesRef)$
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.io.stream.BytesStreamOutput;
-import org.elasticsearch.common.io.stream.StreamInput;
 $endif$
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BitArray;
-$if(BytesRef)$
-import org.elasticsearch.common.util.BytesRefArray;
-$endif$
 import org.elasticsearch.common.util.$Hash$;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.aggregation.SeenGroupIds;
@@ -58,8 +53,6 @@ $endif$
 import org.elasticsearch.core.ReleasableIterator;
 
 $if(BytesRef)$
-import java.io.IOException;
-
 $else$
 import java.util.BitSet;
 
@@ -250,26 +243,21 @@ $if(BytesRef)$
          * without and still read from the block.
          */
         // TODO replace with takeBytesRefsOwnership ?!
+        final BytesRef spare = new BytesRef();
         if (seenNull) {
             try (var builder = blockFactory.newBytesRefBlockBuilder(Math.toIntExact(hash.size() + 1))) {
                 builder.appendNull();
-                BytesRef spare = new BytesRef();
                 for (long i = 0; i < hash.size(); i++) {
                     builder.appendBytesRef(hash.get(i, spare));
                 }
                 return new BytesRefBlock[] { builder.build() };
             }
         }
-
-        final int size = Math.toIntExact(hash.size());
-        try (BytesStreamOutput out = new BytesStreamOutput()) {
-            hash.getBytesRefs().writeTo(out);
-            try (StreamInput in = out.bytes().streamInput()) {
-                return new BytesRefBlock[] {
-                    blockFactory.newBytesRefArrayVector(new BytesRefArray(in, BigArrays.NON_RECYCLING_INSTANCE), size).asBlock() };
+        try (var builder = blockFactory.newBytesRefBlockBuilder(Math.toIntExact(hash.size()))) {
+            for (long i = 0; i < hash.size(); i++) {
+                builder.appendBytesRef(hash.get(i, spare));
             }
-        } catch (IOException e) {
-            throw new IllegalStateException(e);
+            return new BytesRefBlock[] { builder.build() };
         }
 $else$
         if (seenNull) {