Fix and unmute ExponentialHistogramFieldMapperTests.testFormattedDocValues (#138659)

Author: JonasKunz

muted-tests.yml

@@ -427,9 +427,6 @@ tests:
 - class: org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT
   method: testCreatesEisChatCompletion_DoesNotRemoveEndpointWhenNoLongerAuthorized
   issue: https://github.com/elastic/elasticsearch/issues/138480
-- class: org.elasticsearch.xpack.exponentialhistogram.ExponentialHistogramFieldMapperTests
-  method: testFormattedDocValues
-  issue: https://github.com/elastic/elasticsearch/issues/138504
 - class: org.elasticsearch.xpack.esql.heap_attack.HeapAttackLookupJoinIT
   method: testLookupExplosionBigString
   issue: https://github.com/elastic/elasticsearch/issues/138510

x-pack/plugin/mapper-exponential-histogram/src/test/java/org/elasticsearch/xpack/exponentialhistogram/ExponentialHistogramFieldMapperTests.java

@@ -7,11 +7,13 @@
 
 package org.elasticsearch.xpack.exponentialhistogram;
 
+import org.apache.lucene.document.NumericDocValuesField;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.IndexableField;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.analysis.MockAnalyzer;
 import org.apache.lucene.tests.index.RandomIndexWriter;
@@ -59,6 +61,7 @@
 import java.util.Map;
 import java.util.OptionalDouble;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.IntStream;
 
 import static org.elasticsearch.exponentialhistogram.ExponentialHistogram.MAX_INDEX;
@@ -725,18 +728,33 @@ public void testFormattedDocValues() throws IOException {
 
             IndexWriterConfig config = LuceneTestCase.newIndexWriterConfig(random(), new MockAnalyzer(random()));
             RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory, config);
-            inputHistograms.forEach(histo -> ExponentialHistogramAggregatorTestCase.addHistogramDoc(indexWriter, "field", histo));
+
+            // give each document a ID field because we are not guaranteed to read them in-order later
+            AtomicInteger currentId = new AtomicInteger(0);
+            inputHistograms.forEach(
+                histo -> ExponentialHistogramAggregatorTestCase.addHistogramDoc(
+                    indexWriter,
+                    "field",
+                    histo,
+                    new NumericDocValuesField("histo_index", currentId.getAndIncrement())
+                )
+            );
             indexWriter.close();
 
+            int seenCount = 0;
             try (DirectoryReader reader = DirectoryReader.open(directory)) {
                 for (int i = 0; i < reader.leaves().size(); i++) {
                     LeafReaderContext leaf = reader.leaves().get(i);
-                    int docBase = leaf.docBase;
                     LeafReader leafReader = leaf.reader();
                     int maxDoc = leafReader.maxDoc();
                     FormattedDocValues docValues = ExponentialHistogramFieldMapper.createFormattedDocValues(leafReader, "field");
+                    NumericDocValues histoIndex = leafReader.getNumericDocValues("histo_index");
                     for (int j = 0; j < maxDoc; j++) {
-                        var expectedHistogram = inputHistograms.get(docBase + j);
+                        assertThat(histoIndex.advanceExact(j), equalTo(true));
+                        int index = (int) histoIndex.longValue();
+                        var expectedHistogram = inputHistograms.get(index);
+                        seenCount++;
+
                         if (expectedHistogram == null) {
                             assertThat(docValues.advanceExact(j), equalTo(false));
                             expectThrows(IllegalStateException.class, docValues::nextValue);
@@ -750,6 +768,7 @@ public void testFormattedDocValues() throws IOException {
                     }
                 }
             }
+            assertThat(seenCount, equalTo(inputHistograms.size()));
         }
     }
 

x-pack/plugin/mapper-exponential-histogram/src/test/java/org/elasticsearch/xpack/exponentialhistogram/aggregations/ExponentialHistogramAggregatorTestCase.java

@@ -21,7 +21,6 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
@@ -52,7 +51,7 @@ public static void addHistogramDoc(
     ) {
         try {
             if (histogram == null) {
-                iw.addDocument(Collections.emptyList());
+                iw.addDocument(List.of(additionalFields));
             } else {
                 ExponentialHistogramFieldMapper.HistogramDocValueFields docValues = ExponentialHistogramFieldMapper.buildDocValueFields(
                     fieldName,