Fix incorrect exponential histogram sub-block serialization. (#137279)

Author: JonasKunz

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

@@ -291,12 +291,12 @@ public ExponentialHistogramArrayBlock deepCopy(BlockFactory blockFactory) {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        minima.writeTo(out);
-        maxima.writeTo(out);
-        sums.writeTo(out);
-        valueCounts.writeTo(out);
-        zeroThresholds.writeTo(out);
-        encodedHistograms.writeTo(out);
+        Block.writeTypedBlock(minima, out);
+        Block.writeTypedBlock(maxima, out);
+        Block.writeTypedBlock(sums, out);
+        Block.writeTypedBlock(valueCounts, out);
+        Block.writeTypedBlock(zeroThresholds, out);
+        Block.writeTypedBlock(encodedHistograms, out);
     }
 
     public static ExponentialHistogramArrayBlock readFrom(BlockStreamInput in) throws IOException {
@@ -309,12 +309,12 @@ public static ExponentialHistogramArrayBlock readFrom(BlockStreamInput in) throw
 
         boolean success = false;
         try {
-            minima = DoubleBlock.readFrom(in);
-            maxima = DoubleBlock.readFrom(in);
-            sums = DoubleBlock.readFrom(in);
-            valueCounts = LongBlock.readFrom(in);
-            zeroThresholds = DoubleBlock.readFrom(in);
-            encodedHistograms = BytesRefBlock.readFrom(in);
+            minima = (DoubleBlock) Block.readTypedBlock(in);
+            maxima = (DoubleBlock) Block.readTypedBlock(in);
+            sums = (DoubleBlock) Block.readTypedBlock(in);
+            valueCounts = (LongBlock) Block.readTypedBlock(in);
+            zeroThresholds = (DoubleBlock) Block.readTypedBlock(in);
+            encodedHistograms = (BytesRefBlock) Block.readTypedBlock(in);
             success = true;
         } finally {
             if (success == false) {

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/ExponentialHistogramBlockEqualityTests.java renamed to x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/ExponentialHistogramBlockTests.java

@@ -7,20 +7,65 @@
 
 package org.elasticsearch.compute.data;
 
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.compute.test.BlockTestUtils;
 import org.elasticsearch.compute.test.ComputeTestCase;
 import org.elasticsearch.compute.test.RandomBlock;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.exponentialhistogram.ExponentialHistogram;
 import org.elasticsearch.exponentialhistogram.ExponentialHistogramCircuitBreaker;
 
+import java.io.IOException;
 import java.util.List;
 
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.Matchers.equalTo;
 
-public class ExponentialHistogramBlockEqualityTests extends ComputeTestCase {
+public class ExponentialHistogramBlockTests extends ComputeTestCase {
 
-    public void testEmptyBlock() {
+    public void testPopulatedBlockSerialization() throws IOException {
+        int elementCount = randomIntBetween(1, 100);
+        ExponentialHistogramBlockBuilder builder = blockFactory().newExponentialHistogramBlockBuilder(elementCount);
+        for (int i = 0; i < elementCount; i++) {
+            if (randomBoolean()) {
+                builder.appendNull();
+            } else {
+                builder.append(BlockTestUtils.randomExponentialHistogram());
+            }
+        }
+        ExponentialHistogramBlock block = builder.build();
+        Block deserializedBlock = serializationRoundTrip(block);
+        assertThat(deserializedBlock, equalTo(block));
+        Releasables.close(block, deserializedBlock);
+    }
+
+    public void testNullSerialization() throws IOException {
+        // sub-blocks can be constant null, those should serialize correctly too
+        int elementCount = randomIntBetween(1, 100);
+
+        Block block = new ExponentialHistogramArrayBlock(
+            (DoubleBlock) blockFactory().newConstantNullBlock(elementCount),
+            (DoubleBlock) blockFactory().newConstantNullBlock(elementCount),
+            (DoubleBlock) blockFactory().newConstantNullBlock(elementCount),
+            (LongBlock) blockFactory().newConstantNullBlock(elementCount),
+            (DoubleBlock) blockFactory().newConstantNullBlock(elementCount),
+            (BytesRefBlock) blockFactory().newConstantNullBlock(elementCount)
+        );
+
+        Block deserializedBlock = serializationRoundTrip(block);
+        assertThat(deserializedBlock, equalTo(block));
+        Releasables.close(block, deserializedBlock);
+    }
+
+    private Block serializationRoundTrip(Block block) throws IOException {
+        BytesStreamOutput out = new BytesStreamOutput();
+        Block.writeTypedBlock(block, out);
+        try (BlockStreamInput input = new BlockStreamInput(out.bytes().streamInput(), blockFactory())) {
+            return Block.readTypedBlock(input);
+        }
+    }
+
+    public void testEmptyBlockEquality() {
         List<Block> blocks = List.of(
             blockFactory().newConstantNullBlock(0),
             blockFactory().newExponentialHistogramBlockBuilder(0).build(),

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/BlockTestUtils.java

@@ -374,7 +374,7 @@ public static Page convertBytesRefsToOrdinals(Page page) {
         }
     }
 
-    static ExponentialHistogram randomExponentialHistogram() {
+    public static ExponentialHistogram randomExponentialHistogram() {
         // TODO(b/133393): allow (index,scale) based zero thresholds as soon as we support them in the block
         // ideally Replace this with the shared random generation in ExponentialHistogramTestUtils
         boolean hasNegativeValues = randomBoolean();