diff --git a/server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java b/server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java index 8a413a73d57e5..08325e25d3c65 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java @@ -512,6 +512,8 @@ interface BlockFactory { AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int count); + Block buildAggregateMetricDoubleDirect(Block minBlock, Block maxBlock, Block sumBlock, Block countBlock); + ExponentialHistogramBuilder exponentialHistogramBlockBuilder(int count); Block buildExponentialHistogramBlockDirect( diff --git a/test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java b/test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java index 63cce70c9efcd..95a4ec2496305 100644 --- a/test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java +++ b/test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java @@ -449,6 +449,21 @@ public BlockLoader.AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int return new AggregateMetricDoubleBlockBuilder(expectedSize); } + @Override + public BlockLoader.Block buildAggregateMetricDoubleDirect( + BlockLoader.Block minBlock, + BlockLoader.Block maxBlock, + BlockLoader.Block sumBlock, + BlockLoader.Block countBlock + ) { + return AggregateMetricDoubleBlockBuilder.parseAggMetricsToBlock( + (TestBlock) minBlock, + (TestBlock) maxBlock, + (TestBlock) sumBlock, + (TestBlock) countBlock + ); + } + @Override public BlockLoader.ExponentialHistogramBuilder exponentialHistogramBlockBuilder(int count) { return new ExponentialHistogramBlockBuilder(this, count); @@ -644,6 +659,10 @@ public BlockLoader.Block build() { var sumBlock = sum.build(); var countBlock = count.build(); + return parseAggMetricsToBlock(minBlock, maxBlock, sumBlock, countBlock); + } + + public static TestBlock parseAggMetricsToBlock(TestBlock minBlock, TestBlock maxBlock, TestBlock sumBlock, TestBlock countBlock) { assert minBlock.size() == maxBlock.size(); assert maxBlock.size() == sumBlock.size(); assert sumBlock.size() == countBlock.size(); diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java index 8d395de4f762f..dff9d421988b5 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java @@ -438,6 +438,20 @@ public AggregateMetricDoubleBlockBuilder newAggregateMetricDoubleBlockBuilder(in return new AggregateMetricDoubleBlockBuilder(estimatedSize, this); } + public final AggregateMetricDoubleBlock newAggregateMetricDoubleBlock( + Block minBlock, + Block maxBlock, + Block sumBlock, + Block countBlock + ) { + return new AggregateMetricDoubleArrayBlock( + (DoubleBlock) minBlock, + (DoubleBlock) maxBlock, + (DoubleBlock) sumBlock, + (IntBlock) countBlock + ); + } + public final AggregateMetricDoubleBlock newConstantAggregateMetricDoubleBlock( AggregateMetricDoubleBlockBuilder.AggregateMetricDoubleLiteral value, int positions @@ -469,6 +483,15 @@ public final AggregateMetricDoubleBlock newConstantAggregateMetricDoubleBlock( } } + public BlockLoader.Block newAggregateMetricDoubleBlockFromDocValues( + DoubleBlock minBlock, + DoubleBlock maxBlock, + DoubleBlock sumBlock, + IntBlock countBlock + ) { + return new AggregateMetricDoubleArrayBlock(minBlock, maxBlock, sumBlock, countBlock); + } + public ExponentialHistogramBlockBuilder newExponentialHistogramBlockBuilder(int estimatedSize) { return new ExponentialHistogramBlockBuilder(estimatedSize, this); } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/DelegatingBlockLoaderFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/DelegatingBlockLoaderFactory.java index eec7472f3feaa..53a3b866acdde 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/DelegatingBlockLoaderFactory.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/DelegatingBlockLoaderFactory.java @@ -17,6 +17,7 @@ import org.elasticsearch.compute.data.BytesRefVector; import org.elasticsearch.compute.data.DoubleBlock; import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; import org.elasticsearch.compute.data.IntVector; import org.elasticsearch.compute.data.LongBlock; import org.elasticsearch.compute.data.OrdinalBytesRefVector; @@ -151,6 +152,21 @@ public BlockLoader.AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int return factory.newAggregateMetricDoubleBlockBuilder(count); } + @Override + public BlockLoader.Block buildAggregateMetricDoubleDirect( + BlockLoader.Block minBlock, + BlockLoader.Block maxBlock, + BlockLoader.Block sumBlock, + BlockLoader.Block countBlock + ) { + return factory.newAggregateMetricDoubleBlockFromDocValues( + (DoubleBlock) minBlock, + (DoubleBlock) maxBlock, + (DoubleBlock) sumBlock, + (IntBlock) countBlock + ); + } + @Override public BlockLoader.ExponentialHistogramBuilder exponentialHistogramBlockBuilder(int count) { return factory.newExponentialHistogramBlockBuilder(count); diff --git a/x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleBlockLoader.java b/x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleBlockLoader.java new file mode 100644 index 0000000000000..d6d96c779834d --- /dev/null +++ b/x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleBlockLoader.java @@ -0,0 +1,187 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.aggregatemetric.mapper; + +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.util.NumericUtils; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.index.mapper.NumberFieldMapper; +import org.elasticsearch.index.mapper.blockloader.docvalues.BlockDocValuesReader; + +import java.io.IOException; +import java.util.EnumMap; + +public class AggregateMetricDoubleBlockLoader extends BlockDocValuesReader.DocValuesBlockLoader { + final NumberFieldMapper.NumberFieldType minFieldType; + final NumberFieldMapper.NumberFieldType maxFieldType; + final NumberFieldMapper.NumberFieldType sumFieldType; + final NumberFieldMapper.NumberFieldType countFieldType; + + AggregateMetricDoubleBlockLoader(EnumMap metricsRequested) { + minFieldType = metricsRequested.getOrDefault(AggregateMetricDoubleFieldMapper.Metric.min, null); + maxFieldType = metricsRequested.getOrDefault(AggregateMetricDoubleFieldMapper.Metric.max, null); + sumFieldType = metricsRequested.getOrDefault(AggregateMetricDoubleFieldMapper.Metric.sum, null); + countFieldType = metricsRequested.getOrDefault(AggregateMetricDoubleFieldMapper.Metric.value_count, null); + } + + private static NumericDocValues getNumericDocValues(NumberFieldMapper.NumberFieldType field, LeafReader leafReader) throws IOException { + if (field == null) { + return null; + } + String fieldName = field.name(); + var values = leafReader.getNumericDocValues(fieldName); + if (values != null) { + return values; + } + + var sortedValues = leafReader.getSortedNumericDocValues(fieldName); + return DocValues.unwrapSingleton(sortedValues); + } + + @Override + public AllReader reader(LeafReaderContext context) throws IOException { + NumericDocValues minValues = getNumericDocValues(minFieldType, context.reader()); + NumericDocValues maxValues = getNumericDocValues(maxFieldType, context.reader()); + NumericDocValues sumValues = getNumericDocValues(sumFieldType, context.reader()); + NumericDocValues valueCountValues = getNumericDocValues(countFieldType, context.reader()); + + return new BlockDocValuesReader() { + + private int docID = -1; + + @Override + protected int docId() { + return docID; + } + + @Override + public String toString() { + return "BlockDocValuesReader.AggregateMetricDouble"; + } + + @Override + public Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException { + boolean success = false; + Block minBlock = null; + Block maxBlock = null; + Block sumBlock = null; + Block countBlock = null; + try { + minBlock = readDoubleSubblock(factory, docs, offset, minValues); + maxBlock = readDoubleSubblock(factory, docs, offset, maxValues); + sumBlock = readDoubleSubblock(factory, docs, offset, sumValues); + countBlock = readIntSubblock(factory, docs, offset, valueCountValues); + Block block = factory.buildAggregateMetricDoubleDirect(minBlock, maxBlock, sumBlock, countBlock); + success = true; + return block; + } finally { + if (success == false) { + Releasables.closeExpectNoException(minBlock, maxBlock, sumBlock, countBlock); + } + } + } + + private Block readDoubleSubblock(BlockFactory factory, Docs docs, int offset, NumericDocValues values) throws IOException { + int count = docs.count() - offset; + if (values == null) { + return factory.constantNulls(count); + } + try (DoubleBuilder builder = factory.doubles(count)) { + copyDoubleValuesToBuilder(docs, offset, builder, values); + return builder.build(); + } + } + + private Block readIntSubblock(BlockFactory factory, Docs docs, int offset, NumericDocValues values) throws IOException { + int count = docs.count() - offset; + if (values == null) { + return factory.constantNulls(count); + } + try (IntBuilder builder = factory.ints(count)) { + copyIntValuesToBuilder(docs, offset, builder, values); + return builder.build(); + } + } + + private void copyDoubleValuesToBuilder(Docs docs, int offset, DoubleBuilder builder, NumericDocValues values) + throws IOException { + int lastDoc = -1; + for (int i = offset; i < docs.count(); i++) { + int doc = docs.get(i); + if (doc < lastDoc) { + throw new IllegalStateException("docs within same block must be in order"); + } + if (values == null || values.advanceExact(doc) == false) { + builder.appendNull(); + } else { + double value = NumericUtils.sortableLongToDouble(values.longValue()); + lastDoc = doc; + this.docID = doc; + builder.appendDouble(value); + } + } + } + + private void copyIntValuesToBuilder(Docs docs, int offset, IntBuilder builder, NumericDocValues values) throws IOException { + int lastDoc = -1; + for (int i = offset; i < docs.count(); i++) { + int doc = docs.get(i); + if (doc < lastDoc) { + throw new IllegalStateException("docs within same block must be in order"); + } + if (values == null || values.advanceExact(doc) == false) { + builder.appendNull(); + } else { + int value = Math.toIntExact(values.longValue()); + lastDoc = doc; + this.docID = doc; + builder.appendInt(value); + } + } + } + + @Override + public void read(int docId, StoredFields storedFields, Builder builder) throws IOException { + var blockBuilder = (AggregateMetricDoubleBuilder) builder; + this.docID = docId; + readSingleRow(docId, blockBuilder); + } + + private void readSingleRow(int docId, AggregateMetricDoubleBuilder builder) throws IOException { + if (minValues != null && minValues.advanceExact(docId)) { + builder.min().appendDouble(NumericUtils.sortableLongToDouble(minValues.longValue())); + } else { + builder.min().appendNull(); + } + if (maxValues != null && maxValues.advanceExact(docId)) { + builder.max().appendDouble(NumericUtils.sortableLongToDouble(maxValues.longValue())); + } else { + builder.max().appendNull(); + } + if (sumValues != null && sumValues.advanceExact(docId)) { + builder.sum().appendDouble(NumericUtils.sortableLongToDouble(sumValues.longValue())); + } else { + builder.sum().appendNull(); + } + if (valueCountValues != null && valueCountValues.advanceExact(docId)) { + builder.count().appendInt(Math.toIntExact(valueCountValues.longValue())); + } else { + builder.count().appendNull(); + } + } + }; + } + + @Override + public Builder builder(BlockFactory factory, int expectedCount) { + return factory.aggregateMetricDoubleBuilder(expectedCount); + } +} diff --git a/x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleFieldMapper.java b/x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleFieldMapper.java index fe7527730f4c4..e04aa5f6ad157 100644 --- a/x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleFieldMapper.java +++ b/x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleFieldMapper.java @@ -10,7 +10,6 @@ import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.NumericDocValues; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.Query; import org.apache.lucene.search.SortField; @@ -44,7 +43,6 @@ import org.elasticsearch.index.mapper.TimeSeriesParams; import org.elasticsearch.index.mapper.TimeSeriesParams.MetricType; import org.elasticsearch.index.mapper.ValueFetcher; -import org.elasticsearch.index.mapper.blockloader.docvalues.BlockDocValuesReader; import org.elasticsearch.index.query.QueryRewriteContext; import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.script.ScriptCompiler; @@ -508,143 +506,9 @@ public ValueFetcher valueFetcher(SearchExecutionContext context, String format) return SourceValueFetcher.identity(name(), context, format); } - public class AggregateMetricDoubleBlockLoader extends BlockDocValuesReader.DocValuesBlockLoader { - NumberFieldMapper.NumberFieldType minFieldType = metricFields.get(Metric.min); - NumberFieldMapper.NumberFieldType maxFieldType = metricFields.get(Metric.max); - NumberFieldMapper.NumberFieldType sumFieldType = metricFields.get(Metric.sum); - NumberFieldMapper.NumberFieldType countFieldType = metricFields.get(Metric.value_count); - - private AggregateMetricDoubleBlockLoader() {} - - static NumericDocValues getNumericDocValues(NumberFieldMapper.NumberFieldType field, LeafReader leafReader) throws IOException { - if (field == null) { - return null; - } - String fieldName = field.name(); - var values = leafReader.getNumericDocValues(fieldName); - if (values != null) { - return values; - } - - var sortedValues = leafReader.getSortedNumericDocValues(fieldName); - return DocValues.unwrapSingleton(sortedValues); - } - - @Override - public AllReader reader(LeafReaderContext context) throws IOException { - NumericDocValues minValues = getNumericDocValues(minFieldType, context.reader()); - NumericDocValues maxValues = getNumericDocValues(maxFieldType, context.reader()); - NumericDocValues sumValues = getNumericDocValues(sumFieldType, context.reader()); - NumericDocValues valueCountValues = getNumericDocValues(countFieldType, context.reader()); - - return new BlockDocValuesReader() { - - private int docID = -1; - - @Override - protected int docId() { - return docID; - } - - @Override - public String toString() { - return "BlockDocValuesReader.AggregateMetricDouble"; - } - - @Override - public Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException { - try (var builder = factory.aggregateMetricDoubleBuilder(docs.count() - offset)) { - copyDoubleValuesToBuilder(docs, offset, builder.min(), minValues); - copyDoubleValuesToBuilder(docs, offset, builder.max(), maxValues); - copyDoubleValuesToBuilder(docs, offset, builder.sum(), sumValues); - copyIntValuesToBuilder(docs, offset, builder.count(), valueCountValues); - return builder.build(); - } - } - - private void copyDoubleValuesToBuilder( - Docs docs, - int offset, - BlockLoader.DoubleBuilder builder, - NumericDocValues values - ) throws IOException { - int lastDoc = -1; - for (int i = offset; i < docs.count(); i++) { - int doc = docs.get(i); - if (doc < lastDoc) { - throw new IllegalStateException("docs within same block must be in order"); - } - if (values == null || values.advanceExact(doc) == false) { - builder.appendNull(); - } else { - double value = NumericUtils.sortableLongToDouble(values.longValue()); - lastDoc = doc; - this.docID = doc; - builder.appendDouble(value); - } - } - } - - private void copyIntValuesToBuilder(Docs docs, int offset, BlockLoader.IntBuilder builder, NumericDocValues values) - throws IOException { - int lastDoc = -1; - for (int i = offset; i < docs.count(); i++) { - int doc = docs.get(i); - if (doc < lastDoc) { - throw new IllegalStateException("docs within same block must be in order"); - } - if (values == null || values.advanceExact(doc) == false) { - builder.appendNull(); - } else { - int value = Math.toIntExact(values.longValue()); - lastDoc = doc; - this.docID = doc; - builder.appendInt(value); - } - } - } - - @Override - public void read(int docId, StoredFields storedFields, Builder builder) throws IOException { - var blockBuilder = (AggregateMetricDoubleBuilder) builder; - this.docID = docId; - readSingleRow(docId, blockBuilder); - } - - private void readSingleRow(int docId, AggregateMetricDoubleBuilder builder) throws IOException { - if (minValues != null && minValues.advanceExact(docId)) { - builder.min().appendDouble(NumericUtils.sortableLongToDouble(minValues.longValue())); - } else { - builder.min().appendNull(); - } - if (maxValues != null && maxValues.advanceExact(docId)) { - builder.max().appendDouble(NumericUtils.sortableLongToDouble(maxValues.longValue())); - } else { - builder.max().appendNull(); - } - if (sumValues != null && sumValues.advanceExact(docId)) { - builder.sum().appendDouble(NumericUtils.sortableLongToDouble(sumValues.longValue())); - } else { - builder.sum().appendNull(); - } - if (valueCountValues != null && valueCountValues.advanceExact(docId)) { - builder.count().appendInt(Math.toIntExact(valueCountValues.longValue())); - } else { - builder.count().appendNull(); - } - } - }; - } - - @Override - public Builder builder(BlockFactory factory, int expectedCount) { - return factory.aggregateMetricDoubleBuilder(expectedCount); - } - } - @Override public BlockLoader blockLoader(BlockLoaderContext blContext) { - return new AggregateMetricDoubleBlockLoader(); + return new AggregateMetricDoubleBlockLoader(metricFields); } /** @@ -935,4 +799,5 @@ public boolean advanceToDoc(int docId) throws IOException { } } } + } diff --git a/x-pack/plugin/mapper-aggregate-metric/src/test/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoublePartialBlockLoaderTests.java b/x-pack/plugin/mapper-aggregate-metric/src/test/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoublePartialBlockLoaderTests.java new file mode 100644 index 0000000000000..0ea57da5f2565 --- /dev/null +++ b/x-pack/plugin/mapper-aggregate-metric/src/test/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoublePartialBlockLoaderTests.java @@ -0,0 +1,205 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.aggregatemetric.mapper; + +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.elasticsearch.index.mapper.BlockLoader; +import org.elasticsearch.index.mapper.NumberFieldMapper; +import org.elasticsearch.index.mapper.TestBlock; +import org.elasticsearch.test.ESTestCase; +import org.hamcrest.Matcher; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.EnumMap; +import java.util.HashMap; +import java.util.List; + +import static org.elasticsearch.xpack.aggregatemetric.mapper.AggregateMetricDoubleFieldMapper.subfieldName; +import static org.hamcrest.Matchers.hasToString; + +public class AggregateMetricDoublePartialBlockLoaderTests extends ESTestCase { + @ParametersFactory(argumentFormatting = "missingValues=%s, loadMin=%s, loadMax=%s, loadSum=%s, loadCount=%s") + public static List parameters() throws IOException { + List parameters = new ArrayList<>(); + for (boolean missingValues : new boolean[] { true, false }) { + for (boolean loadMin : new boolean[] { true, false }) { + for (boolean loadMax : new boolean[] { true, false }) { + for (boolean loadSum : new boolean[] { true, false }) { + for (boolean loadCount : new boolean[] { true, false }) { + parameters.add(new Object[] { missingValues, loadMin, loadMax, loadSum, loadCount }); + } + } + } + } + } + return parameters; + } + + protected final boolean missingValues; + protected final boolean loadMin; + protected final boolean loadMax; + protected final boolean loadSum; + protected final boolean loadCount; + + public AggregateMetricDoublePartialBlockLoaderTests( + boolean missingValues, + boolean loadMin, + boolean loadMax, + boolean loadSum, + boolean loadCount + ) { + this.missingValues = missingValues; + this.loadMin = loadMin; + this.loadMax = loadMax; + this.loadSum = loadSum; + this.loadCount = loadCount; + } + + public void test() throws IOException { + try (Directory dir = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), dir)) { + int docCount = 10_000; + for (int i = 0; i < docCount; i++) { + List doc = new ArrayList<>(4); + doc.add(new NumericDocValuesField("field.min", Double.doubleToLongBits(i))); + doc.add(new NumericDocValuesField("field.max", Double.doubleToLongBits(i))); + doc.add(new NumericDocValuesField("field.sum", Double.doubleToLongBits(i))); + doc.add(new NumericDocValuesField("field.value_count", i)); + iw.addDocument(doc); + } + if (missingValues) { + iw.addDocument(List.of()); + } + iw.forceMerge(1); + try (DirectoryReader dr = iw.getReader()) { + LeafReaderContext ctx = getOnlyLeafReader(dr).getContext(); + innerTest(ctx); + } + } + } + + private void innerTest(LeafReaderContext ctx) throws IOException { + var allMetricFields = generateMetricFields(); + var metricFieldsToLoad = generateMetricsToLoad(allMetricFields); + AggregateMetricDoubleBlockLoader aggMetricPartialLoader = new AggregateMetricDoubleBlockLoader(metricFieldsToLoad); + AggregateMetricDoubleBlockLoader aggMetricAllLoader = new AggregateMetricDoubleBlockLoader(allMetricFields); + var aggMetricPartialReader = aggMetricPartialLoader.reader(ctx); + var aggMetricAllReader = aggMetricAllLoader.reader(ctx); + assertThat(aggMetricPartialReader, readerMatcher()); + assertThat(aggMetricAllReader, readerMatcher()); + + BlockLoader.Docs docs = TestBlock.docs(ctx); + try ( + TestBlock aggMetricsPartial = read(aggMetricPartialLoader, aggMetricPartialReader, ctx, docs); + TestBlock aggMetricsAll = read(aggMetricAllLoader, aggMetricAllReader, ctx, docs) + ) { + checkBlocks(aggMetricsAll, aggMetricsPartial); + } + + aggMetricPartialReader = aggMetricPartialLoader.reader(ctx); + aggMetricAllReader = aggMetricAllLoader.reader(ctx); + for (int i = 0; i < ctx.reader().numDocs(); i += 10) { + int[] docsArray = new int[Math.min(10, ctx.reader().numDocs() - i)]; + for (int d = 0; d < docsArray.length; d++) { + docsArray[d] = i + d; + } + docs = TestBlock.docs(docsArray); + try ( + TestBlock aggMetricsPartial = read(aggMetricPartialLoader, aggMetricPartialReader, ctx, docs); + TestBlock aggMetricsAll = read(aggMetricAllLoader, aggMetricAllReader, ctx, docs) + ) { + checkBlocks(aggMetricsAll, aggMetricsPartial); + } + } + } + + private EnumMap generateMetricsToLoad( + EnumMap allMetricFields + ) { + EnumMap metricFieldsToLoad = new EnumMap<>( + AggregateMetricDoubleFieldMapper.Metric.class + ); + if (loadMin) { + metricFieldsToLoad.put( + AggregateMetricDoubleFieldMapper.Metric.min, + allMetricFields.get(AggregateMetricDoubleFieldMapper.Metric.min) + ); + } + if (loadMax) { + metricFieldsToLoad.put( + AggregateMetricDoubleFieldMapper.Metric.max, + allMetricFields.get(AggregateMetricDoubleFieldMapper.Metric.max) + ); + } + if (loadSum) { + metricFieldsToLoad.put( + AggregateMetricDoubleFieldMapper.Metric.sum, + allMetricFields.get(AggregateMetricDoubleFieldMapper.Metric.sum) + ); + } + if (loadCount) { + metricFieldsToLoad.put( + AggregateMetricDoubleFieldMapper.Metric.value_count, + allMetricFields.get(AggregateMetricDoubleFieldMapper.Metric.value_count) + ); + } + return metricFieldsToLoad; + } + + private EnumMap generateMetricFields() { + EnumMap metricFields = new EnumMap<>( + AggregateMetricDoubleFieldMapper.Metric.class + ); + for (AggregateMetricDoubleFieldMapper.Metric m : List.of(AggregateMetricDoubleFieldMapper.Metric.values())) { + String subfieldName = subfieldName("field", m); + NumberFieldMapper.NumberFieldType subfield = new NumberFieldMapper.NumberFieldType( + subfieldName, + m == AggregateMetricDoubleFieldMapper.Metric.value_count + ? NumberFieldMapper.NumberType.INTEGER + : NumberFieldMapper.NumberType.DOUBLE + ); + metricFields.put(m, subfield); + } + return metricFields; + } + + private Matcher readerMatcher() { + return hasToString("BlockDocValuesReader.AggregateMetricDouble"); + } + + @SuppressWarnings("unchecked") + private void checkBlocks(TestBlock allMetrics, TestBlock partialMetrics) { + for (int i = 0; i < allMetrics.size(); i++) { + var all = (HashMap) (allMetrics.get(i)); + var partial = (HashMap) (partialMetrics.get(i)); + checkMetric(loadMin, AggregateMetricDoubleFieldMapper.Metric.min.name(), all, partial); + checkMetric(loadMax, AggregateMetricDoubleFieldMapper.Metric.max.name(), all, partial); + checkMetric(loadSum, AggregateMetricDoubleFieldMapper.Metric.sum.name(), all, partial); + checkMetric(loadCount, AggregateMetricDoubleFieldMapper.Metric.value_count.name(), all, partial); + } + } + + private void checkMetric(boolean loadField, String field, HashMap allMetrics, HashMap partialMetrics) { + if (loadField) { + assertEquals(allMetrics.get(field), partialMetrics.get(field)); + } else { + assertNull(partialMetrics.get(field)); + } + } + + protected final TestBlock read(BlockLoader loader, BlockLoader.AllReader reader, LeafReaderContext ctx, BlockLoader.Docs docs) + throws IOException { + return (TestBlock) reader.read(TestBlock.factory(), docs, 0, false); + } +}