Extensible design to add new query and field type support for Star Tree

CHANGELOG.md

@@ -40,6 +40,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))
 - Propagate the sourceIncludes and excludes fields from fetchSourceContext to FieldsVisitor. ([#17080](https://github.com/opensearch-project/OpenSearch/pull/17080))
 - [Star Tree] [Search] Resolving Date histogram with metric aggregation using star-tree ([#16674](https://github.com/opensearch-project/OpenSearch/pull/16674))
+- [Star Tree] [Search] Extensible design to support different query and field types ([#17137](https://github.com/opensearch-project/OpenSearch/pull/17137))
 
 ### Dependencies
 - Bump `com.google.cloud:google-cloud-core-http` from 2.23.0 to 2.47.0 ([#16504](https://github.com/opensearch-project/OpenSearch/pull/16504))

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

@@ -10,6 +10,7 @@
 import org.apache.lucene.store.RandomAccessInput;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
+import org.opensearch.search.startree.StarTreeNodeCollector;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -192,15 +193,15 @@ public StarTreeNode getChildStarNode() throws IOException {
     }
 
     @Override
-    public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
+    public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException {
         // there will be no children for leaf nodes
         if (isLeaf()) {
             return null;
         }
 
         StarTreeNode resultStarTreeNode = null;
         if (null != dimensionValue) {
-            resultStarTreeNode = binarySearchChild(dimensionValue);
+            resultStarTreeNode = binarySearchChild(dimensionValue, lastMatchedChild);
         }
         return resultStarTreeNode;
     }
@@ -240,21 +241,29 @@ private static FixedLengthStarTreeNode matchStarTreeNodeTypeOrNull(FixedLengthSt
      * @return The child node if found, null otherwise
      * @throws IOException If there's an error reading from the input
      */
-    private FixedLengthStarTreeNode binarySearchChild(long dimensionValue) throws IOException {
+    private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, StarTreeNode lastMatchedNode) 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++;
-        }
-
         int high = getInt(LAST_CHILD_ID_OFFSET);
         // if the current node is null node, decrement the high to reduce the search space
         if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, high), StarTreeNodeType.NULL) != null) {
             high--;
         }
 
+        if (lastMatchedNode instanceof FixedLengthStarTreeNode) {
+            int lastMatchedNodeId = ((FixedLengthStarTreeNode) lastMatchedNode).nodeId();
+            // Start the binary search from node after the last matched as low.
+            if ((lastMatchedNodeId + 1) <= high) {
+                low = lastMatchedNodeId + 1;
+            } else {
+                return null;
+            }
+        } else if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, low), StarTreeNodeType.STAR) != null) {
+            // if the current node is star node, increment the low to reduce the search space
+            low++;
+        }
+
         while (low <= high) {
             int mid = low + (high - low) / 2;
             FixedLengthStarTreeNode midNode = new FixedLengthStarTreeNode(in, mid);
@@ -271,6 +280,100 @@ private FixedLengthStarTreeNode binarySearchChild(long dimensionValue) throws IO
         return null;
     }
 
+    @Override
+    public void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector) throws IOException {
+        if (low <= high) {
+            FixedLengthStarTreeNode lowStarTreeNode = binarySearchChild(low, true, null);
+            if (lowStarTreeNode != null) {
+                FixedLengthStarTreeNode highStarTreeNode = binarySearchChild(high, false, lowStarTreeNode);
+                if (highStarTreeNode != null) {
+                    for (int lowNodeId = lowStarTreeNode.nodeId(); lowNodeId <= highStarTreeNode.nodeId(); ++lowNodeId) {
+                        collector.collectStarTreeNode(new FixedLengthStarTreeNode(in, lowNodeId));
+                    }
+                } else if (lowStarTreeNode.getDimensionValue() <= high) { // Low StarTreeNode is the last default node for that dimension.
+                    collector.collectStarTreeNode(lowStarTreeNode);
+                }
+            }
+        }
+    }
+
+    /**
+     *
+     * @param dimensionValue : The dimension to match.
+     * @param matchNextHighest : If true then we try to return @dimensionValue or the next Highest. Else, we return @dimensionValue or the next Lowest.
+     * @param lastMatchedNode : If not null, we begin the binary search from the node after this.
+     * @return : Matched node or null.
+     * @throws IOException :
+     */
+    private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, boolean matchNextHighest, StarTreeNode lastMatchedNode)
+        throws IOException {
+
+        int low = firstChildId;
+        int tempLow = low;
+        int starNodeId, nullNodeId;
+        starNodeId = nullNodeId = Integer.MIN_VALUE;
+
+        // if the current node is star node, increment the tempLow to reduce the search space
+        if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempLow), StarTreeNodeType.STAR) != null) {
+            starNodeId = tempLow;
+            tempLow++;
+        }
+
+        int high = getInt(LAST_CHILD_ID_OFFSET);
+        int tempHigh = high;
+        // if the current node is null node, decrement the tempHigh to reduce the search space
+        if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempHigh), StarTreeNodeType.NULL) != null) {
+            nullNodeId = tempHigh;
+            tempHigh--;
+        }
+
+        if (lastMatchedNode instanceof FixedLengthStarTreeNode) {
+            int lastMatchedNodeId = ((FixedLengthStarTreeNode) lastMatchedNode).nodeId();
+            // Start the binary search from node after the last matched as low.
+            if ((lastMatchedNodeId + 1) <= tempHigh) {
+                tempLow = lastMatchedNodeId + 1;
+            } else {
+                return null;
+            }
+        }
+
+        while (tempLow <= tempHigh) {
+            int mid = tempLow + (tempHigh - tempLow) / 2;
+            FixedLengthStarTreeNode midNode = new FixedLengthStarTreeNode(in, mid);
+            long midDimensionValue = midNode.getDimensionValue();
+
+            if (midDimensionValue == dimensionValue) {
+                return midNode;
+            } else {
+                if (midDimensionValue < dimensionValue) { // Going to the right from mid to search next
+                    tempLow = mid + 1;
+                    // We are going out of bounds for this dimension on the right side.
+                    if (tempLow > high || tempLow == nullNodeId) {
+                        return matchNextHighest ? null : midNode;
+                    } else {
+                        FixedLengthStarTreeNode nodeGreaterThanMid = new FixedLengthStarTreeNode(in, tempLow);
+                        if (nodeGreaterThanMid.getDimensionValue() > dimensionValue) {
+                            return matchNextHighest ? nodeGreaterThanMid : midNode;
+                        }
+                    }
+                } else {  // Going to the left from mid to search next
+                    tempHigh = mid - 1;
+                    // We are going out of bounds for this dimension on the left side.
+                    if (tempHigh < low || tempHigh == starNodeId) {
+                        return matchNextHighest ? midNode : null;
+                    } else {
+                        FixedLengthStarTreeNode nodeLessThanMid = new FixedLengthStarTreeNode(in, tempHigh);
+                        if (nodeLessThanMid.getDimensionValue() < dimensionValue) {
+                            return matchNextHighest ? midNode : nodeLessThanMid;
+                        }
+                    }
+                }
+            }
+        }
+        return null;
+
+    }
+
     @Override
     public Iterator<FixedLengthStarTreeNode> getChildrenIterator() throws IOException {
         return new Iterator<>() {
@@ -297,4 +400,8 @@ public void remove() {
             }
         };
     }
+
+    public int nodeId() {
+        return nodeId;
+    }
 }

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

@@ -9,6 +9,7 @@
 package org.opensearch.index.compositeindex.datacube.startree.node;
 
 import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.search.startree.StarTreeNodeCollector;
 
 import java.io.IOException;
 import java.util.Iterator;
@@ -107,7 +108,27 @@ public interface StarTreeNode {
      * @return the child node for the given dimension value or null if child is not present
      * @throws IOException if an I/O error occurs while retrieving the child node
      */
-    StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException;
+    default StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
+        return getChildForDimensionValue(dimensionValue, null);
+    }
+
+    /**
+     * Matches the given @dimensionValue amongst the child default nodes for this node.
+     * @param dimensionValue : Value to match
+     * @param lastMatchedChild : If not null, binary search will use this as the start/low
+     * @return : Matched StarTreeNode or null if not found
+     * @throws IOException : Any exception in reading the node data from index.
+     */
+    StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException;
+
+    /**
+     * Collects all matching child nodes whose dimension values lie within the range of low and high, both inclusive.
+     * @param low : Starting of the range ( inclusive )
+     * @param high : End of the range ( inclusive )
+     * @param collector : Collector to collect the matched child StarTreeNode's
+     * @throws IOException : Any exception in reading the node data from index.
+     */
+    void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector) throws IOException;
 
     /**
      * Returns the child star node for a node in the star-tree.