add query changes to support unsigned-long in star tree

Signed-off-by: Shailesh Singh <[email protected]>

server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionDataType.java

@@ -19,7 +19,7 @@
 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 +34,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 +48,5 @@ int compare(Long a, Long b) {
         }
     };
 
-    abstract int compare(Long a, Long b);
+    public abstract int compare(Long a, Long b);
 }

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

@@ -14,6 +14,7 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.Comparator;
 import java.util.Iterator;
 
 /**
@@ -193,15 +194,16 @@ public StarTreeNode getChildStarNode() throws IOException {
     }
 
     @Override
-    public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException {
+    public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild, Comparator<Long> comparator)
+        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, comparator);
         }
         return resultStarTreeNode;
     }
@@ -238,11 +240,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 +272,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 +285,18 @@ 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, Comparator<Long> comparator)
+        throws IOException {
+        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 +308,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 +353,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 +375,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;
                         }
                     }

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

@@ -9,9 +9,11 @@
 package org.opensearch.index.compositeindex.datacube.startree.node;
 
 import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.compositeindex.datacube.DimensionDataType;
 import org.opensearch.search.startree.StarTreeNodeCollector;
 
 import java.io.IOException;
+import java.util.Comparator;
 import java.util.Iterator;
 
 /**
@@ -109,26 +111,29 @@ public interface StarTreeNode {
      * @throws IOException if an I/O error occurs while retrieving the child node
      */
     default StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
-        return getChildForDimensionValue(dimensionValue, null);
+        return getChildForDimensionValue(dimensionValue, null, DimensionDataType.LONG::compare);
     }
 
     /**
      * 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 comparator : Comparator (LONG or UNSIGNED_LONG) to compare the dimension values
      * @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, Comparator<Long> comparator)
+        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 comparator : Comparator (LONG or UNSIGNED_LONG) to compare the dimension values
      * @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, Comparator<Long> comparator) throws IOException;
 
     /**
      * Returns the child star node for a node in the star-tree.

server/src/main/java/org/opensearch/search/startree/filter/ExactMatchDimFilter.java

@@ -18,6 +18,7 @@
 import org.opensearch.search.startree.filter.provider.DimensionFilterMapper;
 
 import java.io.IOException;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
 import java.util.TreeSet;
@@ -35,6 +36,8 @@ public class ExactMatchDimFilter implements DimensionFilter {
     // Order is essential for successive binary search
     private TreeSet<Long> convertedOrdinals;
 
+    Comparator<Long> comparator;
+
     public ExactMatchDimFilter(String dimensionName, List<Object> valuesToMatch) {
         this.dimensionName = dimensionName;
         this.rawValues = valuesToMatch;
@@ -50,6 +53,7 @@ public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext se
         DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
             searchContext.mapperService().fieldType(dimensionName)
         );
+        this.comparator = dimensionFilterMapper.comparator();
         for (Object rawValue : rawValues) {
             Optional<Long> ordinal = dimensionFilterMapper.getMatchingOrdinal(
                 matchedDim.getField(),
@@ -69,7 +73,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, comparator);
                 if (lastMatchedNode != null) {
                     collector.collectStarTreeNode(lastMatchedNode);
                 }

server/src/main/java/org/opensearch/search/startree/filter/RangeMatchDimFilter.java

@@ -16,6 +16,7 @@
 import org.opensearch.search.startree.filter.provider.DimensionFilterMapper;
 
 import java.io.IOException;
+import java.util.Comparator;
 import java.util.Optional;
 
 /**
@@ -37,6 +38,8 @@ public class RangeMatchDimFilter implements DimensionFilter {
 
     private boolean skipRangeCollection = false;
 
+    private Comparator<Long> comparator;
+
     public RangeMatchDimFilter(String dimensionName, Object low, Object high, boolean includeLow, boolean includeHigh) {
         this.dimensionName = dimensionName;
         this.low = low;
@@ -51,6 +54,8 @@ public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext se
         DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
             searchContext.mapperService().fieldType(dimensionName)
         );
+        this.comparator = dimensionFilterMapper.comparator();
+
         lowOrdinal = 0L;
         if (low != null) {
             MatchType lowMatchType = includeLow ? MatchType.GTE : MatchType.GT;
@@ -77,13 +82,13 @@ 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, comparator);
         }
     }
 
     @Override
     public boolean matchDimValue(long ordinal, StarTreeValues starTreeValues) {
-        return lowOrdinal <= ordinal && ordinal <= highOrdinal;
+        return comparator.compare(lowOrdinal, ordinal) <= 0 && comparator.compare(ordinal, highOrdinal) <= 0;
     }
 
 }

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,20 @@ Optional<Long> getMatchingOrdinal(
         DimensionFilter.MatchType matchType
     );
 
+    /**
+     * Returns the dimensionDataType used for comparing dimension values. <br>
+     * This determines how numeric values are compared: <br>
+     *   - DimensionDataType.UNSIGNED_LONG for unsigned long values <br>
+     *   - DimensionDataType.LONG for all other numeric types (DEFAULT)
+     */
+    default DimensionDataType getDimensionDataType() {
+        return DimensionDataType.LONG;
+    }
+
+    default Comparator<Long> comparator() {
+        return (a, b) -> getDimensionDataType().compare(a, b);
+    }
+
     /**
      * Singleton Factory for @{@link DimensionFilterMapper}
      */
@@ -109,7 +127,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) {
@@ -161,14 +181,14 @@ public DimensionFilter getRangeMatchFilter(
         Long parsedLow = rawLow == null ? defaultMinimum() : numberFieldType.numberType().parse(rawLow, true).longValue();
         Long parsedHigh = rawHigh == null ? defaultMaximum() : numberFieldType.numberType().parse(rawHigh, true).longValue();
 
-        boolean lowerTermHasDecimalPart = hasDecimalPart(parsedLow);
+        boolean lowerTermHasDecimalPart = hasDecimalPart(rawLow);
         if ((lowerTermHasDecimalPart == false && includeLow == false) || (lowerTermHasDecimalPart && signum(parsedLow) > 0)) {
             if (parsedLow.equals(defaultMaximum())) {
                 return new MatchNoneFilter();
             }
             ++parsedLow;
         }
-        boolean upperTermHasDecimalPart = hasDecimalPart(parsedHigh);
+        boolean upperTermHasDecimalPart = hasDecimalPart(rawHigh);
         if ((upperTermHasDecimalPart == false && includeHigh == false) || (upperTermHasDecimalPart && signum(parsedHigh) < 0)) {
             if (parsedHigh.equals(defaultMinimum())) {
                 return new MatchNoneFilter();
@@ -208,6 +228,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 DimensionDataType getDimensionDataType() {
+        return DimensionDataType.UNSIGNED_LONG;
+    }
+
+}
+
 abstract class NumericDecimalFieldMapper extends NumericMapper {
 
     @Override