From 1545244b1313e841d0aa6cc4fde0a7917935f621 Mon Sep 17 00:00:00 2001 From: Shailesh Singh Date: Thu, 6 Feb 2025 17:41:48 +0530 Subject: [PATCH] add query changes to support unsigned-long in star tree Signed-off-by: Shailesh Singh --- .../datacube/DimensionDataType.java | 6 +- .../node/FixedLengthStarTreeNode.java | 48 +++++---- .../datacube/startree/node/StarTreeNode.java | 11 +- .../startree/filter/ExactMatchDimFilter.java | 6 +- .../startree/filter/RangeMatchDimFilter.java | 9 +- .../provider/DimensionFilterMapper.java | 45 +++++++- .../FixedLengthStarTreeNodeSearchTests.java | 45 +++++--- .../startree/BigIntegerField.java | 100 ++++++++++++++++++ .../startree/MetricAggregatorTests.java | 30 +++++- 9 files changed, 253 insertions(+), 47 deletions(-) create mode 100644 server/src/test/java/org/opensearch/search/aggregations/startree/BigIntegerField.java 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 index 67138b69c69fa..93fe514c18a3f 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionDataType.java +++ b/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); } 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 index c6c4993290c16..b0f1cf874d43f 100644 --- 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 @@ -14,6 +14,7 @@ import java.io.IOException; import java.io.UncheckedIOException; +import java.util.Comparator; import java.util.Iterator; /** @@ -193,7 +194,8 @@ public StarTreeNode getChildStarNode() throws IOException { } @Override - public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException { + public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild, Comparator comparator) + throws IOException { // there will be no children for leaf nodes if (isLeaf()) { return null; @@ -201,7 +203,7 @@ public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode 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 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 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 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; } } 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 index 40161a942ae4b..cdeec6ab6817b 100644 --- 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 @@ -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 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 comparator) 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 index 28ea261ca1e56..b28a2a81dc550 100644 --- a/server/src/main/java/org/opensearch/search/startree/filter/ExactMatchDimFilter.java +++ b/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 convertedOrdinals; + Comparator comparator; + public ExactMatchDimFilter(String dimensionName, List 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 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); } 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 index fecf1a9ebf76b..45aa1df1b36d5 100644 --- a/server/src/main/java/org/opensearch/search/startree/filter/RangeMatchDimFilter.java +++ b/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 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; } } 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 index 8afdb00864b22..1ba3e9d00b2f2 100644 --- 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,20 @@ Optional getMatchingOrdinal( DimensionFilter.MatchType matchType ); + /** + * Returns the dimensionDataType used for comparing dimension values.
+ * This determines how numeric values are compared:
+ * - DimensionDataType.UNSIGNED_LONG for unsigned long values
+ * - DimensionDataType.LONG for all other numeric types (DEFAULT) + */ + default DimensionDataType getDimensionDataType() { + return DimensionDataType.LONG; + } + + default Comparator 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 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 index 4d95034d80bb7..dadad15ce8753 100644 --- 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 @@ -20,6 +20,7 @@ import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.IndexOutput; +import org.opensearch.index.compositeindex.datacube.DimensionDataType; import org.opensearch.index.compositeindex.datacube.startree.fileformats.StarTreeWriter; import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; @@ -53,18 +54,30 @@ public void testExactMatch() { result &= -1 == lastMatchedNode.getDimensionValue(); // Leaf Node should return null result &= null == lastMatchedNode.getChildForDimensionValue(5L); - result &= null == lastMatchedNode.getChildForDimensionValue(5L, lastMatchedNode); + result &= null == lastMatchedNode.getChildForDimensionValue(5L, lastMatchedNode, DimensionDataType.LONG::compare); // Asserting Last Matched Node works as expected - lastMatchedNode = (FixedLengthStarTreeNode) fixedLengthStarTreeNode.getChildForDimensionValue(1L, lastMatchedNode); + lastMatchedNode = (FixedLengthStarTreeNode) fixedLengthStarTreeNode.getChildForDimensionValue( + 1L, + lastMatchedNode, + DimensionDataType.LONG::compare + ); result &= 1 == lastMatchedNode.getDimensionValue(); - lastMatchedNode = (FixedLengthStarTreeNode) fixedLengthStarTreeNode.getChildForDimensionValue(5L, lastMatchedNode); + lastMatchedNode = (FixedLengthStarTreeNode) fixedLengthStarTreeNode.getChildForDimensionValue( + 5L, + lastMatchedNode, + DimensionDataType.LONG::compare + ); result &= 5 == lastMatchedNode.getDimensionValue(); // Asserting null is returned when last matched node is after the value to search. - lastMatchedNode = (FixedLengthStarTreeNode) fixedLengthStarTreeNode.getChildForDimensionValue(2L, lastMatchedNode); + lastMatchedNode = (FixedLengthStarTreeNode) fixedLengthStarTreeNode.getChildForDimensionValue( + 2L, + lastMatchedNode, + DimensionDataType.LONG::compare + ); result &= null == lastMatchedNode; // When dimension value is null result &= null == fixedLengthStarTreeNode.getChildForDimensionValue(null); - result &= null == fixedLengthStarTreeNode.getChildForDimensionValue(null, null); + result &= null == fixedLengthStarTreeNode.getChildForDimensionValue(null, null, DimensionDataType.LONG::compare); // non-existing dimensionValue result &= null == fixedLengthStarTreeNode.getChildForDimensionValue(4L); result &= null == fixedLengthStarTreeNode.getChildForDimensionValue(randomLongBetween(6, Long.MAX_VALUE)); @@ -136,43 +149,43 @@ public void testRangeMatch() { ArrayBasedCollector collector; // Whole range collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(-20, 26, collector); + fixedLengthStarTreeNode.collectChildrenInRange(-20, 26, collector, DimensionDataType.LONG::compare); result &= collector.matchAllCollectedValues(new long[] { -10, -1, 1, 2, 5, 9, 25 }); // Subset matched from left collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(-2, 1, collector); + fixedLengthStarTreeNode.collectChildrenInRange(-2, 1, collector, DimensionDataType.LONG::compare); result &= collector.matchAllCollectedValues(new long[] { -1, 1 }); // Subset matched from right collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(6, 100, collector); + fixedLengthStarTreeNode.collectChildrenInRange(6, 100, collector, DimensionDataType.LONG::compare); result &= collector.matchAllCollectedValues(new long[] { 9, 25 }); // No match on left collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(-30, -20, collector); + fixedLengthStarTreeNode.collectChildrenInRange(-30, -20, collector, DimensionDataType.LONG::compare); result &= collector.collectedNodeCount() == 0; // No match on right collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(30, 50, collector); + fixedLengthStarTreeNode.collectChildrenInRange(30, 50, collector, DimensionDataType.LONG::compare); result &= collector.collectedNodeCount() == 0; // Low > High collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(50, 10, collector); + fixedLengthStarTreeNode.collectChildrenInRange(50, 10, collector, DimensionDataType.LONG::compare); result &= collector.collectedNodeCount() == 0; // Match leftmost collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(-30, -10, collector); + fixedLengthStarTreeNode.collectChildrenInRange(-30, -10, collector, DimensionDataType.LONG::compare); result &= collector.matchAllCollectedValues(new long[] { -10 }); // Match rightmost collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(10, 25, collector); + fixedLengthStarTreeNode.collectChildrenInRange(10, 25, collector, DimensionDataType.LONG::compare); result &= collector.matchAllCollectedValues(new long[] { 25 }); // Match contains interval which has nothing collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(10, 24, collector); + fixedLengthStarTreeNode.collectChildrenInRange(10, 24, collector, DimensionDataType.LONG::compare); result &= collector.collectedNodeCount() == 0; // Match contains interval which has nothing collector = new ArrayBasedCollector(); - fixedLengthStarTreeNode.collectChildrenInRange(6, 24, collector); + fixedLengthStarTreeNode.collectChildrenInRange(6, 24, collector, DimensionDataType.LONG::compare); result &= collector.matchAllCollectedValues(new long[] { 9 }); return result; } catch (IOException e) { @@ -187,7 +200,7 @@ public void testRangeMatch() { try { ArrayBasedCollector collector = new ArrayBasedCollector(); long low = randomLong(), high = randomLong(); - fixedLengthStarTreeNode.collectChildrenInRange(low, high, collector); + fixedLengthStarTreeNode.collectChildrenInRange(low, high, collector, DimensionDataType.LONG::compare); if (low < high) { Long lowValue = treeSet.ceiling(low); if (lowValue != null) { 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 new file mode 100644 index 0000000000000..d4733379579d6 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/aggregations/startree/BigIntegerField.java @@ -0,0 +1,100 @@ +/* +* SPDX-License-Identifier: Apache-2.0 +* +* The OpenSearch Contributors require contributions made to +* this file be licensed under the Apache-2.0 license or a +* compatible open source license. +*/ +package org.opensearch.search.aggregations.startree; + +import org.apache.lucene.document.Field; +import org.apache.lucene.document.FieldType; +import org.apache.lucene.document.StoredValue; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.NumericUtils; + +import java.math.BigInteger; + +public final class BigIntegerField extends Field { + + private static final FieldType FIELD_TYPE = new FieldType(); + private static final FieldType FIELD_TYPE_STORED; + public static final int BYTES = 16; + + static { + FIELD_TYPE.setDimensions(1, 16); + FIELD_TYPE.setDocValuesType(DocValuesType.SORTED_NUMERIC); + FIELD_TYPE.freeze(); + + FIELD_TYPE_STORED = new FieldType(FIELD_TYPE); + FIELD_TYPE_STORED.setStored(true); + FIELD_TYPE_STORED.freeze(); + } + + private final StoredValue storedValue; + + /** + * Creates a new BigIntegerField, indexing the provided point, storing it as a DocValue, and optionally + * storing it as a stored field. + * + * @param name field name + * @param value the BigInteger value + * @param stored whether to store the field + * @throws IllegalArgumentException if the field name or value is null. + */ + public BigIntegerField(String name, BigInteger value, Field.Store stored) { + super(name, stored == Field.Store.YES ? FIELD_TYPE_STORED : FIELD_TYPE); + fieldsData = value; + if (stored == Field.Store.YES) { + storedValue = new StoredValue(value.longValue()); + } else { + storedValue = null; + } + } + + @Override + public BytesRef binaryValue() { + return pack((BigInteger) fieldsData); + } + + @Override + public StoredValue storedValue() { + return storedValue; + } + + @Override + public void setLongValue(long value) { + super.setLongValue(value); + if (storedValue != null) { + storedValue.setLongValue(value); + } + } + + @Override + public String toString() { + return getClass().getSimpleName() + " <" + name + ':' + fieldsData + '>'; + } + + private static BytesRef pack(BigInteger... point) { + if (point == null) { + throw new IllegalArgumentException("point must not be null"); + } + if (point.length == 0) { + throw new IllegalArgumentException("point must not be 0 dimensions"); + } + byte[] packed = new byte[point.length * BYTES]; + + for (int dim = 0; dim < point.length; dim++) { + encodeDimension(point[dim], packed, dim * BYTES); + } + + return new BytesRef(packed); + } + + /** Encode single BigInteger dimension */ + public static void encodeDimension(BigInteger value, byte[] dest, int offset) { + NumericUtils.bigIntToSortableBytes(value, BYTES, dest, offset); + } + +} 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 index 6e10562c3a846..6f4da8e7ba8cc 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java @@ -78,6 +78,7 @@ import org.junit.Before; import java.io.IOException; +import java.math.BigInteger; import java.util.ArrayList; import java.util.Collections; import java.util.LinkedHashMap; @@ -91,6 +92,7 @@ import java.util.stream.Collectors; import java.util.stream.Stream; +import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.objectToUnsignedLong; import static org.opensearch.search.aggregations.AggregationBuilders.avg; import static org.opensearch.search.aggregations.AggregationBuilders.count; import static org.opensearch.search.aggregations.AggregationBuilders.max; @@ -146,7 +148,14 @@ public void testStarTreeDocValues() throws IOException { new DimensionFieldData("long_field", () -> random().nextInt(50), DimensionTypes.LONG), new DimensionFieldData("half_float_field", () -> random().nextFloat(50), DimensionTypes.HALF_FLOAT), new DimensionFieldData("float_field", () -> random().nextFloat(50), DimensionTypes.FLOAT), - new DimensionFieldData("double_field", () -> random().nextDouble(50), DimensionTypes.DOUBLE) + new DimensionFieldData("double_field", () -> random().nextDouble(50), DimensionTypes.DOUBLE), + new DimensionFieldData("unsigned_long_field", () -> { + long queryValue = randomBoolean() + ? 9223372036854775807L - random().nextInt(100000) + : -9223372036854775808L + random().nextInt(100000); + + return objectToUnsignedLong(asUnsignedDecimalString(queryValue), false); + }, DimensionTypes.UNSIGNED_LONG) ); for (Supplier maxLeafDocsSupplier : MAX_LEAF_DOC_VARIATIONS) { testStarTreeDocValuesInternal( @@ -624,6 +633,17 @@ public MappedFieldType getMappedField(String fieldName) { public Dimension getDimension(String fieldName) { return new OrdinalDimension(fieldName); } + }), + UNSIGNED_LONG(new NumericDimensionFieldDataSupplier() { + @Override + NumberFieldMapper.NumberType numberType() { + return NumberFieldMapper.NumberType.UNSIGNED_LONG; + } + + @Override + public IndexableField getField(String fieldName, Supplier valueSupplier) { + return new BigIntegerField(fieldName, (BigInteger) valueSupplier.get(), Field.Store.YES); + } }); private final DimensionFieldDataSupplier dimensionFieldDataSupplier; @@ -638,4 +658,12 @@ public DimensionFieldDataSupplier getFieldDataSupplier() { } + private String asUnsignedDecimalString(long l) { + BigInteger b = BigInteger.valueOf(l); + if (b.signum() < 0) { + b = b.add(BigInteger.ONE.shiftLeft(64)); + } + return b.toString(); + } + }