[Star tree] Changes to handle derived metrics such as avg as part of …

…star tree mapping (opensearch-project#15152) --------- Signed-off-by: Bharathwaj G <[email protected]>
dk2k · Oct 21, 2024 · 21be0fc · 21be0fc
1 parent 7f8a228
commit 21be0fc
Show file tree

Hide file tree

Showing 6 changed files with 247 additions and 47 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java b/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java
@@ -265,13 +265,9 @@ public void testValidCompositeIndex() {
                     assertEquals(expectedTimeUnits, dateDim.getIntervals());
                     assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField());
                     assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField());
-                    List<MetricStat> expectedMetrics = Arrays.asList(
-                        MetricStat.AVG,
-                        MetricStat.VALUE_COUNT,
-                        MetricStat.SUM,
-                        MetricStat.MAX,
-                        MetricStat.MIN
-                    );
+
+                    // Assert default metrics
+                    List<MetricStat> expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG);
                     assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics());
                     assertEquals(10000, starTreeFieldType.getStarTreeConfig().maxLeafDocs());
                     assertEquals(
@@ -349,13 +345,9 @@ public void testUpdateIndexWhenMappingIsSame() {
                     assertEquals(expectedTimeUnits, dateDim.getIntervals());
                     assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField());
                     assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField());
-                    List<MetricStat> expectedMetrics = Arrays.asList(
-                        MetricStat.AVG,
-                        MetricStat.VALUE_COUNT,
-                        MetricStat.SUM,
-                        MetricStat.MAX,
-                        MetricStat.MIN
-                    );
+
+                    // Assert default metrics
+                    List<MetricStat> expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG);
                     assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics());
                     assertEquals(10000, starTreeFieldType.getStarTreeConfig().maxLeafDocs());
                     assertEquals(

diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/MetricStat.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/MetricStat.java
@@ -10,6 +10,9 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 
+import java.util.Arrays;
+import java.util.List;
+
 /**
  * Supported metric types for composite index
  *
@@ -18,21 +21,39 @@
 @ExperimentalApi
 public enum MetricStat {
     VALUE_COUNT("value_count"),
-    AVG("avg"),
     SUM("sum"),
     MIN("min"),
-    MAX("max");
+    MAX("max"),
+    AVG("avg", VALUE_COUNT, SUM);
 
     private final String typeName;
+    private final MetricStat[] baseMetrics;
 
-    MetricStat(String typeName) {
+    MetricStat(String typeName, MetricStat... baseMetrics) {
         this.typeName = typeName;
+        this.baseMetrics = baseMetrics;
     }
 
     public String getTypeName() {
         return typeName;
     }
 
+    /**
+     * Return the list of metrics that this metric is derived from
+     * For example, AVG is derived from COUNT and SUM
+     */
+    public List<MetricStat> getBaseMetrics() {
+        return Arrays.asList(baseMetrics);
+    }
+
+    /**
+     * Return true if this metric is derived from other metrics
+     * For example, AVG is derived from COUNT and SUM
+     */
+    public boolean isDerivedMetric() {
+        return baseMetrics != null && baseMetrics.length > 0;
+    }
+
     public static MetricStat fromTypeName(String typeName) {
         for (MetricStat metric : MetricStat.values()) {
             if (metric.getTypeName().equalsIgnoreCase(typeName)) {

diff --git a/...ain/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeIndexSettings.java b/...ain/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeIndexSettings.java
@@ -15,6 +15,7 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Function;
 
 /**
  * Index settings for star tree fields. The settings are final as right now
@@ -93,16 +94,10 @@ public class StarTreeIndexSettings {
     /**
      * Default metrics for metrics as part of star tree fields
      */
-    public static final Setting<List<MetricStat>> DEFAULT_METRICS_LIST = Setting.listSetting(
+    public static final Setting<List<String>> DEFAULT_METRICS_LIST = Setting.listSetting(
         "index.composite_index.star_tree.field.default.metrics",
-        Arrays.asList(
-            MetricStat.AVG.toString(),
-            MetricStat.VALUE_COUNT.toString(),
-            MetricStat.SUM.toString(),
-            MetricStat.MAX.toString(),
-            MetricStat.MIN.toString()
-        ),
-        MetricStat::fromTypeName,
+        Arrays.asList(MetricStat.VALUE_COUNT.toString(), MetricStat.SUM.toString()),
+        Function.identity(),
         Setting.Property.IndexScope,
         Setting.Property.Final
     );

diff --git a/...va/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/...va/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java
@@ -118,6 +118,9 @@ public List<MetricAggregatorInfo> generateMetricAggregatorInfos(MapperService ma
         List<MetricAggregatorInfo> metricAggregatorInfos = new ArrayList<>();
         for (Metric metric : this.starTreeField.getMetrics()) {
             for (MetricStat metricStat : metric.getMetrics()) {
+                if (metricStat.isDerivedMetric()) {
+                    continue;
+                }
                 IndexNumericFieldData.NumericType numericType;
                 Mapper fieldMapper = mapperService.documentMapper().mappers().getMapper(metric.getField());
                 if (fieldMapper instanceof NumberFieldMapper) {

diff --git a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
@@ -28,6 +28,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Queue;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -262,17 +263,50 @@ private Metric getMetric(String name, Map<String, Object> metric, Mapper.TypePar
                 .collect(Collectors.toList());
             metric.remove(STATS);
             if (metricStrings.isEmpty()) {
-                metricTypes = new ArrayList<>(StarTreeIndexSettings.DEFAULT_METRICS_LIST.get(context.getSettings()));
-            } else {
-                Set<MetricStat> metricSet = new LinkedHashSet<>();
-                for (String metricString : metricStrings) {
-                    metricSet.add(MetricStat.fromTypeName(metricString));
-                }
-                metricTypes = new ArrayList<>(metricSet);
+                metricStrings = new ArrayList<>(StarTreeIndexSettings.DEFAULT_METRICS_LIST.get(context.getSettings()));
+            }
+            // Add all required metrics initially
+            Set<MetricStat> metricSet = new LinkedHashSet<>();
+            for (String metricString : metricStrings) {
+                MetricStat metricStat = MetricStat.fromTypeName(metricString);
+                metricSet.add(metricStat);
+                addBaseMetrics(metricStat, metricSet);
             }
+            addEligibleDerivedMetrics(metricSet);
+            metricTypes = new ArrayList<>(metricSet);
             return new Metric(name, metricTypes);
         }
 
+        /**
+         * Add base metrics of derived metric to metric set
+         */
+        private void addBaseMetrics(MetricStat metricStat, Set<MetricStat> metricSet) {
+            if (metricStat.isDerivedMetric()) {
+                Queue<MetricStat> metricQueue = new LinkedList<>(metricStat.getBaseMetrics());
+                while (metricQueue.isEmpty() == false) {
+                    MetricStat metric = metricQueue.poll();
+                    if (metric.isDerivedMetric() && !metricSet.contains(metric)) {
+                        metricQueue.addAll(metric.getBaseMetrics());
+                    }
+                    metricSet.add(metric);
+                }
+            }
+        }
+
+        /**
+         * Add derived metrics if all associated base metrics are present
+         */
+        private void addEligibleDerivedMetrics(Set<MetricStat> metricStats) {
+            for (MetricStat metric : MetricStat.values()) {
+                if (metric.isDerivedMetric() && !metricStats.contains(metric)) {
+                    List<MetricStat> sourceMetrics = metric.getBaseMetrics();
+                    if (metricStats.containsAll(sourceMetrics)) {
+                        metricStats.add(metric);
+                    }
+                }
+            }
+        }
+
         @Override
         protected List<Parameter<?>> getParameters() {
             return List.of(config);