[Star tree] Star tree merge changes (opensearch-project#14652)

---------

Signed-off-by: Bharathwaj G <bharath78910@gmail.com>

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

@@ -17,9 +17,9 @@
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.index.SortedSetDocValues;
 import org.opensearch.common.annotation.ExperimentalApi;
-import org.opensearch.index.mapper.CompositeMappedFieldType;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -74,15 +74,13 @@ public void close() throws IOException {
     }
 
     @Override
-    public List<String> getCompositeIndexFields() {
+    public List<CompositeIndexFieldInfo> getCompositeIndexFields() {
         // todo : read from file formats and get the field names.
-        throw new UnsupportedOperationException();
-
+        return new ArrayList<>();
     }
 
     @Override
-    public CompositeIndexValues getCompositeIndexValues(String field, CompositeMappedFieldType.CompositeFieldType fieldType)
-        throws IOException {
+    public CompositeIndexValues getCompositeIndexValues(CompositeIndexFieldInfo compositeIndexFieldInfo) throws IOException {
         // TODO : read compositeIndexValues [starTreeValues] from star tree files
         throw new UnsupportedOperationException();
     }

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

@@ -8,20 +8,29 @@
 
 package org.opensearch.index.codec.composite;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.apache.lucene.codecs.DocValuesConsumer;
 import org.apache.lucene.codecs.DocValuesProducer;
+import org.apache.lucene.index.DocValues;
+import org.apache.lucene.index.DocValuesType;
+import org.apache.lucene.index.EmptyDocValuesProducer;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.MergeState;
 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.MapperService;
-import org.opensearch.index.mapper.StarTreeMapper;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
@@ -40,8 +49,10 @@ public class Composite99DocValuesWriter extends DocValuesConsumer {
     AtomicReference<MergeState> mergeState = new AtomicReference<>();
     private final Set<CompositeMappedFieldType> compositeMappedFieldTypes;
     private final Set<String> compositeFieldSet;
+    private final Set<String> segmentFieldSet;
 
     private final Map<String, DocValuesProducer> fieldProducerMap = new HashMap<>();
+    private static final Logger logger = LogManager.getLogger(Composite99DocValuesWriter.class);
 
     public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState segmentWriteState, MapperService mapperService) {
 
@@ -50,6 +61,12 @@ public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState
         this.mapperService = mapperService;
         this.compositeMappedFieldTypes = mapperService.getCompositeFieldTypes();
         compositeFieldSet = new HashSet<>();
+        segmentFieldSet = new HashSet<>();
+        for (FieldInfo fi : segmentWriteState.fieldInfos) {
+            if (DocValuesType.SORTED_NUMERIC.equals(fi.getDocValuesType())) {
+                segmentFieldSet.add(fi.name);
+            }
+        }
         for (CompositeMappedFieldType type : compositeMappedFieldTypes) {
             compositeFieldSet.addAll(type.fields());
         }
@@ -95,23 +112,91 @@ private void createCompositeIndicesIfPossible(DocValuesProducer valuesProducer,
             fieldProducerMap.put(field.name, valuesProducer);
             compositeFieldSet.remove(field.name);
         }
+        segmentFieldSet.remove(field.name);
+        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);
+            }
+        }
         // we have all the required fields to build composite fields
         if (compositeFieldSet.isEmpty()) {
             for (CompositeMappedFieldType mappedType : compositeMappedFieldTypes) {
-                if (mappedType instanceof StarTreeMapper.StarTreeFieldType) {
-                    try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(fieldProducerMap, state, mapperService)) {
-                        starTreesBuilder.build();
+                if (mappedType.getCompositeIndexType().equals(CompositeMappedFieldType.CompositeFieldType.STAR_TREE)) {
+                    try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(state, mapperService)) {
+                        starTreesBuilder.build(fieldProducerMap);
                     }
                 }
             }
         }
+
     }
 
     @Override
     public void merge(MergeState mergeState) throws IOException {
         this.mergeState.compareAndSet(null, mergeState);
         super.merge(mergeState);
-        // TODO : handle merge star tree
-        // mergeStarTreeFields(mergeState);
+        mergeCompositeFields(mergeState);
+    }
+
+    /**
+     * Merges composite fields from multiple segments
+     * @param mergeState merge state
+     */
+    private void mergeCompositeFields(MergeState mergeState) throws IOException {
+        mergeStarTreeFields(mergeState);
+    }
+
+    /**
+     * Merges star tree data fields from multiple segments
+     * @param mergeState merge state
+     */
+    private void mergeStarTreeFields(MergeState mergeState) throws IOException {
+        Map<String, List<StarTreeValues>> starTreeSubsPerField = new HashMap<>();
+        StarTreeField starTreeField = null;
+        for (int i = 0; i < mergeState.docValuesProducers.length; i++) {
+            CompositeIndexReader reader = null;
+            if (mergeState.docValuesProducers[i] == null) {
+                continue;
+            }
+            if (mergeState.docValuesProducers[i] instanceof CompositeIndexReader) {
+                reader = (CompositeIndexReader) mergeState.docValuesProducers[i];
+            } else {
+                continue;
+            }
+
+            List<CompositeIndexFieldInfo> compositeFieldInfo = reader.getCompositeIndexFields();
+            for (CompositeIndexFieldInfo fieldInfo : compositeFieldInfo) {
+                if (fieldInfo.getType().equals(CompositeMappedFieldType.CompositeFieldType.STAR_TREE)) {
+                    CompositeIndexValues compositeIndexValues = reader.getCompositeIndexValues(fieldInfo);
+                    if (compositeIndexValues instanceof StarTreeValues) {
+                        StarTreeValues starTreeValues = (StarTreeValues) compositeIndexValues;
+                        List<StarTreeValues> fieldsList = starTreeSubsPerField.getOrDefault(fieldInfo.getField(), Collections.emptyList());
+                        if (starTreeField == null) {
+                            starTreeField = starTreeValues.getStarTreeField();
+                        }
+                        // assert star tree configuration is same across segments
+                        else {
+                            if (starTreeField.equals(starTreeValues.getStarTreeField()) == false) {
+                                throw new IllegalArgumentException(
+                                    "star tree field configuration must match the configuration of the field being merged"
+                                );
+                            }
+                        }
+                        fieldsList.add(starTreeValues);
+                        starTreeSubsPerField.put(fieldInfo.getField(), fieldsList);
+                    }
+                }
+            }
+        }
+        try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(state, mapperService)) {
+            starTreesBuilder.buildDuringMerge(starTreeSubsPerField);
+        }
     }
 }

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

@@ -0,0 +1,37 @@
+
+/*
+ * 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.codec.composite;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.mapper.CompositeMappedFieldType;
+
+/**
+ * Field info details of composite index fields
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class CompositeIndexFieldInfo {
+    private final String field;
+    private final CompositeMappedFieldType.CompositeFieldType type;
+
+    public CompositeIndexFieldInfo(String field, CompositeMappedFieldType.CompositeFieldType type) {
+        this.field = field;
+        this.type = type;
+    }
+
+    public String getField() {
+        return field;
+    }
+
+    public CompositeMappedFieldType.CompositeFieldType getType() {
+        return type;
+    }
+}

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

@@ -9,7 +9,6 @@
 package org.opensearch.index.codec.composite;
 
 import org.opensearch.common.annotation.ExperimentalApi;
-import org.opensearch.index.mapper.CompositeMappedFieldType;
 
 import java.io.IOException;
 import java.util.List;
@@ -25,10 +24,10 @@ public interface CompositeIndexReader {
      * Get list of composite index fields from the segment
      *
      */
-    List<String> getCompositeIndexFields();
+    List<CompositeIndexFieldInfo> getCompositeIndexFields();
 
     /**
      * Get composite index values based on the field name and the field type
      */
-    CompositeIndexValues getCompositeIndexValues(String field, CompositeMappedFieldType.CompositeFieldType fieldType) throws IOException;
+    CompositeIndexValues getCompositeIndexValues(CompositeIndexFieldInfo fieldInfo) throws IOException;
 }

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

@@ -8,10 +8,13 @@
 
 package org.opensearch.index.codec.composite.datacube.startree;
 
+import org.apache.lucene.search.DocIdSetIterator;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.codec.composite.CompositeIndexValues;
+import org.opensearch.index.compositeindex.datacube.startree.StarTreeField;
+import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
 
-import java.util.List;
+import java.util.Map;
 
 /**
  * Concrete class that holds the star tree associated values from the segment
@@ -20,16 +23,48 @@
  */
 @ExperimentalApi
 public class StarTreeValues implements CompositeIndexValues {
-    private final List<String> dimensionsOrder;
+    private final StarTreeField starTreeField;
+    private final StarTreeNode root;
+    private final Map<String, DocIdSetIterator> dimensionDocValuesIteratorMap;
+    private final Map<String, DocIdSetIterator> metricDocValuesIteratorMap;
+    private final Map<String, String> attributes;
 
-    // TODO : come up with full set of vales such as dimensions and metrics doc values + star tree
-    public StarTreeValues(List<String> dimensionsOrder) {
-        super();
-        this.dimensionsOrder = List.copyOf(dimensionsOrder);
+    public StarTreeValues(
+        StarTreeField starTreeField,
+        StarTreeNode root,
+        Map<String, DocIdSetIterator> dimensionDocValuesIteratorMap,
+        Map<String, DocIdSetIterator> metricDocValuesIteratorMap,
+        Map<String, String> attributes
+    ) {
+        this.starTreeField = starTreeField;
+        this.root = root;
+        this.dimensionDocValuesIteratorMap = dimensionDocValuesIteratorMap;
+        this.metricDocValuesIteratorMap = metricDocValuesIteratorMap;
+        this.attributes = attributes;
     }
 
     @Override
     public CompositeIndexValues getValues() {
         return this;
     }
+
+    public StarTreeField getStarTreeField() {
+        return starTreeField;
+    }
+
+    public StarTreeNode getRoot() {
+        return root;
+    }
+
+    public Map<String, DocIdSetIterator> getDimensionDocValuesIteratorMap() {
+        return dimensionDocValuesIteratorMap;
+    }
+
+    public Map<String, DocIdSetIterator> getMetricDocValuesIteratorMap() {
+        return metricDocValuesIteratorMap;
+    }
+
+    public Map<String, String> getAttributes() {
+        return attributes;
+    }
 }

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

@@ -18,6 +18,11 @@
 public class CountValueAggregator implements ValueAggregator<Long> {
     public static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.LONG;
     public static final long DEFAULT_INITIAL_VALUE = 1L;
+    private StarTreeNumericType starTreeNumericType;
+
+    public CountValueAggregator(StarTreeNumericType starTreeNumericType) {
+        this.starTreeNumericType = starTreeNumericType;
+    }
 
     @Override
     public MetricStat getAggregationType() {
@@ -30,12 +35,12 @@ public StarTreeNumericType getAggregatedValueType() {
     }
 
     @Override
-    public Long getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue, StarTreeNumericType starTreeNumericType) {
+    public Long getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue) {
         return DEFAULT_INITIAL_VALUE;
     }
 
     @Override
-    public Long mergeAggregatedValueAndSegmentValue(Long value, Long segmentDocValue, StarTreeNumericType starTreeNumericType) {
+    public Long mergeAggregatedValueAndSegmentValue(Long value, Long segmentDocValue) {
         return value + 1;
     }
 
@@ -60,7 +65,7 @@ public Long toLongValue(Long value) {
     }
 
     @Override
-    public Long toStarTreeNumericTypeValue(Long value, StarTreeNumericType type) {
+    public Long toStarTreeNumericTypeValue(Long value) {
         return value;
     }
 }