separating child star node from children

Signed-off-by: Sarthak Aggarwal <[email protected]>
opensearch-project · Aug 22, 2024 · 0d7768e · 0d7768e
1 parent f5336aa
commit 0d7768e
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 33 deletions.
diff --git a/...va/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/...va/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java
@@ -696,7 +696,12 @@ private void constructStarTree(InMemoryTreeNode node, int startDocId, int endDoc
 
         // Construct star-node if required
         if (!skipStarNodeCreationForDimensions.contains(childDimensionId) && children.size() > 1) {
-            children.put((long) ALL, constructStarNode(startDocId, endDocId, childDimensionId));
+            node.childStarNode = constructStarNode(startDocId, endDocId, childDimensionId);
+        }
+
+        // Further split star node if needed
+        if (node.childStarNode != null && (node.childStarNode.endDocId - node.childStarNode.startDocId > maxLeafDocuments)) {
+            constructStarTree(node.childStarNode, node.childStarNode.startDocId, node.childStarNode.endDocId);
         }
 
         // Further split on child nodes if required
@@ -705,6 +710,7 @@ private void constructStarTree(InMemoryTreeNode node, int startDocId, int endDoc
                 constructStarTree(child, child.startDocId, child.endDocId);
             }
         }
+
     }
 
     /**
@@ -778,9 +784,10 @@ private InMemoryTreeNode constructStarNode(int startDocId, int endDocId, int dim
      */
     private StarTreeDocument createAggregatedDocs(InMemoryTreeNode node) throws IOException {
         StarTreeDocument aggregatedStarTreeDocument = null;
-        if (node.children == null) {
 
-            // For leaf node
+        // For leaf node
+        if (node.children == null && node.childStarNode == null) {
+
             if (node.startDocId == node.endDocId - 1) {
                 // If it has only one document, use it as the aggregated document
                 aggregatedStarTreeDocument = getStarTreeDocument(node.startDocId);
@@ -801,15 +808,13 @@ private StarTreeDocument createAggregatedDocs(InMemoryTreeNode node) throws IOEx
             }
         } else {
             // For non-leaf node
-            if (node.children.containsKey((long) ALL)) {
+            if (node.childStarNode != null) {
                 // If it has star child, use the star child aggregated document directly
+                aggregatedStarTreeDocument = createAggregatedDocs(node.childStarNode);
+                node.aggregatedDocId = node.childStarNode.aggregatedDocId;
+
                 for (InMemoryTreeNode child : node.children.values()) {
-                    if (child.nodeType == StarTreeNodeType.STAR.getValue()) {
-                        aggregatedStarTreeDocument = createAggregatedDocs(child);
-                        node.aggregatedDocId = child.aggregatedDocId;
-                    } else {
-                        createAggregatedDocs(child);
-                    }
+                    createAggregatedDocs(child);
                 }
             } else {
                 // If no star child exists, aggregate all aggregated documents from non-star children

diff --git a/...pensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeDataWriter.java b/...pensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeDataWriter.java
@@ -26,7 +26,7 @@
 import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.ALL;
 
 /**
- * Utility class for serializing a star-tree data structure.
+ * Utility class for serializing a star-tree data structure.x
  *
  * @opensearch.experimental
  */
@@ -98,21 +98,32 @@ private static void writeStarTreeNodes(IndexOutput output, InMemoryTreeNode root
         while (!queue.isEmpty()) {
             InMemoryTreeNode node = queue.remove();
 
-            if (node.children == null || node.children.isEmpty()) {
+            if ((node.children == null || node.children.isEmpty()) && node.childStarNode == null) {
                 writeStarTreeNode(output, node, ALL, ALL);
             } else {
 
-                // Sort all children nodes based on dimension value
-                List<InMemoryTreeNode> sortedChildren = new ArrayList<>(node.children.values());
-                sortedChildren.sort(
-                    Comparator.comparingInt(InMemoryTreeNode::getNodeType).thenComparingLong(InMemoryTreeNode::getDimensionValue)
-                );
-
+                int totalNumberOfChildren = 0;
                 int firstChildId = currentNodeId + queue.size() + 1;
-                int lastChildId = firstChildId + sortedChildren.size() - 1;
+
+                if (node.childStarNode != null) {
+                    totalNumberOfChildren++;
+                    queue.add(node.childStarNode);
+                }
+
+                if (node.children != null) {
+                    // Sort all children nodes based on dimension value
+                    List<InMemoryTreeNode> sortedChildren = new ArrayList<>(node.children.values());
+                    sortedChildren.sort(
+                        Comparator.comparingInt(InMemoryTreeNode::getNodeType).thenComparingLong(InMemoryTreeNode::getDimensionValue)
+                    );
+
+                    totalNumberOfChildren = totalNumberOfChildren + sortedChildren.size();
+                    queue.addAll(sortedChildren);
+                }
+
+                int lastChildId = firstChildId + totalNumberOfChildren - 1;
                 writeStarTreeNode(output, node, firstChildId, lastChildId);
 
-                queue.addAll(sortedChildren);
             }
 
             currentNodeId++;

diff --git a/...a/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNode.java b/...a/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNode.java
@@ -12,8 +12,6 @@
 import java.io.IOException;
 import java.util.Iterator;
 
-import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.ALL;
-
 /**
  * Fixed Length implementation of {@link StarTreeNode}
  *
@@ -129,7 +127,7 @@ public StarTreeNode getChildForDimensionValue(long dimensionValue, boolean isSta
 
     private FixedLengthStarTreeNode handleStarNode() throws IOException {
         FixedLengthStarTreeNode firstNode = new FixedLengthStarTreeNode(in, firstChildId);
-        if (firstNode.getDimensionValue() == ALL) {
+        if (firstNode.getStarTreeNodeType() == StarTreeNodeType.STAR.getValue()) {
             return firstNode;
         } else {
             return null;

diff --git a/...ain/java/org/opensearch/index/compositeindex/datacube/startree/node/InMemoryTreeNode.java b/...ain/java/org/opensearch/index/compositeindex/datacube/startree/node/InMemoryTreeNode.java
@@ -14,7 +14,6 @@
 import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.ALL;
 
 /**
- * /**
  * 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.
  *
@@ -63,6 +62,11 @@ public class InMemoryTreeNode {
      */
     public Map<Long, InMemoryTreeNode> children;
 
+    /**
+     * A map containing the child star node of this star-tree node.
+     */
+    public InMemoryTreeNode childStarNode;
+
     public long getDimensionValue() {
         return dimensionValue;
     }

diff --git a/...rc/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java b/...rc/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java
@@ -19,6 +19,7 @@
 import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTree;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
+import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
 import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator;
 import org.opensearch.index.mapper.CompositeMappedFieldType;
 
@@ -36,6 +37,7 @@
 import static org.opensearch.index.mapper.CompositeMappedFieldType.CompositeFieldType.STAR_TREE;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
@@ -186,19 +188,29 @@ public static void validateFileFormats(
             if (starTreeNode.getChildDimensionId() != -1) {
                 assertFalse(sortedChildren.isEmpty());
                 int childCount = 0;
+                boolean childStarNodeAsserted = false;
                 while (expectedChildrenIterator.hasNext()) {
                     StarTreeNode child = expectedChildrenIterator.next();
-                    InMemoryTreeNode resultChildNode = sortedChildren.get(childCount);
-
-                    assertNotNull(resultChildNode);
-                    assertNotNull(child);
-                    assertStarTreeNode(child, resultChildNode);
+                    InMemoryTreeNode resultChildNode = null;
+                    if (!childStarNodeAsserted && rootNode.childStarNode != null) {
+                        // check if star tree node exists
+                        resultChildNode = rootNode.childStarNode;
+                        assertNotNull(child);
+                        assertStarTreeNode(child, resultChildNode);
+                        childStarNodeAsserted = true;
+                    } else {
+                        resultChildNode = sortedChildren.get(childCount);
+                        assertNotNull(child);
+                        assertNotNull(resultChildNode);
+                        assertStarTreeNode(child, resultChildNode);
+                        assertNotEquals(child.getStarTreeNodeType(), StarTreeNodeType.STAR.getValue());
+                        childCount++;
+                    }
 
                     expectedTreeNodeQueue.add(child);
                     resultTreeNodeQueue.add(resultChildNode);
-
-                    childCount++;
                 }
+
                 assertEquals(childCount, rootNode.children.size());
             } else {
                 assertNull(rootNode.children);
@@ -221,7 +233,7 @@ public static void assertStarTreeNode(StarTreeNode starTreeNode, InMemoryTreeNod
         if (starTreeNode.getChildDimensionId() != -1) {
             assertFalse(starTreeNode.isLeaf());
             if (treeNode.children != null) {
-                assertEquals(starTreeNode.getNumChildren(), treeNode.children.values().size());
+                assertEquals(starTreeNode.getNumChildren(), treeNode.children.values().size() + (treeNode.childStarNode != null ? 1 : 0));
             }
         } else {
             assertTrue(starTreeNode.isLeaf());

diff --git a/...ensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java b/...ensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java
@@ -498,7 +498,6 @@ public void test_sortAndAggregateStarTreeDocuments_nullMetricField() throws IOEx
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/14813")
     public void test_sortAndAggregateStarTreeDocuments_nullAndMinusOneInDimensionField() throws IOException {
         int noOfStarTreeDocuments = 5;
         StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments];