ESQL: Avoid null in intermediate exponential_histogram merge state (#138511)

Author: JonasKunz

x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/HistogramMergeExponentialHistogramAggregatorFunction.java

@@ -64,8 +64,16 @@ public void add(ExponentialHistogram histogram, boolean allowUpscale) {
 
         @Override
         public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
-            assert blocks.length >= offset + 1;
-            blocks[offset] = evaluateFinal(driverContext);
+            assert blocks.length >= offset + 2;
+            BlockFactory blockFactory = driverContext.blockFactory();
+            // in case of error, the blocks are closed by the caller
+            if (merger == null) {
+                blocks[offset] = blockFactory.newConstantExponentialHistogramBlock(ExponentialHistogram.empty(), 1);
+                blocks[offset + 1] = blockFactory.newConstantBooleanBlockWith(false, 1);
+            } else {
+                blocks[offset] = blockFactory.newConstantExponentialHistogramBlock(merger.get(), 1);
+                blocks[offset + 1] = blockFactory.newConstantBooleanBlockWith(true, 1);
+            }
         }
 
         @Override
@@ -127,8 +135,25 @@ private void ensureCapacity(int groupId) {
 
         @Override
         public void toIntermediate(Block[] blocks, int offset, IntVector selected, DriverContext driverContext) {
-            assert blocks.length >= offset + 1 : "blocks=" + blocks.length + ",offset=" + offset;
-            blocks[offset] = evaluateFinal(selected, driverContext);
+            assert blocks.length >= offset + 2 : "blocks=" + blocks.length + ",offset=" + offset;
+            try (
+                var histoBuilder = driverContext.blockFactory().newExponentialHistogramBlockBuilder(selected.getPositionCount());
+                var seenBuilder = driverContext.blockFactory().newBooleanBlockBuilder(selected.getPositionCount());
+            ) {
+                for (int i = 0; i < selected.getPositionCount(); i++) {
+                    int groupId = selected.getInt(i);
+                    ExponentialHistogramMerger state = getOrNull(groupId);
+                    if (state != null) {
+                        seenBuilder.appendBoolean(true);
+                        histoBuilder.append(state.get());
+                    } else {
+                        seenBuilder.appendBoolean(false);
+                        histoBuilder.append(ExponentialHistogram.empty());
+                    }
+                }
+                blocks[offset] = histoBuilder.build();
+                blocks[offset + 1] = seenBuilder.build();
+            }
         }
 
         public Block evaluateFinal(IntVector selected, DriverContext driverContext) {

x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/HistogramMergeExponentialHistogramGroupingAggregatorFunction.java

@@ -16,7 +16,7 @@
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.exponentialhistogram.ExponentialHistogram;
 
-@Aggregator({ @IntermediateState(name = "value", type = "EXPONENTIAL_HISTOGRAM"), })
+@Aggregator({ @IntermediateState(name = "value", type = "EXPONENTIAL_HISTOGRAM"), @IntermediateState(name = "seen", type = "BOOLEAN") })
 @GroupingAggregator
 public class HistogramMergeExponentialHistogramAggregator {
 
@@ -28,8 +28,10 @@ public static void combine(ExponentialHistogramStates.SingleState state, Exponen
         state.add(value, true);
     }
 
-    public static void combineIntermediate(ExponentialHistogramStates.SingleState state, ExponentialHistogram value) {
-        state.add(value, false);
+    public static void combineIntermediate(ExponentialHistogramStates.SingleState state, ExponentialHistogram value, boolean seen) {
+        if (seen) {
+            state.add(value, false);
+        }
     }
 
     public static Block evaluateFinal(ExponentialHistogramStates.SingleState state, DriverContext driverContext) {
@@ -44,8 +46,15 @@ public static void combine(ExponentialHistogramStates.GroupingState current, int
         current.add(groupId, value, true);
     }
 
-    public static void combineIntermediate(ExponentialHistogramStates.GroupingState state, int groupId, ExponentialHistogram value) {
-        state.add(groupId, value, false);
+    public static void combineIntermediate(
+        ExponentialHistogramStates.GroupingState state,
+        int groupId,
+        ExponentialHistogram value,
+        boolean seen
+    ) {
+        if (seen) {
+            state.add(groupId, value, false);
+        }
     }
 
     public static Block evaluateFinal(

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

@@ -108,6 +108,7 @@ private List<Block> getSubBlocks() {
 
     @Override
     public ExponentialHistogram getExponentialHistogram(int valueIndex, ExponentialHistogramScratch scratch) {
+        assert isNull(valueIndex) == false : "tried to get histogram at null position " + valueIndex;
         BytesRef bytes = encodedHistograms.getBytesRef(encodedHistograms.getFirstValueIndex(valueIndex), scratch.bytesRefScratch);
         double zeroThreshold = zeroThresholds.getDouble(zeroThresholds.getFirstValueIndex(valueIndex));
         double valueCount = valueCounts.getDouble(valueCounts.getFirstValueIndex(valueIndex));

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

@@ -430,11 +430,11 @@ protected boolean supportsExponentialHistograms() {
         try {
             return RestEsqlTestCase.hasCapabilities(
                 client(),
-                List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V3.capabilityName())
+                List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V4.capabilityName())
             )
                 && RestEsqlTestCase.hasCapabilities(
                     remoteClusterClient(),
-                    List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V3.capabilityName())
+                    List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V4.capabilityName())
                 );
         } catch (IOException e) {
             throw new RuntimeException(e);

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

@@ -58,7 +58,7 @@ protected boolean supportsSourceFieldMapping() {
     protected boolean supportsExponentialHistograms() {
         return RestEsqlTestCase.hasCapabilities(
             client(),
-            List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V3.capabilityName())
+            List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V4.capabilityName())
         );
     }
 

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java

@@ -289,7 +289,7 @@ protected boolean supportsSourceFieldMapping() throws IOException {
     protected boolean supportsExponentialHistograms() {
         return RestEsqlTestCase.hasCapabilities(
             client(),
-            List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V3.capabilityName())
+            List.of(EsqlCapabilities.Cap.EXPONENTIAL_HISTOGRAM_PRE_TECH_PREVIEW_V4.capabilityName())
         );
     }