[Star tree][Bug] Fix for derived metrics

server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesReader.java

@@ -164,7 +164,7 @@ public Composite99DocValuesReader(DocValuesProducer producer, SegmentReadState r
 
                             // adding metric fields
                             for (Metric metric : starTreeMetadata.getMetrics()) {
-                                for (MetricStat metricStat : metric.getMetrics()) {
+                                for (MetricStat metricStat : metric.getBaseMetrics()) {
                                     fields.add(
                                         fullyQualifiedFieldNameForStarTreeMetricsDocValues(
                                             compositeFieldName,

server/src/main/java/org/opensearch/index/compositeindex/datacube/Metric.java

@@ -13,6 +13,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 
@@ -23,10 +24,18 @@
 public class Metric implements ToXContent {
     private final String field;
     private final List<MetricStat> metrics;
+    private final List<MetricStat> baseMetrics;
 
     public Metric(String field, List<MetricStat> metrics) {
         this.field = field;
         this.metrics = metrics;
+        this.baseMetrics = new ArrayList<>();
+        for (MetricStat metricStat : metrics) {
+            if (metricStat.isDerivedMetric()) {
+                continue;
+            }
+            baseMetrics.add(metricStat);
+        }
     }
 
     public String getField() {
@@ -37,6 +46,13 @@ public List<MetricStat> getMetrics() {
         return metrics;
     }
 
+    /**
+     * Returns only the base metrics
+     */
+    public List<MetricStat> getBaseMetrics() {
+        return baseMetrics;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java

@@ -148,10 +148,7 @@ public List<MetricAggregatorInfo> generateMetricAggregatorInfos(MapperService ma
                 metricAggregatorInfos.add(metricAggregatorInfo);
                 continue;
             }
-            for (MetricStat metricStat : metric.getMetrics()) {
-                if (metricStat.isDerivedMetric()) {
-                    continue;
-                }
+            for (MetricStat metricStat : metric.getBaseMetrics()) {
                 IndexNumericFieldData.NumericType numericType;
                 Mapper fieldMapper = mapperService.documentMapper().mappers().getMapper(metric.getField());
                 if (fieldMapper instanceof NumberFieldMapper) {
@@ -183,7 +180,7 @@ public List<SequentialDocValuesIterator> getMetricReaders(SegmentWriteState stat
 
         List<SequentialDocValuesIterator> metricReaders = new ArrayList<>();
         for (Metric metric : this.starTreeField.getMetrics()) {
-            for (MetricStat metricStat : metric.getMetrics()) {
+            for (MetricStat metricStat : metric.getBaseMetrics()) {
                 SequentialDocValuesIterator metricReader;
                 FieldInfo metricFieldInfo = state.fieldInfos.fieldInfo(metric.getField());
                 if (metricStat.equals(MetricStat.DOC_COUNT)) {

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java

@@ -158,7 +158,7 @@ Iterator<StarTreeDocument> mergeStarTrees(List<StarTreeValues> starTreeValuesSub
                 List<SequentialDocValuesIterator> metricReaders = new ArrayList<>();
                 // get doc id set iterators for metrics
                 for (Metric metric : starTreeValues.getStarTreeField().getMetrics()) {
-                    for (MetricStat metricStat : metric.getMetrics()) {
+                    for (MetricStat metricStat : metric.getBaseMetrics()) {
                         String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues(
                             starTreeValues.getStarTreeField().getName(),
                             metric.getField(),

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java

@@ -144,7 +144,7 @@ StarTreeDocument[] getSegmentsStarTreeDocuments(List<StarTreeValues> starTreeVal
             List<SequentialDocValuesIterator> metricReaders = new ArrayList<>();
             // get doc id set iterators for metrics
             for (Metric metric : starTreeValues.getStarTreeField().getMetrics()) {
-                for (MetricStat metricStat : metric.getMetrics()) {
+                for (MetricStat metricStat : metric.getBaseMetrics()) {
                     String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues(
                         starTreeValues.getStarTreeField().getName(),
                         metric.getField(),

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadata.java

@@ -220,16 +220,37 @@ private int readMetricsCount() throws IOException {
     private List<Metric> readMetricEntries() throws IOException {
         int metricCount = readMetricsCount();
 
-        Map<String, Metric> starTreeMetricMap = new LinkedHashMap<>();
+        Map<String, List<MetricStat>> starTreeMetricStatMap = new LinkedHashMap<>();
         for (int i = 0; i < metricCount; i++) {
             String metricName = meta.readString();
             int metricStatOrdinal = meta.readVInt();
             MetricStat metricStat = MetricStat.fromMetricOrdinal(metricStatOrdinal);
-            Metric metric = starTreeMetricMap.computeIfAbsent(metricName, field -> new Metric(field, new ArrayList<>()));
-            metric.getMetrics().add(metricStat);
+            List<MetricStat> metricStats = starTreeMetricStatMap.computeIfAbsent(metricName, field -> new ArrayList<>());
+            metricStats.add(metricStat);
         }
+        List<Metric> starTreeMetricMap = new ArrayList<>();
+        for (Map.Entry<String, List<MetricStat>> metricStatsEntry : starTreeMetricStatMap.entrySet()) {
+            addEligibleDerivedMetrics(metricStatsEntry.getValue());
+            starTreeMetricMap.add(new Metric(metricStatsEntry.getKey(), metricStatsEntry.getValue()));
 
-        return new ArrayList<>(starTreeMetricMap.values());
+        }
+        return starTreeMetricMap;
+    }
+
+    /**
+     * Add derived metrics if all associated base metrics are present
+     */
+    private void addEligibleDerivedMetrics(List<MetricStat> metricStatsList) {
+        Set<MetricStat> metricStatsSet = new HashSet<>(metricStatsList);
+        for (MetricStat metric : MetricStat.values()) {
+            if (metric.isDerivedMetric() && !metricStatsSet.contains(metric)) {
+                List<MetricStat> sourceMetrics = metric.getBaseMetrics();
+                if (metricStatsSet.containsAll(sourceMetrics)) {
+                    metricStatsList.add(metric);
+                    metricStatsSet.add(metric);
+                }
+            }
+        }
     }
 
     private int readSegmentAggregatedDocCount() throws IOException {

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java

@@ -171,7 +171,7 @@ public StarTreeValues(
 
         // get doc id set iterators for metrics
         for (Metric metric : starTreeMetadata.getMetrics()) {
-            for (MetricStat metricStat : metric.getMetrics()) {
+            for (MetricStat metricStat : metric.getBaseMetrics()) {
                 String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues(
                     starTreeField.getName(),
                     metric.getField(),

server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java

@@ -292,11 +292,7 @@ private List<Metric> buildMetrics(String fieldName, Map<String, Object> map, Map
             }
             int numBaseMetrics = 0;
             for (Metric metric : metrics) {
-                for (MetricStat metricStat : metric.getMetrics()) {
-                    if (metricStat.isDerivedMetric() == false) {
-                        numBaseMetrics++;
-                    }
-                }
+                numBaseMetrics += metric.getBaseMetrics().size();
             }
             if (numBaseMetrics > context.getSettings()
                 .getAsInt(

server/src/test/java/org/opensearch/index/codec/composite99/datacube/startree/StarTreeDocValuesFormatTests.java

@@ -119,23 +119,23 @@ public void testStarTreeDocValues() throws IOException {
         Document doc = new Document();
         doc.add(new SortedNumericDocValuesField("sndv", 1));
         doc.add(new SortedNumericDocValuesField("dv", 1));
-        doc.add(new SortedNumericDocValuesField("field", 1));
+        doc.add(new SortedNumericDocValuesField("field", -1));
         iw.addDocument(doc);
         doc = new Document();
         doc.add(new SortedNumericDocValuesField("sndv", 1));
         doc.add(new SortedNumericDocValuesField("dv", 1));
-        doc.add(new SortedNumericDocValuesField("field", 1));
+        doc.add(new SortedNumericDocValuesField("field", -1));
         iw.addDocument(doc);
         doc = new Document();
         iw.forceMerge(1);
         doc.add(new SortedNumericDocValuesField("sndv", 2));
         doc.add(new SortedNumericDocValuesField("dv", 2));
-        doc.add(new SortedNumericDocValuesField("field", 2));
+        doc.add(new SortedNumericDocValuesField("field", -2));
         iw.addDocument(doc);
         doc = new Document();
         doc.add(new SortedNumericDocValuesField("sndv", 2));
         doc.add(new SortedNumericDocValuesField("dv", 2));
-        doc.add(new SortedNumericDocValuesField("field", 2));
+        doc.add(new SortedNumericDocValuesField("field", -2));
         iw.addDocument(doc);
         iw.forceMerge(1);
         iw.close();
@@ -144,11 +144,39 @@ public void testStarTreeDocValues() throws IOException {
         TestUtil.checkReader(ir);
         assertEquals(1, ir.leaves().size());
 
+        // Segment documents
+        /**
+         * sndv dv field
+         * [1,  1,   -1]
+         * [1,  1,   -1]
+         * [2,  2,   -2]
+         * [2,  2,   -2]
+         */
+        // Star tree docuements
+        /**
+         * sndv dv | [ sum, value_count, min, max[field]] , [ sum, value_count, min, max[sndv]], doc_count
+         * [1, 1] | [-2.0, 2.0, -1.0, -1.0, 2.0, 2.0, 1.0, 1.0, 2.0]
+         * [2, 2] | [-4.0, 2.0, -2.0, -2.0, 4.0, 2.0, 2.0, 2.0, 2.0]
+         * [null, 1] | [-2.0, 2.0, -1.0, -1.0, 2.0, 2.0, 1.0, 1.0, 2.0]
+         * [null, 2] | [-4.0, 2.0, -2.0, -2.0, 4.0, 2.0, 2.0, 2.0, 2.0]
+         */
         StarTreeDocument[] expectedStarTreeDocuments = new StarTreeDocument[4];
-        expectedStarTreeDocuments[0] = new StarTreeDocument(new Long[] { 1L, 1L }, new Double[] { 2.0, 2.0, 2.0 });
-        expectedStarTreeDocuments[1] = new StarTreeDocument(new Long[] { 2L, 2L }, new Double[] { 4.0, 2.0, 4.0 });
-        expectedStarTreeDocuments[2] = new StarTreeDocument(new Long[] { null, 1L }, new Double[] { 2.0, 2.0, 2.0 });
-        expectedStarTreeDocuments[3] = new StarTreeDocument(new Long[] { null, 2L }, new Double[] { 4.0, 2.0, 4.0 });
+        expectedStarTreeDocuments[0] = new StarTreeDocument(
+            new Long[] { 1L, 1L },
+            new Double[] { -2.0, 2.0, -1.0, -1.0, 2.0, 2.0, 1.0, 1.0, 2.0 }
+        );
+        expectedStarTreeDocuments[1] = new StarTreeDocument(
+            new Long[] { 2L, 2L },
+            new Double[] { -4.0, 2.0, -2.0, -2.0, 4.0, 2.0, 2.0, 2.0, 2.0 }
+        );
+        expectedStarTreeDocuments[2] = new StarTreeDocument(
+            new Long[] { null, 1L },
+            new Double[] { -2.0, 2.0, -1.0, -1.0, 2.0, 2.0, 1.0, 1.0, 2.0 }
+        );
+        expectedStarTreeDocuments[3] = new StarTreeDocument(
+            new Long[] { null, 2L },
+            new Double[] { -4.0, 2.0, -2.0, -2.0, 4.0, 2.0, 2.0, 2.0, 2.0 }
+        );
 
         for (LeafReaderContext context : ir.leaves()) {
             SegmentReader reader = Lucene.segmentReader(context.reader());
@@ -159,7 +187,17 @@ public void testStarTreeDocValues() throws IOException {
                 StarTreeValues starTreeValues = (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(compositeIndexFieldInfo);
                 StarTreeDocument[] starTreeDocuments = StarTreeTestUtils.getSegmentsStarTreeDocuments(
                     List.of(starTreeValues),
-                    List.of(StarTreeNumericType.DOUBLE, StarTreeNumericType.LONG, StarTreeNumericType.LONG),
+                    List.of(
+                        StarTreeNumericType.DOUBLE,
+                        StarTreeNumericType.LONG,
+                        StarTreeNumericType.DOUBLE,
+                        StarTreeNumericType.DOUBLE,
+                        StarTreeNumericType.DOUBLE,
+                        StarTreeNumericType.LONG,
+                        StarTreeNumericType.DOUBLE,
+                        StarTreeNumericType.DOUBLE,
+                        StarTreeNumericType.LONG
+                    ),
                     reader.maxDoc()
                 );
                 assertStarTreeDocuments(starTreeDocuments, expectedStarTreeDocuments);
@@ -190,6 +228,19 @@ private XContentBuilder getExpandedMapping() throws IOException {
             b.startArray("stats");
             b.value("sum");
             b.value("value_count");
+            b.value("avg");
+            b.value("min");
+            b.value("max");
+            b.endArray();
+            b.endObject();
+            b.startObject();
+            b.field("name", "sndv");
+            b.startArray("stats");
+            b.value("sum");
+            b.value("value_count");
+            b.value("avg");
+            b.value("min");
+            b.value("max");
             b.endArray();
             b.endObject();
             b.endArray();

server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java

@@ -61,6 +61,9 @@ public static StarTreeDocument[] getSegmentsStarTreeDocuments(
             // get doc id set iterators for metrics
             for (Metric metric : starTreeValues.getStarTreeField().getMetrics()) {
                 for (MetricStat metricStat : metric.getMetrics()) {
+                    if (metricStat.isDerivedMetric()) {
+                        continue;
+                    }
                     String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues(
                         starTreeValues.getStarTreeField().getName(),
                         metric.getField(),
@@ -125,18 +128,18 @@ public static void assertStarTreeDocuments(StarTreeDocument[] starTreeDocuments,
             assertNotNull(resultStarTreeDocument.dimensions);
             assertNotNull(resultStarTreeDocument.metrics);
 
-            assertEquals(resultStarTreeDocument.dimensions.length, expectedStarTreeDocument.dimensions.length);
-            assertEquals(resultStarTreeDocument.metrics.length, expectedStarTreeDocument.metrics.length);
+            assertEquals(expectedStarTreeDocument.dimensions.length, resultStarTreeDocument.dimensions.length);
+            assertEquals(expectedStarTreeDocument.metrics.length, resultStarTreeDocument.metrics.length);
 
             for (int di = 0; di < resultStarTreeDocument.dimensions.length; di++) {
-                assertEquals(resultStarTreeDocument.dimensions[di], expectedStarTreeDocument.dimensions[di]);
+                assertEquals(expectedStarTreeDocument.dimensions[di], resultStarTreeDocument.dimensions[di]);
             }
 
             for (int mi = 0; mi < resultStarTreeDocument.metrics.length; mi++) {
                 if (expectedStarTreeDocument.metrics[mi] instanceof Long) {
-                    assertEquals(resultStarTreeDocument.metrics[mi], ((Long) expectedStarTreeDocument.metrics[mi]).doubleValue());
+                    assertEquals(((Long) expectedStarTreeDocument.metrics[mi]).doubleValue(), resultStarTreeDocument.metrics[mi]);
                 } else {
-                    assertEquals(resultStarTreeDocument.metrics[mi], expectedStarTreeDocument.metrics[mi]);
+                    assertEquals(expectedStarTreeDocument.metrics[mi], resultStarTreeDocument.metrics[mi]);
                 }
             }
         }
@@ -267,9 +270,34 @@ public static void assertStarTreeMetadata(StarTreeMetadata expectedStarTreeMetad
             Metric expectedMetric = expectedStarTreeMetadata.getMetrics().get(i);
             Metric resultMetric = resultStarTreeMetadata.getMetrics().get(i);
             assertEquals(expectedMetric.getField(), resultMetric.getField());
+            List<MetricStat> metricStats = new ArrayList<>();
+            for (MetricStat metricStat : expectedMetric.getMetrics()) {
+                if (metricStat.isDerivedMetric()) {
+                    continue;
+                }
+                metricStats.add(metricStat);
+            }
+            Metric expectedMetricWithoutDerivedMetrics = new Metric(expectedMetric.getField(), metricStats);
+            metricStats = new ArrayList<>();
+            for (MetricStat metricStat : resultMetric.getMetrics()) {
+                if (metricStat.isDerivedMetric()) {
+                    continue;
+                }
+                metricStats.add(metricStat);
+            }
+            Metric resultantMetricWithoutDerivedMetrics = new Metric(resultMetric.getField(), metricStats);
+
+            // assert base metrics are in order in metadata
+            for (int j = 0; j < expectedMetricWithoutDerivedMetrics.getMetrics().size(); j++) {
+                assertEquals(
+                    expectedMetricWithoutDerivedMetrics.getMetrics().get(j),
+                    resultantMetricWithoutDerivedMetrics.getMetrics().get(j)
+                );
+            }
 
+            // assert all metrics ( including derived metrics are present )
             for (int j = 0; j < expectedMetric.getMetrics().size(); j++) {
-                assertEquals(expectedMetric.getMetrics().get(j), resultMetric.getMetrics().get(j));
+                assertTrue(resultMetric.getMetrics().contains(expectedMetric.getMetrics().get(j)));
             }
 
         }