improved binary search to reduce search space in the presence of null…

… or star node

Signed-off-by: Sarthak Aggarwal <[email protected]>

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java

@@ -214,7 +214,20 @@ public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOExce
      */
     private FixedLengthStarTreeNode handleStarNode() throws IOException {
         FixedLengthStarTreeNode firstNode = new FixedLengthStarTreeNode(in, firstChildId);
-        if (firstNode.getStarTreeNodeType() == StarTreeNodeType.STAR.getValue()) {
+        return matchStarTreeNodeTypeOrNull(firstNode, StarTreeNodeType.STAR);
+    }
+
+    /**
+     * Checks if the given node matches the specified StarTreeNodeType.
+     *
+     * @param firstNode The FixedLengthStarTreeNode to check.
+     * @param starTreeNodeType The StarTreeNodeType to match against.
+     * @return The firstNode if its type matches the targetType, null otherwise.
+     * @throws IOException If an I/O error occurs during the operation.
+     */
+    private static FixedLengthStarTreeNode matchStarTreeNodeTypeOrNull(FixedLengthStarTreeNode firstNode, StarTreeNodeType starTreeNodeType)
+        throws IOException {
+        if (firstNode.getStarTreeNodeType() == starTreeNodeType.getValue()) {
             return firstNode;
         } else {
             return null;
@@ -229,7 +242,19 @@ private FixedLengthStarTreeNode handleStarNode() throws IOException {
      * @throws IOException If there's an error reading from the input
      */
     private FixedLengthStarTreeNode binarySearchChild(long dimensionValue) throws IOException {
+
         int low = firstChildId;
+
+        // if the current node is star node, increment the low to reduce the search space
+        if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, firstChildId), StarTreeNodeType.STAR) != null) {
+            low++;
+        }
+
+        // if the current node is null node, increment the low to reduce the search space
+        if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, low), StarTreeNodeType.NULL) != null) {
+            low++;
+        }
+
         int high = getInt(LAST_CHILD_ID_OFFSET);
 
         while (low <= high) {

server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java

@@ -34,6 +34,7 @@ public class FixedLengthStarTreeNodeTests extends OpenSearchTestCase {
     private Directory directory;
     InMemoryTreeNode node;
     InMemoryTreeNode starChild;
+    InMemoryTreeNode nullChild;
     FixedLengthStarTreeNode starTreeNode;
 
     @Before
@@ -63,6 +64,17 @@ public void setup() throws IOException {
         starChild.children = new HashMap<>();
         node.children.put(-1L, starChild);
 
+        nullChild = new InMemoryTreeNode();
+        nullChild.dimensionId = node.dimensionId + 1;
+        nullChild.dimensionValue = -1;
+        nullChild.startDocId = randomInt();
+        nullChild.endDocId = randomInt();
+        nullChild.childDimensionId = -1;
+        nullChild.aggregatedDocId = randomInt();
+        nullChild.nodeType = (byte) -1;
+        nullChild.children = new HashMap<>();
+        node.children.put(null, nullChild);
+
         for (int i = 1; i < randomIntBetween(2, 5); i++) {
             InMemoryTreeNode child = new InMemoryTreeNode();
             child.dimensionId = node.dimensionId + 1;
@@ -140,7 +152,7 @@ public void testGetStarTreeNodeType() throws IOException {
     }
 
     public void testGetChildForDimensionValue() throws IOException {
-        long dimensionValue = randomIntBetween(0, node.children.size() - 2);
+        long dimensionValue = randomIntBetween(0, node.children.size() - 3);
         FixedLengthStarTreeNode childNode = (FixedLengthStarTreeNode) starTreeNode.getChildForDimensionValue(dimensionValue);
         assertNotNull(childNode);
         assertEquals(dimensionValue, childNode.getDimensionValue());