[Star tree] Doc count field support in star tree

server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java

@@ -264,6 +264,7 @@ public void testValidCompositeIndex() {
                     );
                     assertEquals(expectedTimeUnits, dateDim.getIntervals());
                     assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField());
+                    assertEquals(2, starTreeFieldType.getMetrics().size());
                     assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField());
                     List<MetricStat> expectedMetrics = Arrays.asList(
                         MetricStat.AVG,
@@ -273,6 +274,10 @@ public void testValidCompositeIndex() {
                         MetricStat.MIN
                     );
                     assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics());
+
+                    assertEquals("_doc_count", starTreeFieldType.getMetrics().get(1).getField());
+                    assertEquals(List.of(MetricStat.DOC_COUNT), starTreeFieldType.getMetrics().get(1).getMetrics());
+
                     assertEquals(10000, starTreeFieldType.getStarTreeConfig().maxLeafDocs());
                     assertEquals(
                         StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP,

server/src/main/java/org/opensearch/index/codec/composite/Composite99DocValuesWriter.java

@@ -15,13 +15,15 @@
 import org.apache.lucene.index.EmptyDocValuesProducer;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.MergeState;
+import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SegmentWriteState;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.codec.composite.datacube.startree.StarTreeValues;
 import org.opensearch.index.compositeindex.datacube.startree.StarTreeField;
 import org.opensearch.index.compositeindex.datacube.startree.builder.StarTreesBuilder;
 import org.opensearch.index.mapper.CompositeMappedFieldType;
+import org.opensearch.index.mapper.DocCountFieldMapper;
 import org.opensearch.index.mapper.MapperService;
 
 import java.io.IOException;
@@ -60,21 +62,29 @@ public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState
         this.compositeMappedFieldTypes = mapperService.getCompositeFieldTypes();
         compositeFieldSet = new HashSet<>();
         segmentFieldSet = new HashSet<>();
+        // TODO : add integ test for this
         for (FieldInfo fi : segmentWriteState.fieldInfos) {
             if (DocValuesType.SORTED_NUMERIC.equals(fi.getDocValuesType())) {
                 segmentFieldSet.add(fi.name);
+            } else if (fi.name.equals(DocCountFieldMapper.NAME)) {
+                segmentFieldSet.add(fi.name);
             }
         }
         for (CompositeMappedFieldType type : compositeMappedFieldTypes) {
             compositeFieldSet.addAll(type.fields());
         }
         // check if there are any composite fields which are part of the segment
+        // TODO : add integ test where there are no composite fields in a segment, test both flush and merge cases
         segmentHasCompositeFields = Collections.disjoint(segmentFieldSet, compositeFieldSet) == false;
     }
 
     @Override
     public void addNumericField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException {
         delegate.addNumericField(field, valuesProducer);
+        // Perform this only during flush flow
+        if (mergeState.get() == null && segmentHasCompositeFields) {
+            createCompositeIndicesIfPossible(valuesProducer, field);
+        }
     }
 
     @Override
@@ -116,13 +126,7 @@ private void createCompositeIndicesIfPossible(DocValuesProducer valuesProducer,
         if (segmentFieldSet.isEmpty()) {
             Set<String> compositeFieldSetCopy = new HashSet<>(compositeFieldSet);
             for (String compositeField : compositeFieldSetCopy) {
-                fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() {
-                    @Override
-                    public SortedNumericDocValues getSortedNumeric(FieldInfo field) {
-                        return DocValues.emptySortedNumeric();
-                    }
-                });
-                compositeFieldSet.remove(compositeField);
+                addDocValuesForEmptyField(compositeField);
             }
         }
         // we have all the required fields to build composite fields
@@ -135,7 +139,28 @@ public SortedNumericDocValues getSortedNumeric(FieldInfo field) {
                 }
             }
         }
+    }
 
+    /**
+     * Add empty doc values for fields not present in segment
+     */
+    private void addDocValuesForEmptyField(String compositeField) {
+        if (compositeField.equals(DocCountFieldMapper.NAME)) {
+            fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() {
+                @Override
+                public NumericDocValues getNumeric(FieldInfo field) {
+                    return DocValues.emptyNumeric();
+                }
+            });
+        } else {
+            fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() {
+                @Override
+                public SortedNumericDocValues getSortedNumeric(FieldInfo field) {
+                    return DocValues.emptySortedNumeric();
+                }
+            });
+        }
+        compositeFieldSet.remove(compositeField);
     }
 
     @Override

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

@@ -21,12 +21,21 @@ public enum MetricStat {
     AVG("avg"),
     SUM("sum"),
     MIN("min"),
-    MAX("max");
+    MAX("max"),
+    DOC_COUNT("doc_count", true);
 
     private final String typeName;
 
+    // System field stats cannot be used as input for user metric types
+    private final boolean isSystemFieldStat;
+
     MetricStat(String typeName) {
+        this(typeName, false);
+    }
+
+    MetricStat(String typeName, boolean isSystemFieldStat) {
         this.typeName = typeName;
+        this.isSystemFieldStat = isSystemFieldStat;
     }
 
     public String getTypeName() {
@@ -35,7 +44,8 @@ public String getTypeName() {
 
     public static MetricStat fromTypeName(String typeName) {
         for (MetricStat metric : MetricStat.values()) {
-            if (metric.getTypeName().equalsIgnoreCase(typeName)) {
+            // prevent system fields to be entered as user input
+            if (metric.getTypeName().equalsIgnoreCase(typeName) && metric.isSystemFieldStat == false) {
                 return metric;
             }
         }

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

@@ -14,6 +14,7 @@
 import org.opensearch.index.compositeindex.datacube.Dimension;
 import org.opensearch.index.compositeindex.datacube.Metric;
 import org.opensearch.index.mapper.CompositeMappedFieldType;
+import org.opensearch.index.mapper.DocCountFieldMapper;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.StarTreeMapper;
@@ -78,7 +79,7 @@ public static void validate(MapperService mapperService, CompositeIndexSettings
                         String.format(Locale.ROOT, "unknown metric field [%s] as part of star tree field", metric.getField())
                     );
                 }
-                if (ft.isAggregatable() == false) {
+                if (ft.isAggregatable() == false && ft instanceof DocCountFieldMapper.DocCountFieldType == false) {
                     throw new IllegalArgumentException(
                         String.format(
                             Locale.ROOT,

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

@@ -17,12 +17,9 @@
 class CountValueAggregator implements ValueAggregator<Long> {
 
     public static final long DEFAULT_INITIAL_VALUE = 1L;
-    private final StarTreeNumericType starTreeNumericType;
     private static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.LONG;
 
-    public CountValueAggregator(StarTreeNumericType starTreeNumericType) {
-        this.starTreeNumericType = starTreeNumericType;
-    }
+    public CountValueAggregator() {}
 
     @Override
     public StarTreeNumericType getAggregatedValueType() {

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

@@ -0,0 +1,70 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.compositeindex.datacube.startree.aggregators;
+
+import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType;
+
+/**
+ * Aggregator to handle '_doc_count' field
+ *
+ * @opensearch.experimental
+ */
+public class DocCountAggregator implements ValueAggregator<Long> {
+
+    private static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.LONG;
+
+    public DocCountAggregator() {}
+
+    @Override
+    public StarTreeNumericType getAggregatedValueType() {
+        return VALUE_AGGREGATOR_TYPE;
+    }
+
+    /**
+     * If _doc_count field for a doc is missing, we increment the _doc_count by '1' for the associated doc
+     * otherwise take the actual value present in the field
+     */
+    @Override
+    public Long getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue) {
+        if (segmentDocValue == null) {
+            return getIdentityMetricValue();
+        }
+        return segmentDocValue;
+    }
+
+    @Override
+    public Long mergeAggregatedValueAndSegmentValue(Long value, Long segmentDocValue) {
+        assert value != null;
+        return mergeAggregatedValues(value, segmentDocValue);
+    }
+
+    @Override
+    public Long mergeAggregatedValues(Long value, Long aggregatedValue) {
+        if (value == null) {
+            value = getIdentityMetricValue();
+        }
+        if (aggregatedValue == null) {
+            aggregatedValue = getIdentityMetricValue();
+        }
+        return value + aggregatedValue;
+    }
+
+    @Override
+    public Long toAggregatedValueType(Long rawValue) {
+        return rawValue;
+    }
+
+    /**
+     * If _doc_count field for a doc is missing, we increment the _doc_count by '1' for the associated doc
+     */
+    @Override
+    public Long getIdentityMetricValue() {
+        return 1L;
+    }
+}

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

@@ -31,11 +31,13 @@ public static ValueAggregator getValueAggregator(MetricStat aggregationType, Sta
             case SUM:
                 return new SumValueAggregator(starTreeNumericType);
             case VALUE_COUNT:
-                return new CountValueAggregator(starTreeNumericType);
+                return new CountValueAggregator();
             case MIN:
                 return new MinValueAggregator(starTreeNumericType);
             case MAX:
                 return new MaxValueAggregator(starTreeNumericType);
+            case DOC_COUNT:
+                return new DocCountAggregator();
             default:
                 throw new IllegalStateException("Unsupported aggregation type: " + aggregationType);
         }

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

@@ -10,6 +10,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.codecs.DocValuesProducer;
+import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.IndexOptions;
@@ -28,6 +29,7 @@
 import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator;
 import org.opensearch.index.compositeindex.datacube.startree.utils.TreeNode;
 import org.opensearch.index.fielddata.IndexNumericFieldData;
+import org.opensearch.index.mapper.DocCountFieldMapper;
 import org.opensearch.index.mapper.Mapper;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.NumberFieldMapper;
@@ -117,6 +119,16 @@ protected BaseStarTreeBuilder(StarTreeField starTreeField, SegmentWriteState sta
     public List<MetricAggregatorInfo> generateMetricAggregatorInfos(MapperService mapperService) {
         List<MetricAggregatorInfo> metricAggregatorInfos = new ArrayList<>();
         for (Metric metric : this.starTreeField.getMetrics()) {
+            if (metric.getField().equals(DocCountFieldMapper.NAME)) {
+                MetricAggregatorInfo metricAggregatorInfo = new MetricAggregatorInfo(
+                    MetricStat.DOC_COUNT,
+                    metric.getField(),
+                    starTreeField.getName(),
+                    IndexNumericFieldData.NumericType.LONG
+                );
+                metricAggregatorInfos.add(metricAggregatorInfo);
+                continue;
+            }
             for (MetricStat metricStat : metric.getMetrics()) {
                 IndexNumericFieldData.NumericType numericType;
                 Mapper fieldMapper = mapperService.documentMapper().mappers().getMapper(metric.getField());
@@ -426,7 +438,7 @@ public void build(Map<String, DocValuesProducer> fieldProducerMap) throws IOExce
             String dimension = dimensionsSplitOrder.get(i).getField();
             FieldInfo dimensionFieldInfo = state.fieldInfos.fieldInfo(dimension);
             if (dimensionFieldInfo == null) {
-                dimensionFieldInfo = getFieldInfo(dimension);
+                dimensionFieldInfo = getFieldInfo(dimension, DocValuesType.SORTED_NUMERIC);
             }
             dimensionReaders[i] = new SequentialDocValuesIterator(
                 fieldProducerMap.get(dimensionFieldInfo.name).getSortedNumeric(dimensionFieldInfo)
@@ -438,15 +450,15 @@ public void build(Map<String, DocValuesProducer> fieldProducerMap) throws IOExce
         logger.debug("Finished Building star-tree in ms : {}", (System.currentTimeMillis() - startTime));
     }
 
-    private static FieldInfo getFieldInfo(String field) {
+    private static FieldInfo getFieldInfo(String field, DocValuesType docValuesType) {
         return new FieldInfo(
             field,
-            1,
+            1, // This is filled as part of doc values creation and is not used otherwise
             false,
             false,
             false,
             IndexOptions.NONE,
-            DocValuesType.SORTED_NUMERIC,
+            docValuesType,
             -1,
             Collections.emptyMap(),
             0,
@@ -470,20 +482,38 @@ public List<SequentialDocValuesIterator> getMetricReaders(SegmentWriteState stat
         List<SequentialDocValuesIterator> metricReaders = new ArrayList<>();
         for (Metric metric : this.starTreeField.getMetrics()) {
             for (MetricStat metricStat : metric.getMetrics()) {
+                SequentialDocValuesIterator metricReader = null;
                 FieldInfo metricFieldInfo = state.fieldInfos.fieldInfo(metric.getField());
-                if (metricFieldInfo == null) {
-                    metricFieldInfo = getFieldInfo(metric.getField());
+                if (metricStat.equals(MetricStat.DOC_COUNT)) {
+                    metricReader = getDocCountMetricReader(fieldProducerMap, metricFieldInfo);
+                } else {
+                    if (metricFieldInfo == null) {
+                        metricFieldInfo = getFieldInfo(metric.getField(), DocValuesType.SORTED_NUMERIC);
+                    }
+                    metricReader = new SequentialDocValuesIterator(
+                        fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo)
+                    );
                 }
-
-                SequentialDocValuesIterator metricReader = new SequentialDocValuesIterator(
-                    fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo)
-                );
                 metricReaders.add(metricReader);
             }
         }
         return metricReaders;
     }
 
+    private SequentialDocValuesIterator getDocCountMetricReader(Map<String, DocValuesProducer> fieldProducerMap, FieldInfo metricFieldInfo)
+        throws IOException {
+        if (metricFieldInfo == null) {
+            metricFieldInfo = getFieldInfo(DocCountFieldMapper.NAME, DocValuesType.NUMERIC);
+        }
+        SequentialDocValuesIterator metricReader;
+        assert fieldProducerMap.containsKey(metricFieldInfo.name);
+        // _doc_count is numeric field , so we need to get convert to sortedNumericDocValues
+        metricReader = new SequentialDocValuesIterator(
+            DocValues.singleton(fieldProducerMap.get(metricFieldInfo.name).getNumeric(metricFieldInfo))
+        );
+        return metricReader;
+    }
+
     /**
      * Builds the star tree using Star-Tree Document
      *

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

@@ -225,6 +225,10 @@ private List<Metric> buildMetrics(String fieldName, Map<String, Object> map, Map
                 for (Object metric : metricsList) {
                     Map<String, Object> metricMap = (Map<String, Object>) metric;
                     String name = (String) XContentMapValues.extractValue(CompositeDataCubeFieldType.NAME, metricMap);
+                    // Handle _doc_count metric separately at the end
+                    if (name.equals(DocCountFieldMapper.NAME)) {
+                        continue;
+                    }
                     metricMap.remove(CompositeDataCubeFieldType.NAME);
                     if (objbuilder == null || objbuilder.mappersBuilders == null) {
                         metrics.add(getMetric(name, metricMap, context));
@@ -249,7 +253,8 @@ private List<Metric> buildMetrics(String fieldName, Map<String, Object> map, Map
             } else {
                 throw new MapperParsingException(String.format(Locale.ROOT, "unable to parse metrics for star tree field [%s]", this.name));
             }
-
+            Metric docCountMetric = new Metric(DocCountFieldMapper.NAME, List.of(MetricStat.DOC_COUNT));
+            metrics.add(docCountMetric);
             return metrics;
         }