Move serialization logic into block

Author: JonasKunz

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

@@ -92,39 +92,42 @@ void loadValue(int valueIndex, CompressedExponentialHistogram resultHistogram, B
         }
     }
 
-    /**
-     * Returns the minimum for the histogram at the given value index.
-     * @param valueIndex the value index, obtained via {@link #getFirstValueIndex(int)} and {@link #getValueCount(int)}
-     * @return the minimum, or Double.NaN if the histogram is null or has no values
-     */
-    double getHistogramMin(int valueIndex) {
+    void serializeValue(int valueIndex, SerializedOutput out, BytesRef tempBytesRef) {
+        long valueCount = getHistogramValueCount(valueIndex);
+        out.appendLong(valueCount);
+        out.appendDouble(getHistogramSum(valueIndex));
+        out.appendDouble(getHistogramZeroThreshold(valueIndex));
+        if (valueCount > 0) {
+            // min / max are only non-null for non-empty histograms
+            out.appendDouble(getHistogramMin(valueIndex));
+            out.appendDouble(getHistogramMax(valueIndex));
+        }
+        out.appendBytesRef(getEncodedHistogramBytes(valueIndex, tempBytesRef));
+    }
+
+    private double getHistogramMin(int valueIndex) {
         int minimaValIndex = minima.getFirstValueIndex(valueIndex);
         return minima.isNull(minimaValIndex) ? Double.NaN : minima.getDouble(minimaValIndex);
     }
 
-    /**
-     * Returns the maximum for the histogram at the given value index.
-     * @param valueIndex the value index, obtained via {@link #getFirstValueIndex(int)} and {@link #getValueCount(int)}
-     * @return the maximum, or Double.NaN if the histogram is null or has no values
-     */
-    double getHistogramMax(int valueIndex) {
+    private double getHistogramMax(int valueIndex) {
         int maximaValIndex = maxima.getFirstValueIndex(valueIndex);
         return maxima.isNull(maximaValIndex) ? Double.NaN : maxima.getDouble(maximaValIndex);
     }
 
-    double getHistogramSum(int valueIndex) {
+    private double getHistogramSum(int valueIndex) {
         return sums.getDouble(sums.getFirstValueIndex(valueIndex));
     }
 
-    long getHistogramValueCount(int valueIndex) {
+    private long getHistogramValueCount(int valueIndex) {
         return valueCounts.getLong(valueCounts.getFirstValueIndex(valueIndex));
     }
 
-    double getHistogramZeroThreshold(int valueIndex) {
+    private double getHistogramZeroThreshold(int valueIndex) {
         return zeroThresholds.getDouble(zeroThresholds.getFirstValueIndex(valueIndex));
     }
 
-    BytesRef getEncodedHistogramBytes(int valueIndex, BytesRef tempBytesRef) {
+    private BytesRef getEncodedHistogramBytes(int valueIndex, BytesRef tempBytesRef) {
         return encodedHistograms.getBytesRef(encodedHistograms.getFirstValueIndex(valueIndex), tempBytesRef);
     }
 

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.compute.data;
 
+import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.exponentialhistogram.ExponentialHistogram;
 
 /**
@@ -15,6 +16,28 @@
  */
 public sealed interface ExponentialHistogramBlock extends Block permits ConstantNullBlock, ExponentialHistogramArrayBlock {
 
+    /**
+     * Abstraction to use for writing individual values via the {@link ExponentialHistogramBlockAccessor#serializeValue(int, SerializedOutput)}.
+     */
+    public interface SerializedOutput {
+        void appendDouble(double value);
+
+        void appendLong(long value);
+
+        void appendBytesRef(BytesRef bytesRef);
+    }
+
+    /**
+     * Abstraction to use for writing individual values via the {@link ExponentialHistogramBlockAccessor#serializeValue(int, SerializedOutput)}.
+     */
+    public interface SerializedInput {
+        double readDouble();
+
+        long readLong();
+
+        BytesRef readBytesRef(BytesRef tempBytesRef);
+    }
+
     static boolean equals(ExponentialHistogramBlock blockA, ExponentialHistogramBlock blockB) {
         if (blockA == blockB) {
             return true;

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

@@ -47,52 +47,20 @@ public ExponentialHistogram get(int valueIndex) {
         return reusedHistogram;
     }
 
-    public double getSum(int valueIndex) {
-        assert block.isNull(valueIndex) == false;
-        assert block.isReleased() == false;
-        ExponentialHistogramArrayBlock arrayBlock = (ExponentialHistogramArrayBlock) block;
-        return arrayBlock.getHistogramSum(valueIndex);
-    }
-
-    public long getValueCount(int valueIndex) {
-        assert block.isNull(valueIndex) == false;
-        assert block.isReleased() == false;
-        ExponentialHistogramArrayBlock arrayBlock = (ExponentialHistogramArrayBlock) block;
-        return arrayBlock.getHistogramValueCount(valueIndex);
-    }
-
-    public double getMin(int valueIndex) {
-        assert block.isNull(valueIndex) == false;
-        assert block.isReleased() == false;
-        ExponentialHistogramArrayBlock arrayBlock = (ExponentialHistogramArrayBlock) block;
-        return arrayBlock.getHistogramMin(valueIndex);
-    }
-
-    public double getMax(int valueIndex) {
-        assert block.isNull(valueIndex) == false;
-        assert block.isReleased() == false;
-        ExponentialHistogramArrayBlock arrayBlock = (ExponentialHistogramArrayBlock) block;
-        return arrayBlock.getHistogramMax(valueIndex);
-    }
-
-    public double getZeroThreshold(int valueIndex) {
-        assert block.isNull(valueIndex) == false;
-        assert block.isReleased() == false;
-        ExponentialHistogramArrayBlock arrayBlock = (ExponentialHistogramArrayBlock) block;
-        return arrayBlock.getHistogramZeroThreshold(valueIndex);
-    }
-
-    public BytesRef getEncodedHistogramBytes(int valueIndex) {
+    /**
+     * Encodes and appends a histogram value, so that it can be later deserialized
+     * via {@link ExponentialHistogramBlockBuilder#deserializeAndAppend(ExponentialHistogramBlock.SerializedInput)}.
+     *
+     * @param output the output to deserialize into
+     */
+    public void serializeValue(int valueIndex, ExponentialHistogramBlock.SerializedOutput output) {
         assert block.isNull(valueIndex) == false;
         assert block.isReleased() == false;
         ExponentialHistogramArrayBlock arrayBlock = (ExponentialHistogramArrayBlock) block;
         if (tempBytesRef == null) {
             tempBytesRef = new BytesRef();
         }
-        arrayBlock.getEncodedHistogramBytes(valueIndex, tempBytesRef);
-        return tempBytesRef;
+        arrayBlock.serializeValue(valueIndex, output, tempBytesRef);
     }
 
-
-
 }

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

@@ -25,6 +25,8 @@ public class ExponentialHistogramBlockBuilder implements Block.Builder {
     private final DoubleBlock.Builder zeroThresholdsBuilder;
     private final BytesRefBlock.Builder encodedHistogramsBuilder;
 
+    private final BytesRef scratch = new BytesRef();
+
     ExponentialHistogramBlockBuilder(int estimatedSize, BlockFactory blockFactory) {
         DoubleBlock.Builder minimaBuilder = null;
         DoubleBlock.Builder maximaBuilder = null;
@@ -101,23 +103,25 @@ public ExponentialHistogramBlockBuilder append(ExponentialHistogram histogram) {
         return this;
     }
 
-    public ExponentialHistogramBlockBuilder appendEncoded(double sum, long valueCount, double min, double max, double zeroThreshold, BytesRef encodedHistogram) {
-        sumsBuilder.appendDouble(sum);
+    /**
+     * Decodes and appends a value serialized with
+     *  {@link ExponentialHistogramBlockAccessor#serializeValue(int, ExponentialHistogramBlock.SerializedOutput)}.
+     *
+     * @param input the input to deserialize from
+     */
+    public void deserializeAndAppend(ExponentialHistogramBlock.SerializedInput input) {
+        long valueCount = input.readLong();
         valueCountsBuilder.appendLong(valueCount);
-        zeroThresholdsBuilder.appendDouble(zeroThreshold);
-        encodedHistogramsBuilder.appendBytesRef(encodedHistogram);
-        if (valueCount == 0) {
-            assert Double.isNaN(min) : "min should be NaN for empty histogram";
-            assert Double.isNaN(max) : "max should be NaN for empty histogram";
+        sumsBuilder.appendDouble(input.readDouble());
+        zeroThresholdsBuilder.appendDouble(input.readDouble());
+        if (valueCount > 0) {
+            minimaBuilder.appendDouble(input.readDouble());
+            maximaBuilder.appendDouble(input.readDouble());
+        } else {
             minimaBuilder.appendNull();
             maximaBuilder.appendNull();
-        } else {
-            assert Double.isNaN(min) == false : "min should be non-NaN for non-empty histogram";
-            assert Double.isNaN(max) == false : "max should be non-NaN for non-empty histogram";
-            minimaBuilder.appendDouble(min);
-            maximaBuilder.appendDouble(max);
         }
-        return this;
+        encodedHistogramsBuilder.appendBytesRef(input.readBytesRef(scratch));
     }
 
     @Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ResultBuilder.java

@@ -55,7 +55,7 @@ static ResultBuilder resultBuilderFor(
             case NULL -> new ResultBuilderForNull(blockFactory);
             case DOC -> new ResultBuilderForDoc(blockFactory, (DocVectorEncoder) encoder, positions);
             case AGGREGATE_METRIC_DOUBLE -> new ResultBuilderForAggregateMetricDouble(blockFactory, positions);
-            case EXPONENTIAL_HISTOGRAM ->  new ResultBuilderForExponentialHistogram(blockFactory, positions);
+            case EXPONENTIAL_HISTOGRAM -> new ResultBuilderForExponentialHistogram(blockFactory, positions);
             default -> {
                 assert false : "Result builder for [" + elementType + "]";
                 throw new UnsupportedOperationException("Result builder for [" + elementType + "]");

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ResultBuilderForExponentialHistogram.java

@@ -8,26 +8,23 @@
 package org.elasticsearch.compute.operator.topn;
 
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.compute.data.AggregateMetricDoubleBlockBuilder;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.ExponentialHistogramBlock;
 import org.elasticsearch.compute.data.ExponentialHistogramBlockBuilder;
-import org.elasticsearch.index.mapper.BlockLoader;
-
-import java.util.List;
 
 public class ResultBuilderForExponentialHistogram implements ResultBuilder {
 
     private final ExponentialHistogramBlockBuilder builder;
-    private final BytesRef scratch = new BytesRef();
+    private final ReusableTopNEncoderInput reusableInput = new ReusableTopNEncoderInput();
 
     ResultBuilderForExponentialHistogram(BlockFactory blockFactory, int positions) {
         this.builder = blockFactory.newExponentialHistogramBlockBuilder(positions);
     }
 
     @Override
     public void decodeKey(BytesRef keys) {
-        throw new AssertionError("AggregateMetricDoubleBlock can't be a key");
+        throw new AssertionError("ExponentialHistogramBlock can't be a key");
     }
 
     @Override
@@ -37,21 +34,9 @@ public void decodeValue(BytesRef values) {
             builder.appendNull();
             return;
         }
-        assert count == 1: "ExponentialHistogramBlock does not support multi values";
-        long histogramValueCount = TopNEncoder.DEFAULT_UNSORTABLE.decodeLong(values);
-        double sum = TopNEncoder.DEFAULT_UNSORTABLE.decodeDouble(values);
-        double min = TopNEncoder.DEFAULT_UNSORTABLE.decodeDouble(values);
-        double max = TopNEncoder.DEFAULT_UNSORTABLE.decodeDouble(values);
-        double zeroThreshold = TopNEncoder.DEFAULT_UNSORTABLE.decodeDouble(values);
-        BytesRef encodedHistogramBytes = TopNEncoder.DEFAULT_UNSORTABLE.decodeBytesRef(values, scratch);
-        builder.appendEncoded(
-            sum,
-            histogramValueCount,
-            min,
-            max,
-            zeroThreshold,
-            encodedHistogramBytes
-        );
+        assert count == 1 : "ExponentialHistogramBlock does not support multi values";
+        reusableInput.inputValues = values;
+        builder.deserializeAndAppend(reusableInput);
     }
 
     @Override
@@ -68,4 +53,23 @@ public String toString() {
     public void close() {
         builder.close();
     }
+
+    private class ReusableTopNEncoderInput implements ExponentialHistogramBlock.SerializedInput {
+        BytesRef inputValues;
+
+        @Override
+        public double readDouble() {
+            return TopNEncoder.DEFAULT_UNSORTABLE.decodeDouble(inputValues);
+        }
+
+        @Override
+        public long readLong() {
+            return TopNEncoder.DEFAULT_UNSORTABLE.decodeLong(inputValues);
+        }
+
+        @Override
+        public BytesRef readBytesRef(BytesRef tempBytesRef) {
+            return TopNEncoder.DEFAULT_UNSORTABLE.decodeBytesRef(inputValues, tempBytesRef);
+        }
+    }
 }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ValueExtractor.java

@@ -54,7 +54,7 @@ static ValueExtractor extractorFor(ElementType elementType, TopNEncoder encoder,
             case NULL -> new ValueExtractorForNull();
             case DOC -> new ValueExtractorForDoc(encoder, ((DocBlock) block).asVector());
             case AGGREGATE_METRIC_DOUBLE -> new ValueExtractorForAggregateMetricDouble(encoder, (AggregateMetricDoubleBlock) block);
-            case EXPONENTIAL_HISTOGRAM ->  new ValueExtractorForExponentialHistogram(encoder, (ExponentialHistogramBlock)  block);
+            case EXPONENTIAL_HISTOGRAM -> new ValueExtractorForExponentialHistogram(encoder, (ExponentialHistogramBlock) block);
             default -> {
                 assert false : "No value extractor for [" + block.elementType() + "]";
                 throw new UnsupportedOperationException("No value extractor for [" + block.elementType() + "]");

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ValueExtractorForExponentialHistogram.java

@@ -7,18 +7,15 @@
 
 package org.elasticsearch.compute.operator.topn;
 
-import org.elasticsearch.compute.data.AggregateMetricDoubleBlock;
-import org.elasticsearch.compute.data.DoubleBlock;
+import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.compute.data.ExponentialHistogramBlock;
 import org.elasticsearch.compute.data.ExponentialHistogramBlockAccessor;
-import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.operator.BreakingBytesRefBuilder;
 
-import java.util.List;
-
 public class ValueExtractorForExponentialHistogram implements ValueExtractor {
     private final ExponentialHistogramBlock block;
     private final ExponentialHistogramBlockAccessor accessor;
+    private final ReusableTopNEncoderOutput reusableOutput = new ReusableTopNEncoderOutput();
 
     ValueExtractorForExponentialHistogram(TopNEncoder encoder, ExponentialHistogramBlock block) {
         assert encoder == TopNEncoder.DEFAULT_UNSORTABLE;
@@ -32,20 +29,35 @@ public void writeValue(BreakingBytesRefBuilder values, int position) {
         if (block.isNull(position)) {
             TopNEncoder.DEFAULT_UNSORTABLE.encodeVInt(0, values);
         } else {
-            assert block.getValueCount(position) == 1: "Multi-valued ExponentialHistogram not supported in TopN";
+            assert block.getValueCount(position) == 1 : "Multi-valued ExponentialHistogram blcoks are not supported in TopN";
             TopNEncoder.DEFAULT_UNSORTABLE.encodeVInt(1, values);
             int valueIndex = block.getFirstValueIndex(position);
-            TopNEncoder.DEFAULT_UNSORTABLE.encodeLong(accessor.getValueCount(valueIndex), values);
-            TopNEncoder.DEFAULT_UNSORTABLE.encodeDouble(accessor.getSum(valueIndex), values);
-            TopNEncoder.DEFAULT_UNSORTABLE.encodeDouble(accessor.getMin(valueIndex), values);
-            TopNEncoder.DEFAULT_UNSORTABLE.encodeDouble(accessor.getMax(valueIndex), values);
-            TopNEncoder.DEFAULT_UNSORTABLE.encodeDouble(accessor.getZeroThreshold(valueIndex), values);
-            TopNEncoder.DEFAULT_UNSORTABLE.encodeBytesRef(accessor.getEncodedHistogramBytes(valueIndex), values);
+            reusableOutput.target = values;
+            accessor.serializeValue(valueIndex, reusableOutput);
         }
     }
 
     @Override
     public String toString() {
         return "ValueExtractorForExponentialHistogram";
     }
+
+    private static final class ReusableTopNEncoderOutput implements ExponentialHistogramBlock.SerializedOutput {
+        BreakingBytesRefBuilder target;
+
+        @Override
+        public void appendDouble(double value) {
+            TopNEncoder.DEFAULT_UNSORTABLE.encodeDouble(value, target);
+        }
+
+        @Override
+        public void appendLong(long value) {
+            TopNEncoder.DEFAULT_UNSORTABLE.encodeLong(value, target);
+        }
+
+        @Override
+        public void appendBytesRef(BytesRef value) {
+            TopNEncoder.DEFAULT_UNSORTABLE.encodeBytesRef(value, target);
+        }
+    }
 }