ES|QL block type for exponential histograms (#133393)

Author: JonasKunz

libs/exponential-histogram/src/main/java/org/elasticsearch/exponentialhistogram/ExponentialHistogramXContent.java

@@ -21,9 +21,16 @@
 
 package org.elasticsearch.exponentialhistogram;
 
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.Types;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiConsumer;
 
 /**
  * Handles the serialization of an {@link ExponentialHistogram} to XContent.
@@ -48,7 +55,11 @@ public class ExponentialHistogramXContent {
      * @param histogram the ExponentialHistogram to serialize
      * @throws IOException if the XContentBuilder throws an IOException
      */
-    public static void serialize(XContentBuilder builder, ExponentialHistogram histogram) throws IOException {
+    public static void serialize(XContentBuilder builder, @Nullable ExponentialHistogram histogram) throws IOException {
+        if (histogram == null) {
+            builder.nullValue();
+            return;
+        }
         builder.startObject();
 
         builder.field(SCALE_FIELD, histogram.scale());
@@ -101,4 +112,63 @@ private static void writeBuckets(XContentBuilder b, String fieldName, Exponentia
         b.endObject();
     }
 
+    /**
+     * Parses an {@link ExponentialHistogram} from the provided {@link XContentParser}.
+     * This method is neither optimized, nor does it do any validation of the parsed content.
+     * No estimation for missing sum/min/max is done.
+     * Therefore only intended for testing!
+     *
+     * @param xContent the serialized histogram to read
+     * @return the deserialized histogram
+     * @throws IOException if the XContentParser throws an IOException
+     */
+    public static ExponentialHistogram parseForTesting(XContentParser xContent) throws IOException {
+        if (xContent.currentToken() == null) {
+            xContent.nextToken();
+        }
+        if (xContent.currentToken() == XContentParser.Token.VALUE_NULL) {
+            return null;
+        }
+        return parseForTesting(xContent.map());
+    }
+
+    /**
+     * Parses an {@link ExponentialHistogram} from a {@link Map}.
+     * This method is neither optimized, nor does it do any validation of the parsed content.
+     * No estimation for missing sum/min/max is done.
+     * Therefore only intended for testing!
+     *
+     * @param xContent the serialized histogram as a map
+     * @return the deserialized histogram
+     */
+    public static ExponentialHistogram parseForTesting(@Nullable Map<String, Object> xContent) {
+        if (xContent == null) {
+            return null;
+        }
+        int scale = ((Number) xContent.get(SCALE_FIELD)).intValue();
+        ExponentialHistogramBuilder builder = ExponentialHistogram.builder(scale, ExponentialHistogramCircuitBreaker.noop());
+
+        Map<String, Number> zero = Types.forciblyCast(xContent.getOrDefault(ZERO_FIELD, Collections.emptyMap()));
+        double zeroThreshold = zero.getOrDefault(ZERO_THRESHOLD_FIELD, 0).doubleValue();
+        long zeroCount = zero.getOrDefault(ZERO_COUNT_FIELD, 0).longValue();
+        builder.zeroBucket(ZeroBucket.create(zeroThreshold, zeroCount));
+
+        builder.sum(((Number) xContent.getOrDefault(SUM_FIELD, 0)).doubleValue());
+        builder.min(((Number) xContent.getOrDefault(MIN_FIELD, Double.NaN)).doubleValue());
+        builder.max(((Number) xContent.getOrDefault(MAX_FIELD, Double.NaN)).doubleValue());
+
+        parseBuckets(Types.forciblyCast(xContent.getOrDefault(NEGATIVE_FIELD, Collections.emptyMap())), builder::setNegativeBucket);
+        parseBuckets(Types.forciblyCast(xContent.getOrDefault(POSITIVE_FIELD, Collections.emptyMap())), builder::setPositiveBucket);
+
+        return builder.build();
+    }
+
+    private static void parseBuckets(Map<String, List<Number>> serializedBuckets, BiConsumer<Long, Long> bucketSetter) {
+        List<Number> indices = serializedBuckets.getOrDefault(BUCKET_INDICES_FIELD, Collections.emptyList());
+        List<Number> counts = serializedBuckets.getOrDefault(BUCKET_COUNTS_FIELD, Collections.emptyList());
+        assert indices.size() == counts.size();
+        for (int i = 0; i < indices.size(); i++) {
+            bucketSetter.accept(indices.get(i).longValue(), counts.get(i).longValue());
+        }
+    }
 }

libs/exponential-histogram/src/test/java/org/elasticsearch/exponentialhistogram/ExponentialHistogramXContentTests.java

@@ -23,6 +23,8 @@
 
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
 import org.elasticsearch.xcontent.json.JsonXContent;
 
 import java.io.IOException;
@@ -31,6 +33,11 @@
 
 public class ExponentialHistogramXContentTests extends ExponentialHistogramTestCase {
 
+    public void testNullHistogram() {
+        assertThat(toJson(null), equalTo("null"));
+        checkRoundTrip(null);
+    }
+
     public void testEmptyHistogram() {
         ExponentialHistogram emptyHistogram = ExponentialHistogram.empty();
         assertThat(toJson(emptyHistogram), equalTo("{\"scale\":" + emptyHistogram.scale() + ",\"sum\":0.0}"));
@@ -62,18 +69,21 @@ public void testFullHistogram() {
                     + "}"
             )
         );
+        checkRoundTrip(histo);
     }
 
     public void testOnlyZeroThreshold() {
         ExponentialHistogram histo = createAutoReleasedHistogram(b -> b.scale(3).sum(1.1).zeroBucket(ZeroBucket.create(5.0, 0)));
         assertThat(toJson(histo), equalTo("{\"scale\":3,\"sum\":1.1,\"zero\":{\"threshold\":5.0}}"));
+        checkRoundTrip(histo);
     }
 
     public void testOnlyZeroCount() {
         ExponentialHistogram histo = createAutoReleasedHistogram(
             b -> b.zeroBucket(ZeroBucket.create(0.0, 7)).scale(2).sum(1.1).min(0).max(0)
         );
         assertThat(toJson(histo), equalTo("{\"scale\":2,\"sum\":1.1,\"min\":0.0,\"max\":0.0,\"zero\":{\"count\":7}}"));
+        checkRoundTrip(histo);
     }
 
     public void testOnlyPositiveBuckets() {
@@ -84,6 +94,7 @@ public void testOnlyPositiveBuckets() {
             toJson(histo),
             equalTo("{\"scale\":4,\"sum\":1.1,\"min\":0.5,\"max\":2.5,\"positive\":{\"indices\":[-1,2],\"counts\":[3,5]}}")
         );
+        checkRoundTrip(histo);
     }
 
     public void testOnlyNegativeBuckets() {
@@ -94,6 +105,7 @@ public void testOnlyNegativeBuckets() {
             toJson(histo),
             equalTo("{\"scale\":5,\"sum\":1.1,\"min\":-0.5,\"max\":-0.25,\"negative\":{\"indices\":[-1,2],\"counts\":[4,6]}}")
         );
+        checkRoundTrip(histo);
     }
 
     private static String toJson(ExponentialHistogram histo) {
@@ -105,4 +117,17 @@ private static String toJson(ExponentialHistogram histo) {
         }
     }
 
+    private static void checkRoundTrip(ExponentialHistogram histo) {
+        try (XContentBuilder builder = JsonXContent.contentBuilder()) {
+            ExponentialHistogramXContent.serialize(builder, histo);
+            String json = Strings.toString(builder);
+            try (XContentParser parser = JsonXContent.jsonXContent.createParser(XContentParserConfiguration.EMPTY, json)) {
+                ExponentialHistogram parsed = ExponentialHistogramXContent.parseForTesting(parser);
+                assertThat(parsed, equalTo(histo));
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
 }

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/plugin/EsqlCorePlugin.java

@@ -7,9 +7,11 @@
 
 package org.elasticsearch.xpack.esql.core.plugin;
 
+import org.elasticsearch.common.util.FeatureFlag;
 import org.elasticsearch.plugins.ExtensiblePlugin;
 import org.elasticsearch.plugins.Plugin;
 
 public class EsqlCorePlugin extends Plugin implements ExtensiblePlugin {
 
+    public static final FeatureFlag EXPONENTIAL_HISTOGRAM_FEATURE_FLAG = new FeatureFlag("esql_exponential_histogram");
 }

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DataType.java

@@ -340,6 +340,14 @@ public enum DataType implements Writeable {
             .estimatedSize(Double.BYTES * 3 + Integer.BYTES)
             .supportedSince(DataTypesTransportVersions.ESQL_AGGREGATE_METRIC_DOUBLE_CREATED_VERSION)
     ),
+
+    EXPONENTIAL_HISTOGRAM(
+        builder().esType("exponential_histogram")
+            .estimatedSize(16 * 160)// guess 160 buckets (OTEL default for positive values only histograms) with 16 bytes per bucket
+            .docValues()
+            .underConstruction()
+    ),
+
     /**
      * Fields with this type are dense vectors, represented as an array of float values.
      */

x-pack/plugin/esql/build.gradle

@@ -42,6 +42,7 @@ dependencies {
   implementation project('compute:ann')
   implementation project(':libs:dissect')
   implementation project(':libs:grok')
+  implementation project(':libs:exponential-histogram')
   api "org.apache.lucene:lucene-spatial3d:${versions.lucene}"
   api project(":libs:h3")
   implementation project('arrow')

x-pack/plugin/esql/compute/build.gradle

@@ -16,6 +16,7 @@ dependencies {
   compileOnly project(xpackModule('ml'))
   annotationProcessor project('gen')
   implementation 'com.carrotsearch:hppc:0.8.1'
+  api project(':libs:exponential-histogram')
 
   testImplementation(project(':modules:analysis-common'))
   testImplementation(project(':test:framework'))

x-pack/plugin/esql/compute/src/main/java/module-info.java

@@ -21,6 +21,7 @@
     requires org.elasticsearch.geo;
     requires org.elasticsearch.xcore;
     requires hppc;
+    requires org.elasticsearch.exponentialhistogram;
 
     exports org.elasticsearch.compute;
     exports org.elasticsearch.compute.aggregation;

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
 import org.elasticsearch.compute.data.Block.MvOrdering;
+import org.elasticsearch.exponentialhistogram.ExponentialHistogram;
 
 import java.util.BitSet;
 
@@ -467,6 +468,19 @@ public final AggregateMetricDoubleBlock newConstantAggregateMetricDoubleBlock(
         }
     }
 
+    public ExponentialHistogramBlockBuilder newExponentialHistogramBlockBuilder(int estimatedSize) {
+        return new ExponentialHistogramBlockBuilder(estimatedSize, this);
+    }
+
+    public final ExponentialHistogramBlock newConstantExponentialHistogramBlock(ExponentialHistogram value, int positionCount) {
+        try (ExponentialHistogramBlockBuilder builder = newExponentialHistogramBlockBuilder(positionCount)) {
+            for (int i = 0; i < positionCount; i++) {
+                builder.append(value);
+            }
+            return builder.build();
+        }
+    }
+
     public final AggregateMetricDoubleBlock newAggregateMetricDoubleBlock(
         double[] minValues,
         double[] maxValues,

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

@@ -13,6 +13,8 @@
 import org.elasticsearch.compute.data.AggregateMetricDoubleBlockBuilder.AggregateMetricDoubleLiteral;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
+import org.elasticsearch.exponentialhistogram.ExponentialHistogram;
+import org.elasticsearch.exponentialhistogram.ExponentialHistogramCircuitBreaker;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -224,6 +226,7 @@ public static void appendValue(Block.Builder builder, Object val, ElementType ty
             case DOUBLE -> ((DoubleBlock.Builder) builder).appendDouble((Double) val);
             case BOOLEAN -> ((BooleanBlock.Builder) builder).appendBoolean((Boolean) val);
             case AGGREGATE_METRIC_DOUBLE -> ((AggregateMetricDoubleBlockBuilder) builder).appendLiteral((AggregateMetricDoubleLiteral) val);
+            case EXPONENTIAL_HISTOGRAM -> ((ExponentialHistogramBlockBuilder) builder).append((ExponentialHistogram) val);
             default -> throw new UnsupportedOperationException("unsupported element type [" + type + "]");
         }
     }
@@ -253,6 +256,7 @@ private static Block constantBlock(BlockFactory blockFactory, ElementType type,
             case BOOLEAN -> blockFactory.newConstantBooleanBlockWith((boolean) val, size);
             case AGGREGATE_METRIC_DOUBLE -> blockFactory.newConstantAggregateMetricDoubleBlock((AggregateMetricDoubleLiteral) val, size);
             case FLOAT -> blockFactory.newConstantFloatBlockWith((float) val, size);
+            case EXPONENTIAL_HISTOGRAM -> blockFactory.newConstantExponentialHistogramBlock((ExponentialHistogram) val, size);
             default -> throw new UnsupportedOperationException("unsupported element type [" + type + "]");
         };
     }
@@ -306,6 +310,12 @@ yield new AggregateMetricDoubleLiteral(
                     aggBlock.countBlock().getInt(offset)
                 );
             }
+            case EXPONENTIAL_HISTOGRAM -> {
+                ExponentialHistogramBlock histoBlock = (ExponentialHistogramBlock) block;
+                ExponentialHistogram histogram = new ExponentialHistogramBlockAccessor(histoBlock).get(offset);
+                // return a copy so that the returned value is not bound to the lifetime of the block
+                yield ExponentialHistogram.builder(histogram, ExponentialHistogramCircuitBreaker.noop()).build();
+            }
             case UNKNOWN -> throw new IllegalArgumentException("can't read values from [" + block + "]");
         };
     }

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

@@ -26,7 +26,8 @@ public final class ConstantNullBlock extends AbstractNonThreadSafeRefCounted
         FloatBlock,
         DoubleBlock,
         BytesRefBlock,
-        AggregateMetricDoubleBlock {
+        AggregateMetricDoubleBlock,
+        ExponentialHistogramBlock {
 
     private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(ConstantNullBlock.class);
     private final int positionCount;