From e386545e78cc0e0efd6857e38cbc7fbac3fc7750 Mon Sep 17 00:00:00 2001 From: expani Date: Mon, 27 Jan 2025 17:30:54 -0800 Subject: [PATCH] Increased UT coverage and addressed PR Comments Signed-off-by: expani --- .../node/FixedLengthStarTreeNode.java | 14 ---- .../datacube/startree/node/StarTreeNode.java | 13 ++- .../SortedSetStarTreeValuesIterator.java | 1 + .../search/startree/StarTreeQueryContext.java | 12 +-- .../search/startree/StarTreeQueryHelper.java | 8 +- .../startree/filter/ExactMatchDimFilter.java | 2 +- .../provider/StarTreeFilterProvider.java | 81 ++++++++----------- .../FixedLengthStarTreeNodeSearchTests.java | 48 ++++++----- .../search/SearchServiceStarTreeTests.java | 74 +++++++++++++++-- .../startree/ArrayBasedCollector.java | 39 +++++++++ .../DimensionFilterAndMapperTests.java | 81 +++++++++++++++++++ 11 files changed, 267 insertions(+), 106 deletions(-) create mode 100644 server/src/test/java/org/opensearch/search/aggregations/startree/ArrayBasedCollector.java 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 26390faa20dc1..c6c4993290c16 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 @@ -192,20 +192,6 @@ public StarTreeNode getChildStarNode() throws IOException { return handleStarNode(); } - @Override - public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException { - // there will be no children for leaf nodes - if (isLeaf()) { - return null; - } - - StarTreeNode resultStarTreeNode = null; - if (null != 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 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 dfdbdff09a988..9e8ded7b65eab 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 @@ -108,8 +108,17 @@ public interface StarTreeNode { * @return the child node for the given dimension value or null if child is not present * @throws IOException if an I/O error occurs while retrieving the child node */ - StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException; + default StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException { + return getChildForDimensionValue(dimensionValue, null); + } + /** + * Matches the given @dimensionValue amongst the child default nodes for this node. + * @param dimensionValue : Value to match + * @param lastMatchedChild : If not null, binary search will use this as the start/low + * @return : Matched StarTreeNode or null if not found + * @throws IOException : Any exception in reading the node data from index. + */ StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException; /** @@ -117,7 +126,7 @@ public interface StarTreeNode { * @param low : Starting of the range ( inclusive ) * @param high : End of the range ( inclusive ) * @param collector : Collector to collect the matched child StarTreeNode's - * @throws IOException : + * @throws IOException : Any exception in reading the node data from index. */ void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector) throws IOException; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java index f312dc48a9fd3..1605bd9cfc014 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java @@ -39,6 +39,7 @@ public boolean advanceExact(int target) throws IOException { return ((SortedSetDocValues) docIdSetIterator).advanceExact(target); } + // TODO : Remove this and merge @org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator to use value() public long nextOrd() throws IOException { return ((SortedSetDocValues) docIdSetIterator).nextOrd(); } 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 0ad761eeb77e5..3cb914d3e3e70 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java @@ -8,7 +8,6 @@ package org.opensearch.search.startree; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.util.FixedBitSet; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; @@ -68,6 +67,7 @@ public StarTreeQueryContext(SearchContext context, QueryBuilder baseQueryBuilder } } + // TODO : Make changes to change visibility into package private. Handle the same in @org.opensearch.search.SearchServiceStarTreeTests public StarTreeQueryContext(CompositeDataCubeFieldType compositeMappedFieldType, QueryBuilder baseQueryBuilder, int cacheSize) { this.compositeMappedFieldType = compositeMappedFieldType; this.baseQueryBuilder = baseQueryBuilder; @@ -82,17 +82,17 @@ public CompositeIndexFieldInfo getStarTree() { return new CompositeIndexFieldInfo(compositeMappedFieldType.name(), compositeMappedFieldType.getCompositeIndexType()); } - public FixedBitSet getStarTreeValue(LeafReaderContext ctx) { - return starTreeValues != null ? starTreeValues[ctx.ord] : null; + public FixedBitSet maybeGetCachedNodeIdsForSegment(int ordinal) { + return starTreeValues != null ? starTreeValues[ordinal] : null; } - public FixedBitSet[] getStarTreeValues() { + public FixedBitSet[] getAllCachedValues() { return starTreeValues; } - public void setStarTreeValues(LeafReaderContext ctx, FixedBitSet values) { + public void maybeSetCachedNodeIdsForSegment(int key, FixedBitSet values) { if (starTreeValues != null) { - starTreeValues[ctx.ord] = values; + starTreeValues[key] = values; } } 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 b0fc8896b1b94..4213e339b41a3 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java @@ -81,13 +81,11 @@ public static LeafBucketCollector getStarTreeLeafCollector( String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric); - assert starTreeValues != null; SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator( metricName ); // Obtain a FixedBitSet of matched star tree document IDs FixedBitSet filteredValues = getStarTreeFilteredValues(context, ctx, starTreeValues); - assert filteredValues != null; int numBits = filteredValues.length(); // Get the number of the filtered values (matching docs) if (numBits > 0) { @@ -122,7 +120,7 @@ public static LeafBucketCollector getStarTreeLeafCollector( */ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { - FixedBitSet result = context.getQueryShardContext().getStarTreeQueryContext().getStarTreeValue(ctx); + FixedBitSet result = context.getQueryShardContext().getStarTreeQueryContext().maybeGetCachedNodeIdsForSegment(ctx.ord); if (result == null) { result = StarTreeTraversalUtil.getStarTreeResult( starTreeValues, @@ -130,11 +128,11 @@ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafR context ); } - context.getQueryShardContext().getStarTreeQueryContext().setStarTreeValues(ctx, result); + context.getQueryShardContext().getStarTreeQueryContext().maybeSetCachedNodeIdsForSegment(ctx.ord, result); return result; } - public static Dimension getMatchingDimensionOrError(String dimensionName, List orderedDimensions) { + public static Dimension getMatchingDimensionOrThrow(String dimensionName, List orderedDimensions) { Dimension matchingDimension = getMatchingDimensionOrNull(dimensionName, orderedDimensions); if (matchingDimension == null) { throw new IllegalStateException("No matching dimension found for [" + dimensionName + "]"); 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 fe221856e5688..28ea261ca1e56 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 @@ -43,7 +43,7 @@ public ExactMatchDimFilter(String dimensionName, List valuesToMatch) { @Override public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext searchContext) { convertedOrdinals = new TreeSet<>(); - Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError( + Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrThrow( dimensionName, starTreeValues.getStarTreeField().getDimensionsOrder() ); diff --git a/server/src/main/java/org/opensearch/search/startree/filter/provider/StarTreeFilterProvider.java b/server/src/main/java/org/opensearch/search/startree/filter/provider/StarTreeFilterProvider.java index 9932c974cefd6..61f44ba3f163e 100644 --- a/server/src/main/java/org/opensearch/search/startree/filter/provider/StarTreeFilterProvider.java +++ b/server/src/main/java/org/opensearch/search/startree/filter/provider/StarTreeFilterProvider.java @@ -8,8 +8,6 @@ package org.opensearch.search.startree.filter.provider; -import org.apache.lucene.search.MatchNoDocsQuery; -import org.apache.lucene.search.Query; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.mapper.CompositeDataCubeFieldType; @@ -53,20 +51,20 @@ StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, Composit */ class SingletonFactory { - private static final Map, StarTreeFilterProvider> QUERY_BUILDERS_TO_STF_PROVIDER = Map.of( - MatchAllQueryBuilder.class, + private static final Map QUERY_BUILDERS_TO_STF_PROVIDER = Map.of( + MatchAllQueryBuilder.NAME, MATCH_ALL_PROVIDER, - TermQueryBuilder.class, + TermQueryBuilder.NAME, new TermStarTreeFilterProvider(), - TermsQueryBuilder.class, + TermsQueryBuilder.NAME, new TermsStarTreeFilterProvider(), - RangeQueryBuilder.class, + RangeQueryBuilder.NAME, new RangeStarTreeFilterProvider() ); public static StarTreeFilterProvider getProvider(QueryBuilder query) { if (query != null) { - return QUERY_BUILDERS_TO_STF_PROVIDER.get(query.getClass()); + return QUERY_BUILDERS_TO_STF_PROVIDER.get(query.getName()); } return MATCH_ALL_PROVIDER; } @@ -83,23 +81,16 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C TermQueryBuilder termQueryBuilder = (TermQueryBuilder) rawFilter; String field = termQueryBuilder.fieldName(); MappedFieldType mappedFieldType = context.mapperService().fieldType(field); - DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType); + DimensionFilterMapper dimensionFilterMapper = mappedFieldType != null + ? DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType) + : null; Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions()); if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) { return null; // Indicates Aggregators to fallback to default implementation. } else { - // FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere. - Query query = termQueryBuilder.toQuery(context.getQueryShardContext()); - if (query instanceof MatchNoDocsQuery) { - return new StarTreeFilter(Collections.emptyMap()); - } else { - return new StarTreeFilter( - Map.of( - field, - List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, List.of(termQueryBuilder.value()))) - ) - ); - } + return new StarTreeFilter( + Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, List.of(termQueryBuilder.value())))) + ); } } } @@ -115,18 +106,15 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C String field = termsQueryBuilder.fieldName(); Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions()); MappedFieldType mappedFieldType = context.mapperService().fieldType(field); - DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType); + DimensionFilterMapper dimensionFilterMapper = mappedFieldType != null + ? DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType) + : null; if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) { return null; // Indicates Aggregators to fallback to default implementation. } else { - Query query = termsQueryBuilder.toQuery(context.getQueryShardContext()); - if (query instanceof MatchNoDocsQuery) { - return new StarTreeFilter(Collections.emptyMap()); - } else { - return new StarTreeFilter( - Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, termsQueryBuilder.values()))) - ); - } + return new StarTreeFilter( + Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, termsQueryBuilder.values()))) + ); } } } @@ -143,29 +131,26 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C String field = rangeQueryBuilder.fieldName(); Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions()); MappedFieldType mappedFieldType = context.mapperService().fieldType(field); - DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType); + DimensionFilterMapper dimensionFilterMapper = mappedFieldType == null + ? null + : DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType); if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) { return null; } else { - Query query = rangeQueryBuilder.toQuery(context.getQueryShardContext()); - if (query instanceof MatchNoDocsQuery) { - return new StarTreeFilter(Collections.emptyMap()); - } else { - return new StarTreeFilter( - Map.of( - field, - List.of( - dimensionFilterMapper.getRangeMatchFilter( - mappedFieldType, - rangeQueryBuilder.from(), - rangeQueryBuilder.to(), - rangeQueryBuilder.includeLower(), - rangeQueryBuilder.includeUpper() - ) + return new StarTreeFilter( + Map.of( + field, + List.of( + dimensionFilterMapper.getRangeMatchFilter( + mappedFieldType, + rangeQueryBuilder.from(), + rangeQueryBuilder.to(), + rangeQueryBuilder.includeLower(), + rangeQueryBuilder.includeUpper() ) ) - ); - } + ) + ); } } 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 dfa5823e2b21a..96e96d744ab08 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 @@ -24,12 +24,10 @@ import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeFactory; -import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode; -import org.opensearch.search.startree.StarTreeNodeCollector; +import org.opensearch.search.aggregations.startree.ArrayBasedCollector; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; -import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.TreeSet; @@ -314,27 +312,27 @@ private void createStarTreeForDimension( } - private static class ArrayBasedCollector implements StarTreeNodeCollector { - - private final List nodes = new ArrayList<>(); - - @Override - public void collectStarTreeNode(StarTreeNode node) { - nodes.add(node); - } - - public boolean matchValues(long[] values) throws IOException { - boolean matches = true; - for (int i = 0; i < values.length; i++) { - matches &= nodes.get(i).getDimensionValue() == values[i]; - } - return matches; - } - - public int collectedNodeCount() { - return nodes.size(); - } - - } + // private static class ArrayBasedCollector implements StarTreeNodeCollector { + // + // private final List nodes = new ArrayList<>(); + // + // @Override + // public void collectStarTreeNode(StarTreeNode node) { + // nodes.add(node); + // } + // + // public boolean matchValues(long[] values) throws IOException { + // boolean matches = true; + // for (int i = 0; i < values.length; i++) { + // matches &= nodes.get(i).getDimensionValue() == values[i]; + // } + // return matches; + // } + // + // public int collectedNodeCount() { + // return nodes.size(); + // } + // + // } } diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java index 42af6c147cf69..caaf7ad669ea2 100644 --- a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java +++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java @@ -8,6 +8,7 @@ package org.opensearch.search; +import org.apache.lucene.util.FixedBitSet; import org.opensearch.action.OriginalIndices; import org.opensearch.action.admin.indices.create.CreateIndexRequestBuilder; import org.opensearch.action.search.SearchRequest; @@ -29,9 +30,11 @@ import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings; import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.engine.Segment; import org.opensearch.index.mapper.CompositeDataCubeFieldType; import org.opensearch.index.mapper.CompositeMappedFieldType; import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.StarTreeMapper; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.QueryBuilder; @@ -61,6 +64,7 @@ import java.io.IOException; import java.util.Collections; import java.util.List; +import java.util.Set; import static org.opensearch.search.aggregations.AggregationBuilders.dateHistogram; import static org.opensearch.search.aggregations.AggregationBuilders.max; @@ -137,7 +141,7 @@ public void testQueryParsingForMetricAggregations() throws IOException { // Case 2: MatchAllQuery present but no aggregations, should not use star tree sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder()); assertStarTreeContext(request, sourceBuilder, null, -1); - + // Case 3: MatchAllQuery and aggregations present, should use star tree baseQuery = new MatchAllQueryBuilder(); sourceBuilder = new SearchSourceBuilder().size(0).query(baseQuery).aggregation(AggregationBuilders.max("test").field("field")); @@ -270,14 +274,14 @@ public void testQueryParsingForDateHistogramAggregations() throws IOException { SumAggregationBuilder sumAggSub = sum("sum").field(FIELD_NAME).subAggregation(maxAggNoSub); MedianAbsoluteDeviationAggregationBuilder medianAgg = medianAbsoluteDeviation("median").field(FIELD_NAME); - QueryBuilder baseQuery; - SearchContext searchContext = createSearchContext(indexService); StarTreeFieldConfiguration starTreeFieldConfiguration = new StarTreeFieldConfiguration( 1, Collections.emptySet(), StarTreeFieldConfiguration.StarTreeBuildMode.ON_HEAP ); + QueryBuilder baseQuery; + SearchContext searchContext = createSearchContext(indexService); // Case 1: No query or aggregations, should not use star tree SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); assertStarTreeContext(request, sourceBuilder, null, -1); @@ -423,6 +427,65 @@ public void testQueryParsingForDateHistogramAggregations() throws IOException { setStarTreeIndexSetting(null); } + public void testCacheCreationInStarTreeQueryContext() throws IOException { + FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build()); + setStarTreeIndexSetting("true"); + StarTreeFieldConfiguration starTreeFieldConfiguration = new StarTreeFieldConfiguration( + 1, + Collections.emptySet(), + StarTreeFieldConfiguration.StarTreeBuildMode.ON_HEAP + ); + CompositeDataCubeFieldType compositeDataCubeFieldType = new StarTreeMapper.StarTreeFieldType( + "star_tree", + new StarTreeField( + "star_tree", + List.of(new OrdinalDimension("field")), + List.of(new Metric("metricField", List.of(MetricStat.SUM, MetricStat.MAX))), + starTreeFieldConfiguration + ) + ); + + QueryBuilder baseQuery = new MatchAllQueryBuilder(); + SearchContext searchContext = mock(SearchContext.class); + MapperService mapperService = mock(MapperService.class); + IndexShard indexShard = mock(IndexShard.class); + Segment segment = mock(Segment.class); + SearchContextAggregations searchContextAggregations = mock(SearchContextAggregations.class); + AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class); + + when(mapperService.getCompositeFieldTypes()).thenReturn(Set.of(compositeDataCubeFieldType)); + when(searchContext.mapperService()).thenReturn(mapperService); + when(searchContext.indexShard()).thenReturn(indexShard); + when(indexShard.segments(false)).thenReturn(List.of(segment, segment)); + when(searchContext.aggregations()).thenReturn(searchContextAggregations); + when(searchContextAggregations.factories()).thenReturn(aggregatorFactories); + when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] { null, null }); + StarTreeQueryContext starTreeQueryContext = new StarTreeQueryContext(searchContext, baseQuery); + + assertEquals(2, starTreeQueryContext.getAllCachedValues().length); + + // Asserting null values are ignored + when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] {}); + starTreeQueryContext = new StarTreeQueryContext(searchContext, baseQuery); + starTreeQueryContext.maybeSetCachedNodeIdsForSegment(-1, null); + assertNull(starTreeQueryContext.getAllCachedValues()); + assertNull(starTreeQueryContext.maybeGetCachedNodeIdsForSegment(0)); + + // Assert correct cached value is returned + when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] { null, null }); + starTreeQueryContext = new StarTreeQueryContext(searchContext, baseQuery); + FixedBitSet cachedValues = new FixedBitSet(22); + starTreeQueryContext.maybeSetCachedNodeIdsForSegment(0, cachedValues); + assertEquals(2, starTreeQueryContext.getAllCachedValues().length); + assertEquals(22, starTreeQueryContext.maybeGetCachedNodeIdsForSegment(0).length()); + + starTreeQueryContext = new StarTreeQueryContext(compositeDataCubeFieldType, new MatchAllQueryBuilder(), 2); + assertEquals(2, starTreeQueryContext.getAllCachedValues().length); + + setStarTreeIndexSetting(null); + mapperService.close(); + } + /** * Test query parsing for date histogram aggregations on star-tree index when @timestamp field does not exist */ @@ -507,9 +570,9 @@ private void assertStarTreeContext( actualContext.getBaseQueryStarTreeFilter().getDimensions() ); if (expectedCacheUsage > -1) { - assertEquals(expectedCacheUsage, actualContext.getStarTreeValues().length); + assertEquals(expectedCacheUsage, actualContext.getAllCachedValues().length); } else { - assertNull(actualContext.getStarTreeValues()); + assertNull(actualContext.getAllCachedValues()); } } searchService.doStop(); @@ -547,6 +610,7 @@ private StarTreeQueryContext getStarTreeQueryContext( boolean consolidated = starTreeQueryContext.consolidateAllFilters(searchContext); if (assertConsolidation) { assertTrue(consolidated); + searchContext.getQueryShardContext().setStarTreeQueryContext(starTreeQueryContext); } return starTreeQueryContext; } diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/ArrayBasedCollector.java b/server/src/test/java/org/opensearch/search/aggregations/startree/ArrayBasedCollector.java new file mode 100644 index 0000000000000..60d07b41bf632 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/aggregations/startree/ArrayBasedCollector.java @@ -0,0 +1,39 @@ +/* + * 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.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode; +import org.opensearch.search.startree.StarTreeNodeCollector; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +public class ArrayBasedCollector implements StarTreeNodeCollector { + + private final List nodes = new ArrayList<>(); + + @Override + public void collectStarTreeNode(StarTreeNode node) { + nodes.add(node); + } + + public boolean matchValues(long[] values) throws IOException { + boolean matches = true; + for (int i = 0; i < values.length; i++) { + matches &= nodes.get(i).getDimensionValue() == values[i]; + } + return matches; + } + + public int collectedNodeCount() { + return nodes.size(); + } + +} diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/DimensionFilterAndMapperTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/DimensionFilterAndMapperTests.java index 824e2b6c36d81..e89bc8e60e9da 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/startree/DimensionFilterAndMapperTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/startree/DimensionFilterAndMapperTests.java @@ -10,14 +10,33 @@ import org.apache.lucene.index.TermsEnum; import org.apache.lucene.util.BytesRef; +import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.OrdinalDimension; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedSetStarTreeValuesIterator; +import org.opensearch.index.mapper.CompositeDataCubeFieldType; import org.opensearch.index.mapper.KeywordFieldMapper; +import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.StarTreeMapper; +import org.opensearch.index.mapper.WildcardFieldMapper; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.RangeQueryBuilder; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.index.query.TermsQueryBuilder; +import org.opensearch.search.internal.SearchContext; +import org.opensearch.search.startree.filter.DimensionFilter; import org.opensearch.search.startree.filter.DimensionFilter.MatchType; +import org.opensearch.search.startree.filter.MatchNoneFilter; import org.opensearch.search.startree.filter.provider.DimensionFilterMapper; +import org.opensearch.search.startree.filter.provider.StarTreeFilterProvider; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; +import java.util.Collections; +import java.util.List; import java.util.Optional; import static org.mockito.Mockito.mock; @@ -106,7 +125,69 @@ public void testKeywordOrdinalMapping() throws IOException { matchingOrdinal = dimensionFilterMapper.getMatchingOrdinal("field", bytesRef, starTreeValues, matchType); assertFalse(matchingOrdinal.isPresent()); } + } + + public void testStarTreeFilterProviders() throws IOException { + CompositeDataCubeFieldType compositeDataCubeFieldType = new StarTreeMapper.StarTreeFieldType( + "star_tree", + new StarTreeField( + "star_tree", + List.of(new OrdinalDimension("keyword")), + List.of(new Metric("field", List.of(MetricStat.MAX))), + new StarTreeFieldConfiguration( + randomIntBetween(1, 10_000), + Collections.emptySet(), + StarTreeFieldConfiguration.StarTreeBuildMode.ON_HEAP + ) + ) + ); + MapperService mapperService = mock(MapperService.class); + SearchContext searchContext = mock(SearchContext.class); + when(searchContext.mapperService()).thenReturn(mapperService); + + // Null returned when mapper doesn't exist + assertNull(DimensionFilterMapper.Factory.fromMappedFieldType(new WildcardFieldMapper.WildcardFieldType("field"))); + + // Null returned for no mapped field type + assertNull(DimensionFilterMapper.Factory.fromMappedFieldType(null)); + + // Provider for null Query builder + assertEquals(StarTreeFilterProvider.MATCH_ALL_PROVIDER, StarTreeFilterProvider.SingletonFactory.getProvider(null)); + + QueryBuilder[] queryBuilders = new QueryBuilder[] { + new TermQueryBuilder("field", "value"), + new TermsQueryBuilder("field", List.of("value")), + new RangeQueryBuilder("field") }; + + for (QueryBuilder queryBuilder : queryBuilders) { + // Dimension Not Found + StarTreeFilterProvider provider = StarTreeFilterProvider.SingletonFactory.getProvider(queryBuilder); + assertNull(provider.getFilter(searchContext, queryBuilder, compositeDataCubeFieldType)); + } + + queryBuilders = new QueryBuilder[] { + new TermQueryBuilder("keyword", "value"), + new TermsQueryBuilder("keyword", List.of("value")), + new RangeQueryBuilder("keyword") }; + + for (QueryBuilder queryBuilder : queryBuilders) { + // Mapped field type not supported + StarTreeFilterProvider provider = StarTreeFilterProvider.SingletonFactory.getProvider(queryBuilder); + when(mapperService.fieldType("keyword")).thenReturn(new WildcardFieldMapper.WildcardFieldType("keyword")); + assertNull(provider.getFilter(searchContext, queryBuilder, compositeDataCubeFieldType)); + + // Unsupported Mapped Type + when(mapperService.fieldType("keyword")).thenReturn(null); + assertNull(provider.getFilter(searchContext, queryBuilder, compositeDataCubeFieldType)); + } + // Testing MatchNoneFilter + DimensionFilter dimensionFilter = new MatchNoneFilter(); + dimensionFilter.initialiseForSegment(null, null); + ArrayBasedCollector collector = new ArrayBasedCollector(); + assertFalse(dimensionFilter.matchDimValue(1, null)); + dimensionFilter.matchStarTreeNodes(null, null, collector); + assertEquals(0, collector.collectedNodeCount()); } }