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 df2ce9096bfc1..d5b599883ebdc 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 @@ -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; @@ -200,7 +201,21 @@ public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOExce StarTreeNode resultStarTreeNode = null; if (null != dimensionValue) { - resultStarTreeNode = binarySearchChild(dimensionValue); + resultStarTreeNode = binarySearchChild(dimensionValue, null); + } + return resultStarTreeNode; + } + + @Override + 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, lastMatchedChild); } return resultStarTreeNode; } @@ -240,15 +255,19 @@ 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) { + if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, low), StarTreeNodeType.STAR) != null) { low++; } + if (lastMatchedNode instanceof FixedLengthStarTreeNode) { + low = ((FixedLengthStarTreeNode) lastMatchedNode).nodeId(); + } + 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) { @@ -271,6 +290,82 @@ private FixedLengthStarTreeNode binarySearchChild(long dimensionValue) throws IO return null; } + @Override + public void collectChildrenInRange(Long low, Long high, StarTreeNodeCollector collector) throws IOException { + FixedLengthStarTreeNode lowStarTreeNode = binarySearchChild(low, true, null); + FixedLengthStarTreeNode highStarTreeNode = binarySearchChild(high, false, lowStarTreeNode); + if (lowStarTreeNode != null && highStarTreeNode != null) { + for (int lowNodeId = lowStarTreeNode.nodeId(); lowNodeId <= highStarTreeNode.nodeId(); ++lowNodeId) { + collector.collectStarNode(new FixedLengthStarTreeNode(in, lowNodeId)); + } + } + } + + private FixedLengthStarTreeNode binarySearchChild(long dimensionValue, boolean matchNextHighest, StarTreeNode lastMatchedNode) + throws IOException { + + int low = firstChildId; + int tempLow = low; + + // if the current node is star node, increment the tempLow to reduce the search space + if (matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempLow), StarTreeNodeType.STAR) != null) { + tempLow++; + } + + if (lastMatchedNode instanceof FixedLengthStarTreeNode) { + tempLow = ((FixedLengthStarTreeNode) lastMatchedNode).nodeId(); + } + + 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) { + tempHigh--; + } + + 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) + || matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempLow), StarTreeNodeType.NULL) != null) { + return matchNextHighest ? null : midNode; + } else { + FixedLengthStarTreeNode nodeGreaterThanMid = new FixedLengthStarTreeNode(in, tempLow); + if (matchNextHighest && nodeGreaterThanMid.getDimensionValue() > dimensionValue) { + return nodeGreaterThanMid; + } else if (!matchNextHighest && nodeGreaterThanMid.getDimensionValue() > dimensionValue) { + return 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) + || matchStarTreeNodeTypeOrNull(new FixedLengthStarTreeNode(in, tempHigh), StarTreeNodeType.STAR) != null) { + return matchNextHighest ? midNode : null; + } else { + FixedLengthStarTreeNode nodeLessThanMid = new FixedLengthStarTreeNode(in, tempHigh); + if (matchNextHighest && (nodeLessThanMid.getDimensionValue() < dimensionValue)) { + return midNode; + } else if (!matchNextHighest && (nodeLessThanMid.getDimensionValue() > dimensionValue)) { + return nodeLessThanMid; + } + } + } + } + } + return null; + + } + @Override public Iterator getChildrenIterator() throws IOException { return new Iterator<>() { @@ -297,4 +392,8 @@ public void remove() { } }; } + + public int nodeId() { + return nodeId; + } } 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 fce3e30e9ebf6..58d5854b767e0 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,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; @@ -109,6 +110,10 @@ public interface StarTreeNode { */ StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException; + StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException; + + void collectChildrenInRange(Long low, Long high, StarTreeNodeCollector collector) 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/DimensionFilter.java b/server/src/main/java/org/opensearch/search/startree/DimensionFilter.java index 481f3c3b88716..1de8d2dee1125 100644 --- a/server/src/main/java/org/opensearch/search/startree/DimensionFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/DimensionFilter.java @@ -8,7 +8,6 @@ package org.opensearch.search.startree; -import org.apache.lucene.index.LeafReaderContext; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode; @@ -18,7 +17,7 @@ @ExperimentalApi public interface DimensionFilter { - public void initialiseForSegment(LeafReaderContext leafReaderContext, StarTreeValues starTreeValues) throws IOException; + public void initialiseForSegment(StarTreeValues starTreeValues) throws IOException; public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeValues, StarTreeNodeCollector collector) throws IOException; diff --git a/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java b/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java index 1889ccae7c64d..a3230de8bc8f1 100644 --- a/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java @@ -8,20 +8,17 @@ package org.opensearch.search.startree; -import org.apache.lucene.index.LeafReaderContext; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode; import java.io.IOException; -import java.util.ArrayList; import java.util.List; -import java.util.stream.Collectors; +import java.util.TreeSet; import static org.opensearch.search.startree.FieldToDimensionOrdinalMapper.SingletonFactory.getFieldToDimensionOrdinalMapper; - @ExperimentalApi public class ExactMatchDimFilter implements DimensionFilter { @@ -29,7 +26,7 @@ public class ExactMatchDimFilter implements DimensionFilter { private final List rawValues; - private List convertedOrdinals; + private TreeSet convertedOrdinals; public ExactMatchDimFilter(String dimensionName, List valuesToMatch) { this.dimensionName = dimensionName; @@ -37,39 +34,33 @@ public ExactMatchDimFilter(String dimensionName, List valuesToMatch) { } @Override - public void initialiseForSegment(LeafReaderContext leafReaderContext, StarTreeValues starTreeValues) { - convertedOrdinals = new ArrayList<>(); - List matchingDimensions = starTreeValues.getStarTreeField() - .getDimensionsOrder() - .stream() - .filter(x -> x.getField().equals(dimensionName)) - .collect(Collectors.toList()); - if (matchingDimensions.size() != 1) { - throw new IllegalStateException("Expected exactly one dimension but found " + matchingDimensions); - } - Dimension matchedDim = matchingDimensions.get(0); + public void initialiseForSegment(StarTreeValues starTreeValues) { + convertedOrdinals = new TreeSet<>(); + Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(dimensionName, starTreeValues); FieldToDimensionOrdinalMapper fieldToDimensionOrdinalMapper = getFieldToDimensionOrdinalMapper(matchedDim.getDocValuesType()); for (Object rawValue : rawValues) { - convertedOrdinals.add(fieldToDimensionOrdinalMapper.getOrdinal(matchedDim.getField(), rawValue, starTreeValues)); + long ordinal = fieldToDimensionOrdinalMapper.getMatchingOrdinal( + matchedDim.getField(), + rawValue, + starTreeValues, + FieldToDimensionOrdinalMapper.MatchType.EXACT + ); + if (ordinal >= 0) { + convertedOrdinals.add(ordinal); + } } } @Override public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeValues, StarTreeNodeCollector collector) throws IOException { - if (parentNode.getChildStarNode() != null) { - Dimension dimension = starTreeValues.getStarTreeField() - .getDimensionsOrder() - .get(parentNode.getChildStarNode().getDimensionId()); - if (dimension.getField().equals(dimensionName)) { - FieldToDimensionOrdinalMapper queryToDimensionOrdinalMapper = getFieldToDimensionOrdinalMapper(dimension.getDocValuesType()); - // TODO : [Optimisation] Implement storing the last searched StarTreeNode nodeId for successive binary search. - for (Object value : rawValues) { - collector.collectStarNode( - parentNode.getChildForDimensionValue( - queryToDimensionOrdinalMapper.getOrdinal(dimension.getField(), value, starTreeValues) - ) - ); + if (parentNode != null) { + // TODO : [Optimisation] Implement storing the last searched StarTreeNode nodeId for successive binary search. + StarTreeNode lastMatchedNode = null; + for (long ordinal : convertedOrdinals) { + lastMatchedNode = parentNode.getChildForDimensionValue(ordinal, lastMatchedNode); + if (lastMatchedNode != null) { + collector.collectStarNode(lastMatchedNode); } } } diff --git a/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java b/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java index da751992d4248..0c725c1f464cd 100644 --- a/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java +++ b/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java @@ -9,6 +9,7 @@ package org.opensearch.search.startree; import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.TermsEnum; import org.apache.lucene.util.BytesRef; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; @@ -22,24 +23,47 @@ @ExperimentalApi public interface FieldToDimensionOrdinalMapper { - long getOrdinal(String dimensionName, Object value, StarTreeValues starTreeValues); + long getMatchingOrdinal(String dimensionName, Object value, StarTreeValues starTreeValues, MatchType matchType); enum SingletonFactory { - NUMERIC_FIELD_MAPPER((dimensionName, value, starTreeValues) -> { + NUMERIC_FIELD_MAPPER((dimensionName, value, starTreeValues, matchType) -> { StarTreeValuesIterator genericIterator = starTreeValues.getDimensionValuesIterator(dimensionName); if (genericIterator instanceof SortedNumericStarTreeValuesIterator) { - return Long.parseLong(value.toString()); + long parsedValue = Long.parseLong(value.toString()); + switch (matchType) { + case GT: + return parsedValue + 1; + case GTE: + case EXACT: + case LTE: + return parsedValue; + case LT: + return parsedValue - 1; + default: + return -(parsedValue - 1); + } } else { throw new IllegalArgumentException("Unsupported star tree values iterator " + genericIterator.getClass().getName()); } }), - KEYWORD_FIELD_MAPPER((dimensionName, value, starTreeValues) -> { + KEYWORD_FIELD_MAPPER((dimensionName, value, starTreeValues, matchType) -> { StarTreeValuesIterator genericIterator = starTreeValues.getDimensionValuesIterator(dimensionName); if (genericIterator instanceof SortedSetStarTreeValuesIterator) { + SortedSetStarTreeValuesIterator sortedSetIterator = (SortedSetStarTreeValuesIterator) genericIterator; try { - return ((SortedSetStarTreeValuesIterator) genericIterator).lookupTerm((BytesRef) value); + if (matchType == MatchType.EXACT) { + return sortedSetIterator.lookupTerm((BytesRef) value); + } else { + TermsEnum termsEnum = sortedSetIterator.termsEnum(); + TermsEnum.SeekStatus seekStatus = termsEnum.seekCeil((BytesRef) value); + if (matchType == MatchType.GT || matchType == MatchType.GTE) { + return termsEnum.ord(); + } else { + return (seekStatus == TermsEnum.SeekStatus.FOUND) ? termsEnum.ord() : termsEnum.ord() - 1; + } + } } catch (IOException e) { throw new RuntimeException(e); } @@ -71,4 +95,13 @@ public static FieldToDimensionOrdinalMapper getFieldToDimensionOrdinalMapper(Doc } + @ExperimentalApi + enum MatchType { + GT, + LT, + GTE, + LTE, + EXACT; + } + } diff --git a/server/src/main/java/org/opensearch/search/startree/OlderStarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/OlderStarTreeFilter.java index edc74107aabde..203bdbe263854 100644 --- a/server/src/main/java/org/opensearch/search/startree/OlderStarTreeFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/OlderStarTreeFilter.java @@ -111,6 +111,14 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, Map { + TermsQueryBuilder termsQueryBuilder = (TermsQueryBuilder) rawFilter; + String field = termsQueryBuilder.fieldName(); + List matchedDimension = compositeFieldType.getDimensions() + .stream() + .filter(dim -> dim.getField().equals(field)) + .collect(Collectors.toList()); + // FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere. + if (matchedDimension.size() != 1 || matchedDimension.get(0).getDocValuesType() != DocValuesType.SORTED_NUMERIC) { + return null; + } + return new StarTreeFilter(Map.of(field, List.of(new ExactMatchDimFilter(field, termsQueryBuilder.values())))); + }, + RangeQueryBuilder.class, + (rawFilter, compositeFieldType) -> { + RangeQueryBuilder rangeQueryBuilder = (RangeQueryBuilder) rawFilter; + String field = rangeQueryBuilder.fieldName(); + List matchedDimensions = compositeFieldType.getDimensions() + .stream() + .filter(dim -> dim.getField().equals(field)) + .collect(Collectors.toList()); + // FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere. + if (matchedDimensions.size() != 1 || matchedDimensions.get(0).getDocValuesType() != DocValuesType.SORTED_NUMERIC) { + return null; + } + Dimension matchedDimension = matchedDimensions.get(0); + return new StarTreeFilter( + Map.of( + field, + List.of( + new RangeMatchDimFilter( + matchedDimension.getField(), + rangeQueryBuilder.from(), + rangeQueryBuilder.to(), + rangeQueryBuilder.includeLower(), + rangeQueryBuilder.includeUpper() + ) + ) + ) + ); } ); diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java index 004d758421d75..383979966a3d5 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java @@ -38,7 +38,7 @@ public class StarTreeQueryContext { * This is used to cache the results for each leaf reader context * to avoid reading the filtered values from the leaf reader context multiple times */ - private final Map> starTreeValues; + private final FixedBitSet[] starTreeValues; private final QueryBuilder baseQueryBuilder; @@ -55,7 +55,7 @@ public StarTreeQueryContext(SearchContext context, QueryBuilder baseQueryBuilder boolean cacheStarTreeValues = context.aggregations().factories().getFactories().length > 1; int cacheSize = cacheStarTreeValues ? context.indexShard().segments(false).size() : -1; if (cacheSize > -1) { - starTreeValues = new HashMap<>(cacheSize); + starTreeValues = new FixedBitSet[cacheSize]; } else { starTreeValues = null; } @@ -65,9 +65,14 @@ public CompositeIndexFieldInfo getStarTree() { return new CompositeIndexFieldInfo(compositeMappedFieldType.name(), compositeMappedFieldType.getCompositeIndexType()); } - public FixedBitSet getStarTreeValues(LeafReaderContext ctx, Aggregator aggregator) { - Map segmentCache = starTreeValues.get(ctx.ord); - return segmentCache != null ? segmentCache.get(getAggregatorKey(aggregator)) : null; + public FixedBitSet getStarTreeValues(LeafReaderContext ctx) { + return starTreeValues != null ? starTreeValues[ctx.ord] : null; + } + + public void setStarTreeValues(LeafReaderContext ctx, FixedBitSet values) { + if (starTreeValues != null) { + starTreeValues[ctx.ord] = values; + } } public void registerQuery(QueryBuilder query, AggregatorFactory aggregatorFactory) { diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java index 54466af0258cf..722287494ade9 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java @@ -254,15 +254,15 @@ public void collect(int doc, long bucket) { public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { // TODO : Uncomment and implement caching in new STQC - // FixedBitSet result = context.getStarTreeQueryContext().getStarTreeValues(ctx); - // if (result == null) { - return OlderStarTreeFilter.getStarTreeResult2( - starTreeValues, - context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter() - ); - // context.getStarTreeQueryContext().setStarTreeValues(ctx, result); - // } - // return result; + FixedBitSet result = context.getQueryShardContext().getStarTreeQueryContext().getStarTreeValues(ctx); + if (result == null) { + result = OlderStarTreeFilter.getStarTreeResult2( + starTreeValues, + context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter() + ); + } + context.getQueryShardContext().getStarTreeQueryContext().setStarTreeValues(ctx, result); + return result; } public static Set traverseStarTree(StarTreeValues starTreeValues, Map> dimensionFilterMap) @@ -360,6 +360,18 @@ private static StarTreeNode reachClosestParent(StarTreeNode startNode, int dimen return currentNode; } + public static Dimension getMatchingDimensionOrError(String dimensionName, StarTreeValues starTreeValues) { + List matchingDimensions = starTreeValues.getStarTreeField() + .getDimensionsOrder() + .stream() + .filter(x -> x.getField().equals(dimensionName)) + .collect(Collectors.toList()); + if (matchingDimensions.size() != 1) { + throw new IllegalStateException("Expected exactly one dimension but found " + matchingDimensions); + } + return matchingDimensions.get(0); + } + static class UnMatchedDocIdSet { private final int minDimensionIdUnMatched; diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java index 0fbe04c49b0db..5fbc61b3fd9cb 100644 --- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java +++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java @@ -149,7 +149,7 @@ import org.opensearch.search.internal.ContextIndexSearcher; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.lookup.SearchLookup; -import org.opensearch.search.startree.StarTreeQueryContext; +import org.opensearch.search.startree.OlderStarTreeQueryContext; import org.opensearch.test.InternalAggregationTestCase; import org.opensearch.test.OpenSearchTestCase; import org.junit.After; @@ -432,9 +432,9 @@ protected SearchContext createSearchContextWithStarTreeContext( when(searchContext.mapperService()).thenReturn(mapperService); SearchSourceBuilder sb = new SearchSourceBuilder().query(queryBuilder); - StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(searchContext, sb); + OlderStarTreeQueryContext olderStarTreeQueryContext = StarTreeQueryHelper.getOlderStarTreeQueryContext(searchContext, sb); - when(searchContext.getStarTreeQueryContext()).thenReturn(starTreeQueryContext); + when(searchContext.getStarTreeQueryContext()).thenReturn(olderStarTreeQueryContext); return searchContext; }