Add query changes to support unsigned-long in star tree (#17275)

Shailesh-Kumar-Singh · web-flow · commit 3e10fe3f4eac · 2025-04-11T11:34:16.000+05:30
Signed-off-by: Shailesh Singh &lt;shaikumm@amazon.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -35,6 +35,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add SearchService and Search GRPC endpoint ([#17830](https://github.com/opensearch-project/OpenSearch/pull/17830))
 - Add update and delete support in pull-based ingestion ([#17822](https://github.com/opensearch-project/OpenSearch/pull/17822))
 - Allow maxPollSize and pollTimeout in IngestionSource to be configurable ([#17863](https://github.com/opensearch-project/OpenSearch/pull/17863))
+- [Star Tree] [Search] Add query changes to support unsigned-long in star tree ([#17275](https://github.com/opensearch-project/OpenSearch/pull/17275))
 
 ### Changed
 - Migrate BC libs to their FIPS counterparts ([#14912](https://github.com/opensearch-project/OpenSearch/pull/14912))
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionDataType.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionDataType.java
@@ -12,14 +12,16 @@
 
 /**
  * Represents the data type of the dimension value.
+ * TODO: This needs to be eventually merged with DimensionFilterMapper and all indexing related code
+ *  which use this should instead use the mapper
  *
  * @opensearch.experimental
  */
 @ExperimentalApi
 public enum DimensionDataType {
     LONG {
         @Override
-        int compare(Long a, Long b) {
+        public int compare(Long a, Long b) {
             if (a == null && b == null) {
                 return 0;
             }
@@ -34,7 +36,7 @@ int compare(Long a, Long b) {
     },
     UNSIGNED_LONG {
         @Override
-        int compare(Long a, Long b) {
+        public int compare(Long a, Long b) {
             if (a == null && b == null) {
                 return 0;
             }
@@ -48,5 +50,5 @@ int compare(Long a, Long b) {
         }
     };
 
-    abstract int compare(Long a, Long b);
+    public abstract int compare(Long a, Long b);
 }
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java
@@ -8,12 +8,15 @@
 package org.opensearch.index.compositeindex.datacube.startree.fileformats.node;
 
 import org.apache.lucene.store.RandomAccessInput;
+import org.opensearch.index.compositeindex.datacube.DimensionDataType;
 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 org.opensearch.search.startree.filter.provider.DimensionFilterMapper;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.Comparator;
 import java.util.Iterator;
 
 /**
@@ -193,15 +196,23 @@ public StarTreeNode getChildStarNode() throws IOException {
     }
 
     @Override
-    public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException {
+    public StarTreeNode getChildForDimensionValue(
+        Long dimensionValue,
+        StarTreeNode lastMatchedChild,
+        DimensionFilterMapper dimensionFilterMapper
+    ) throws IOException {
         // there will be no children for leaf nodes
         if (isLeaf()) {
             return null;
         }
 
         StarTreeNode resultStarTreeNode = null;
         if (null != dimensionValue) {
-            resultStarTreeNode = binarySearchChild(dimensionValue, lastMatchedChild);
+            resultStarTreeNode = binarySearchChild(
+                dimensionValue,
+                lastMatchedChild,
+                dimensionFilterMapper == null ? DimensionDataType.LONG::compare : dimensionFilterMapper.comparator()
+            );
         }
         return resultStarTreeNode;
     }
@@ -238,11 +249,13 @@ private static FixedLengthStarTreeNode matchStarTreeNodeTypeOrNull(FixedLengthSt
      * Performs a binary search to find a child node with the given dimension value.
      *
      * @param dimensionValue The dimension value to search for
+     * @param lastMatchedNode : If not null, we begin the binary search from the node after this.
+     * @param comparator : Comparator (LONG or UNSIGNED_LONG) to compare the dimension values
      * @return The child node if found, null otherwise
      * @throws IOException If there's an error reading from the input
      */
-    private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, StarTreeNode lastMatchedNode) throws IOException {
-
+    private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, StarTreeNode lastMatchedNode, Comparator<Long> comparator)
+        throws IOException {
         int low = firstChildId;
 
         int high = getInt(LAST_CHILD_ID_OFFSET);
@@ -268,10 +281,10 @@ private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, StarTreeN
             int mid = low + (high - low) / 2;
             FixedLengthStarTreeNode midNode = new FixedLengthStarTreeNode(in, mid);
             long midDimensionValue = midNode.getDimensionValue();
-
-            if (midDimensionValue == dimensionValue) {
+            int compare = comparator.compare(midDimensionValue, dimensionValue);
+            if (compare == 0) {
                 return midNode;
-            } else if (midDimensionValue < dimensionValue) {
+            } else if (compare < 0) {
                 low = mid + 1;
             } else {
                 high = mid - 1;
@@ -281,16 +294,19 @@ private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, StarTreeN
     }
 
     @Override
-    public void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector) throws IOException {
-        if (low <= high) {
-            FixedLengthStarTreeNode lowStarTreeNode = binarySearchChild(low, true, null);
+    public void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector, DimensionFilterMapper dimensionFilterMapper)
+        throws IOException {
+        Comparator<Long> comparator = dimensionFilterMapper.comparator();
+        if (comparator.compare(low, high) <= 0) {
+            FixedLengthStarTreeNode lowStarTreeNode = binarySearchChild(low, true, null, comparator);
             if (lowStarTreeNode != null) {
-                FixedLengthStarTreeNode highStarTreeNode = binarySearchChild(high, false, lowStarTreeNode);
+                FixedLengthStarTreeNode highStarTreeNode = binarySearchChild(high, false, lowStarTreeNode, comparator);
                 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.
+                } else if (comparator.compare(lowStarTreeNode.getDimensionValue(), high) <= 0) { // Low StarTreeNode is the last default//
+                                                                                                 // node for that dimension.
                     collector.collectStarTreeNode(lowStarTreeNode);
                 }
             }
@@ -302,11 +318,16 @@ public void collectChildrenInRange(long low, long high, StarTreeNodeCollector co
      * @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.
+     * @param comparator : Comparator (LONG or UNSIGNED_LONG) to compare the dimension values
      * @return : Matched node or null.
      * @throws IOException :
      */
-    private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, boolean matchNextHighest, StarTreeNode lastMatchedNode)
-        throws IOException {
+    private FixedLengthStarTreeNode binarySearchChild(
+        long dimensionValue,
+        boolean matchNextHighest,
+        StarTreeNode lastMatchedNode,
+        Comparator<Long> comparator
+    ) throws IOException {
 
         int low = firstChildId;
         int tempLow = low;
@@ -342,17 +363,18 @@ private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, boolean m
             FixedLengthStarTreeNode midNode = new FixedLengthStarTreeNode(in, mid);
             long midDimensionValue = midNode.getDimensionValue();
 
-            if (midDimensionValue == dimensionValue) {
+            int compare = comparator.compare(midDimensionValue, dimensionValue);
+            if (compare == 0) {
                 return midNode;
             } else {
-                if (midDimensionValue < dimensionValue) { // Going to the right from mid to search next
+                if (compare < 0) { // 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) {
+                        if (comparator.compare(nodeGreaterThanMid.getDimensionValue(), dimensionValue) > 0) {
                             return matchNextHighest ? nodeGreaterThanMid : midNode;
                         }
                     }
@@ -363,7 +385,7 @@ private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, boolean m
                         return matchNextHighest ? midNode : null;
                     } else {
                         FixedLengthStarTreeNode nodeLessThanMid = new FixedLengthStarTreeNode(in, tempHigh);
-                        if (nodeLessThanMid.getDimensionValue() < dimensionValue) {
+                        if (comparator.compare(nodeLessThanMid.getDimensionValue(), dimensionValue) < 0) {
                             return matchNextHighest ? midNode : nodeLessThanMid;
                         }
                     }
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
@@ -10,6 +10,7 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.search.startree.StarTreeNodeCollector;
+import org.opensearch.search.startree.filter.provider.DimensionFilterMapper;
 
 import java.io.IOException;
 import java.util.Iterator;
@@ -107,28 +108,34 @@ public interface StarTreeNode {
      * @param dimensionValue  the dimension value
      * @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
+     *
+     * TODO: Remove this method - Only used in UTs for Star Tree indexing
      */
     default StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
-        return getChildForDimensionValue(dimensionValue, null);
+        return getChildForDimensionValue(dimensionValue, null, 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
+     * @param dimensionFilterMapper : dimensionFilterMapper object
      * @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;
+    StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild, DimensionFilterMapper dimensionFilterMapper)
+        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
+     *  @param dimensionFilterMapper : dimensionFilterMapper object
      * @throws IOException : Any exception in reading the node data from index.
      */
-    void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector) throws IOException;
+    void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector, DimensionFilterMapper dimensionFilterMapper)
+        throws IOException;
 
     /**
      * Returns the child star node for a node in the star-tree.
diff --git a/server/src/main/java/org/opensearch/search/startree/filter/ExactMatchDimFilter.java b/server/src/main/java/org/opensearch/search/startree/filter/ExactMatchDimFilter.java
@@ -35,6 +35,8 @@ public class ExactMatchDimFilter implements DimensionFilter {
     // Order is essential for successive binary search
     private TreeSet<Long> convertedOrdinals;
 
+    private DimensionFilterMapper dimensionFilterMapper;
+
     public ExactMatchDimFilter(String dimensionName, List<Object> valuesToMatch) {
         this.dimensionName = dimensionName;
         this.rawValues = valuesToMatch;
@@ -47,9 +49,10 @@ public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext se
             dimensionName,
             starTreeValues.getStarTreeField().getDimensionsOrder()
         );
-        DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
+        this.dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
             searchContext.mapperService().fieldType(dimensionName)
         );
+
         for (Object rawValue : rawValues) {
             Optional<Long> ordinal = dimensionFilterMapper.getMatchingOrdinal(
                 matchedDim.getField(),
@@ -69,7 +72,7 @@ public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeV
         if (parentNode != null) {
             StarTreeNode lastMatchedNode = null;
             for (long ordinal : convertedOrdinals) {
-                lastMatchedNode = parentNode.getChildForDimensionValue(ordinal, lastMatchedNode);
+                lastMatchedNode = parentNode.getChildForDimensionValue(ordinal, lastMatchedNode, dimensionFilterMapper);
                 if (lastMatchedNode != null) {
                     collector.collectStarTreeNode(lastMatchedNode);
                 }
diff --git a/server/src/main/java/org/opensearch/search/startree/filter/RangeMatchDimFilter.java b/server/src/main/java/org/opensearch/search/startree/filter/RangeMatchDimFilter.java
@@ -37,6 +37,8 @@ public class RangeMatchDimFilter implements DimensionFilter {
 
     private boolean skipRangeCollection = false;
 
+    private DimensionFilterMapper dimensionFilterMapper;
+
     public RangeMatchDimFilter(String dimensionName, Object low, Object high, boolean includeLow, boolean includeHigh) {
         this.dimensionName = dimensionName;
         this.low = low;
@@ -48,9 +50,10 @@ public RangeMatchDimFilter(String dimensionName, Object low, Object high, boolea
     @Override
     public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext searchContext) {
         skipRangeCollection = false;
-        DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
+        this.dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
             searchContext.mapperService().fieldType(dimensionName)
         );
+
         lowOrdinal = 0L;
         if (low != null) {
             MatchType lowMatchType = includeLow ? MatchType.GTE : MatchType.GT;
@@ -77,13 +80,14 @@ public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext se
     public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeValues, StarTreeNodeCollector collector)
         throws IOException {
         if (parentNode != null && !skipRangeCollection) {
-            parentNode.collectChildrenInRange(lowOrdinal, highOrdinal, collector);
+            parentNode.collectChildrenInRange(lowOrdinal, highOrdinal, collector, dimensionFilterMapper);
         }
     }
 
     @Override
     public boolean matchDimValue(long ordinal, StarTreeValues starTreeValues) {
-        return lowOrdinal <= ordinal && ordinal <= highOrdinal;
+        return dimensionFilterMapper.comparator().compare(lowOrdinal, ordinal) <= 0
+            && dimensionFilterMapper.comparator().compare(ordinal, highOrdinal) <= 0;
     }
 
 }
diff --git a/server/src/main/java/org/opensearch/search/startree/filter/provider/DimensionFilterMapper.java b/server/src/main/java/org/opensearch/search/startree/filter/provider/DimensionFilterMapper.java
@@ -14,9 +14,11 @@
 import org.apache.lucene.sandbox.document.HalfFloatPoint;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.NumericUtils;
+import org.opensearch.common.Numbers;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.lucene.BytesRefs;
 import org.opensearch.common.lucene.Lucene;
+import org.opensearch.index.compositeindex.datacube.DimensionDataType;
 import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
 import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedSetStarTreeValuesIterator;
 import org.opensearch.index.mapper.KeywordFieldMapper.KeywordFieldType;
@@ -29,6 +31,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -40,6 +43,7 @@
 import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.INTEGER;
 import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.LONG;
 import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.SHORT;
+import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.UNSIGNED_LONG;
 import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.hasDecimalPart;
 import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.signum;
 
@@ -88,6 +92,10 @@ Optional<Long> getMatchingOrdinal(
         DimensionFilter.MatchType matchType
     );
 
+    default Comparator<Long> comparator() {
+        return DimensionDataType.LONG::compare;
+    }
+
     /**
      * Singleton Factory for @{@link DimensionFilterMapper}
      */
@@ -109,7 +117,9 @@ class Factory {
             DOUBLE.typeName(),
             new DoubleFieldMapperNumeric(),
             org.opensearch.index.mapper.KeywordFieldMapper.CONTENT_TYPE,
-            new KeywordFieldMapper()
+            new KeywordFieldMapper(),
+            UNSIGNED_LONG.typeName(),
+            new UnsignedLongFieldMapperNumeric()
         );
 
         public static DimensionFilterMapper fromMappedFieldType(MappedFieldType mappedFieldType) {
@@ -208,6 +218,25 @@ Long defaultMaximum() {
     }
 }
 
+class UnsignedLongFieldMapperNumeric extends NumericNonDecimalMapper {
+
+    @Override
+    Long defaultMinimum() {
+        return Numbers.MIN_UNSIGNED_LONG_VALUE_AS_LONG;
+    }
+
+    @Override
+    Long defaultMaximum() {
+        return Numbers.MAX_UNSIGNED_LONG_VALUE_AS_LONG;
+    }
+
+    @Override
+    public Comparator<Long> comparator() {
+        return DimensionDataType.UNSIGNED_LONG::compare;
+    }
+
+}
+
 abstract class NumericDecimalFieldMapper extends NumericMapper {
 
     @Override
diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeSearchTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeSearchTests.java
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/BigIntegerField.java b/server/src/test/java/org/opensearch/search/aggregations/startree/BigIntegerField.java
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
