addressing review comments

Signed-off-by: Sarthak Aggarwal <sarthagg@amazon.com>

server/src/main/java/org/apache/lucene/index/BaseStarTreeBuilder.java

@@ -23,7 +23,7 @@
 import org.opensearch.index.compositeindex.datacube.startree.builder.StarTreeDocValuesIteratorAdapter;
 import org.opensearch.index.compositeindex.datacube.startree.builder.StarTreesBuilder;
 import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeBuilderUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.TreeNode;
 import org.opensearch.index.fielddata.IndexNumericFieldData;
 import org.opensearch.index.mapper.Mapper;
 import org.opensearch.index.mapper.MapperService;
@@ -39,6 +39,8 @@
 import java.util.Map;
 import java.util.Set;
 
+import static org.opensearch.index.compositeindex.datacube.startree.utils.TreeNode.ALL;
+
 /**
  * Builder for star tree. Defines the algorithm to construct star-tree
  * See {@link StarTreesBuilder} for information around the construction of star-trees based on star-tree fields
@@ -64,7 +66,7 @@ public abstract class BaseStarTreeBuilder implements StarTreeBuilder {
     protected int numStarTreeNodes;
     protected final int maxLeafDocuments;
 
-    protected final StarTreeBuilderUtils.TreeNode rootNode = getNewNode();
+    protected final TreeNode rootNode = getNewNode();
 
     protected SequentialDocValuesIterator[] dimensionReaders;
 
@@ -237,7 +239,7 @@ protected StarTreeDocument getSegmentStarTreeDocument(int currentDocId) throws I
      * @return dimension values for each of the star-tree dimension
      * @throws IOException when we are unable to iterate to the next doc for the given dimension readers
      */
-    Long[] getStarTreeDimensionsFromSegment(int currentDocId) throws IOException {
+    private Long[] getStarTreeDimensionsFromSegment(int currentDocId) throws IOException {
         Long[] dimensions = new Long[numDimensions];
         for (int i = 0; i < numDimensions; i++) {
             if (dimensionReaders[i] != null) {
@@ -470,9 +472,9 @@ private void appendToStarTree(StarTreeDocument starTreeDocument) throws IOExcept
      *
      * @return return new star-tree node
      */
-    private StarTreeBuilderUtils.TreeNode getNewNode() {
+    private TreeNode getNewNode() {
         numStarTreeNodes++;
-        return new StarTreeBuilderUtils.TreeNode();
+        return new TreeNode();
     }
 
     /**
@@ -483,7 +485,7 @@ private StarTreeBuilderUtils.TreeNode getNewNode() {
      * @param endDocId   end document id
      * @throws IOException throws an exception if we are unable to construct the tree
      */
-    private void constructStarTree(StarTreeBuilderUtils.TreeNode node, int startDocId, int endDocId) throws IOException {
+    private void constructStarTree(TreeNode node, int startDocId, int endDocId) throws IOException {
 
         int childDimensionId = node.dimensionId + 1;
         if (childDimensionId == numDimensions) {
@@ -492,16 +494,16 @@ private void constructStarTree(StarTreeBuilderUtils.TreeNode node, int startDocI
 
         // Construct all non-star children nodes
         node.childDimensionId = childDimensionId;
-        Map<Long, StarTreeBuilderUtils.TreeNode> children = constructNonStarNodes(startDocId, endDocId, childDimensionId);
+        Map<Long, TreeNode> children = constructNonStarNodes(startDocId, endDocId, childDimensionId);
         node.children = children;
 
         // Construct star-node if required
         if (!skipStarNodeCreationForDimensions.contains(childDimensionId) && children.size() > 1) {
-            children.put((long) StarTreeBuilderUtils.ALL, constructStarNode(startDocId, endDocId, childDimensionId));
+            children.put((long) ALL, constructStarNode(startDocId, endDocId, childDimensionId));
         }
 
         // Further split on child nodes if required
-        for (StarTreeBuilderUtils.TreeNode child : children.values()) {
+        for (TreeNode child : children.values()) {
             if (child.endDocId - child.startDocId > maxLeafDocuments) {
                 constructStarTree(child, child.startDocId, child.endDocId);
             }
@@ -517,15 +519,14 @@ private void constructStarTree(StarTreeBuilderUtils.TreeNode node, int startDocI
      * @return root node with non-star nodes constructed
      * @throws IOException throws an exception if we are unable to construct non-star nodes
      */
-    private Map<Long, StarTreeBuilderUtils.TreeNode> constructNonStarNodes(int startDocId, int endDocId, int dimensionId)
-        throws IOException {
-        Map<Long, StarTreeBuilderUtils.TreeNode> nodes = new HashMap<>();
+    private Map<Long, TreeNode> constructNonStarNodes(int startDocId, int endDocId, int dimensionId) throws IOException {
+        Map<Long, TreeNode> nodes = new HashMap<>();
         int nodeStartDocId = startDocId;
         Long nodeDimensionValue = getDimensionValue(startDocId, dimensionId);
         for (int i = startDocId + 1; i < endDocId; i++) {
             Long dimensionValue = getDimensionValue(i, dimensionId);
             if (!dimensionValue.equals(nodeDimensionValue)) {
-                StarTreeBuilderUtils.TreeNode child = getNewNode();
+                TreeNode child = getNewNode();
                 child.dimensionId = dimensionId;
                 child.dimensionValue = nodeDimensionValue;
                 child.startDocId = nodeStartDocId;
@@ -536,7 +537,7 @@ private Map<Long, StarTreeBuilderUtils.TreeNode> constructNonStarNodes(int start
                 nodeDimensionValue = dimensionValue;
             }
         }
-        StarTreeBuilderUtils.TreeNode lastNode = getNewNode();
+        TreeNode lastNode = getNewNode();
         lastNode.dimensionId = dimensionId;
         lastNode.dimensionValue = nodeDimensionValue;
         lastNode.startDocId = nodeStartDocId;
@@ -554,10 +555,10 @@ private Map<Long, StarTreeBuilderUtils.TreeNode> constructNonStarNodes(int start
      * @return root node with star nodes constructed
      * @throws IOException throws an exception if we are unable to construct non-star nodes
      */
-    private StarTreeBuilderUtils.TreeNode constructStarNode(int startDocId, int endDocId, int dimensionId) throws IOException {
-        StarTreeBuilderUtils.TreeNode starNode = getNewNode();
+    private TreeNode constructStarNode(int startDocId, int endDocId, int dimensionId) throws IOException {
+        TreeNode starNode = getNewNode();
         starNode.dimensionId = dimensionId;
-        starNode.dimensionValue = StarTreeBuilderUtils.ALL;
+        starNode.dimensionValue = ALL;
         starNode.isStarNode = true;
         starNode.startDocId = numStarTreeDocs;
         Iterator<StarTreeDocument> starTreeDocumentIterator = generateStarTreeDocumentsForStarNode(startDocId, endDocId, dimensionId);
@@ -575,11 +576,11 @@ private StarTreeBuilderUtils.TreeNode constructStarNode(int startDocId, int endD
      * @return aggregated star-tree documents
      * @throws IOException throws an exception upon failing to create new aggregated docs based on star tree
      */
-    private StarTreeDocument createAggregatedDocs(StarTreeBuilderUtils.TreeNode node) throws IOException {
+    private StarTreeDocument createAggregatedDocs(TreeNode node) throws IOException {
         StarTreeDocument aggregatedStarTreeDocument = null;
         if (node.children == null) {
-            // For leaf node
 
+            // For leaf node
             if (node.startDocId == node.endDocId - 1) {
                 // If it has only one document, use it as the aggregated document
                 aggregatedStarTreeDocument = getStarTreeDocument(node.startDocId);
@@ -600,9 +601,9 @@ private StarTreeDocument createAggregatedDocs(StarTreeBuilderUtils.TreeNode node
             }
         } else {
             // For non-leaf node
-            if (node.children.containsKey((long) StarTreeBuilderUtils.ALL)) {
+            if (node.children.containsKey((long) ALL)) {
                 // If it has star child, use the star child aggregated document directly
-                for (StarTreeBuilderUtils.TreeNode child : node.children.values()) {
+                for (TreeNode child : node.children.values()) {
                     if (child.isStarNode) {
                         aggregatedStarTreeDocument = createAggregatedDocs(child);
                         node.aggregatedDocId = child.aggregatedDocId;
@@ -612,7 +613,7 @@ private StarTreeDocument createAggregatedDocs(StarTreeBuilderUtils.TreeNode node
                 }
             } else {
                 // If no star child exists, aggregate all aggregated documents from non-star children
-                for (StarTreeBuilderUtils.TreeNode child : node.children.values()) {
+                for (TreeNode child : node.children.values()) {
                     aggregatedStarTreeDocument = reduceStarTreeDocuments(aggregatedStarTreeDocument, createAggregatedDocs(child));
                 }
                 if (null == aggregatedStarTreeDocument) {

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

@@ -14,6 +14,7 @@
 import org.apache.lucene.index.MergeState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.opensearch.common.annotation.ExperimentalApi;
+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;
@@ -98,7 +99,9 @@ private void createCompositeIndicesIfPossible(DocValuesProducer valuesProducer,
         if (compositeFieldSet.isEmpty()) {
             for (CompositeMappedFieldType mappedType : compositeMappedFieldTypes) {
                 if (mappedType instanceof StarTreeMapper.StarTreeFieldType) {
-                    // TODO : Call StarTree builder
+                    StarTreesBuilder starTreesBuilder = new StarTreesBuilder(fieldProducerMap, state, mapperService);
+                    starTreesBuilder.build();
+                    starTreesBuilder.close();
                 }
             }
         }

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

@@ -17,6 +17,7 @@
  */
 public class CountValueAggregator implements ValueAggregator<Long> {
     public static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.LONG;
+    public static final long DEFAULT_INITIAL_VALUE = 1L;
 
     @Override
     public MetricStat getAggregationType() {
@@ -30,7 +31,7 @@ public StarTreeNumericType getAggregatedValueType() {
 
     @Override
     public Long getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue, StarTreeNumericType starTreeNumericType) {
-        return 1L;
+        return DEFAULT_INITIAL_VALUE;
     }
 
     @Override

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

@@ -86,17 +86,10 @@ public Iterator<StarTreeDocument> sortAndAggregateStarTreeDocuments(int numDocs)
      * @return iterator for star-tree documents
      * @throws IOException throws when unable to sort, merge and aggregate star-tree documents
      */
-    public Iterator<StarTreeDocument> sortAndAggregateStarTreeDocuments(StarTreeDocument[] starTreeDocuments) throws IOException {
+    Iterator<StarTreeDocument> sortAndAggregateStarTreeDocuments(StarTreeDocument[] starTreeDocuments) throws IOException {
 
         // sort the documents
-        Arrays.sort(starTreeDocuments, (o1, o2) -> {
-            for (int i = 0; i < numDimensions; i++) {
-                if (!Objects.equals(o1.dimensions[i], o2.dimensions[i])) {
-                    return Long.compare(o1.dimensions[i], o2.dimensions[i]);
-                }
-            }
-            return 0;
-        });
+        Arrays.sort(starTreeDocuments, (o1, o2) -> Arrays.compare(o1.dimensions, o2.dimensions));
 
         // merge the documents
         return mergeStarTreeDocuments(starTreeDocuments);
@@ -153,14 +146,10 @@ public Iterator<StarTreeDocument> generateStarTreeDocumentsForStarNode(int start
         for (int i = 0; i < numDocs; i++) {
             starTreeDocuments[i] = getStarTreeDocument(startDocId + i);
         }
-        Arrays.sort(starTreeDocuments, (o1, o2) -> {
-            for (int i = dimensionId + 1; i < numDimensions; i++) {
-                if (!Objects.equals(o1.dimensions[i], o2.dimensions[i])) {
-                    return Long.compare(o1.dimensions[i], o2.dimensions[i]);
-                }
-            }
-            return 0;
-        });
+
+        // sort the documents
+        Arrays.sort(starTreeDocuments, (o1, o2) -> Arrays.compare(o1.dimensions, o2.dimensions));
+
         return new Iterator<StarTreeDocument>() {
             boolean hasNext = true;
             StarTreeDocument currentStarTreeDocument = starTreeDocuments[0];

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

@@ -0,0 +1,65 @@
+/*
+ * 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.utils;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+
+import java.util.Map;
+
+/**
+ * /**
+ * Represents a node in a tree data structure, specifically designed for a star-tree implementation.
+ * A star-tree node will represent both star and non-star nodes.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class TreeNode {
+
+    public static final int ALL = -1;
+
+    /**
+     * The dimension id for the dimension (field) associated with this star-tree node.
+     */
+    public int dimensionId = ALL;
+
+    /**
+     * The starting document id (inclusive) associated with this star-tree node.
+     */
+    public int startDocId = ALL;
+
+    /**
+     * The ending document id (exclusive) associated with this star-tree node.
+     */
+    public int endDocId = ALL;
+
+    /**
+     * The aggregated document id associated with this star-tree node.
+     */
+    public int aggregatedDocId = ALL;
+
+    /**
+     * The child dimension identifier associated with this star-tree node.
+     */
+    public int childDimensionId = ALL;
+
+    /**
+     * The value of the dimension associated with this star-tree node.
+     */
+    public long dimensionValue = ALL;
+
+    /**
+     * A flag indicating whether this node is a star node (a node that represents an aggregation of all dimensions).
+     */
+    public boolean isStarNode = false;
+
+    /**
+     * A map containing the child nodes of this star-tree node, keyed by their dimension id.
+     */
+    public Map<Long, TreeNode> children;
+}