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 06fcae5642e78..71c07a7c9dd32 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 @@ -32,6 +32,8 @@ import org.opensearch.index.codec.composite.CompositeIndexReader; import org.opensearch.index.codec.composite.composite99.Composite99Codec; import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests; +import org.opensearch.index.compositeindex.datacube.Dimension; +import org.opensearch.index.compositeindex.datacube.NumericDimension; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.NumberFieldMapper; @@ -55,6 +57,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.LinkedList; import java.util.List; import java.util.Random; import java.util.function.BiConsumer; @@ -149,21 +152,58 @@ public void testStarTreeDocValues() throws IOException { ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME); AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME); + List supportedDimensions = new LinkedList<>(); + supportedDimensions.add(new NumericDimension(SNDV)); + supportedDimensions.add(new NumericDimension(DV)); + Query query = new MatchAllDocsQuery(); // match-all query QueryBuilder queryBuilder = null; // no predicates - testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue)); - testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue)); - testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue)); + testCase( + indexSearcher, + query, + queryBuilder, + sumAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalSum::getValue) + ); + testCase( + indexSearcher, + query, + queryBuilder, + maxAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalMax::getValue) + ); + testCase( + indexSearcher, + query, + queryBuilder, + minAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalMin::getValue) + ); testCase( indexSearcher, query, queryBuilder, valueCountAggregationBuilder, starTree, + supportedDimensions, verifyAggregation(InternalValueCount::getValue) ); - testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue)); + testCase( + indexSearcher, + query, + queryBuilder, + avgAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalAvg::getValue) + ); // Numeric-terms query for (int cases = 0; cases < 100; cases++) { @@ -180,18 +220,51 @@ public void testStarTreeDocValues() throws IOException { query = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue); queryBuilder = new TermQueryBuilder(queryField, queryValue); - testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue)); - testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue)); - testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue)); + testCase( + indexSearcher, + query, + queryBuilder, + sumAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalSum::getValue) + ); + testCase( + indexSearcher, + query, + queryBuilder, + maxAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalMax::getValue) + ); + testCase( + indexSearcher, + query, + queryBuilder, + minAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalMin::getValue) + ); testCase( indexSearcher, query, queryBuilder, valueCountAggregationBuilder, starTree, + supportedDimensions, verifyAggregation(InternalValueCount::getValue) ); - testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue)); + testCase( + indexSearcher, + query, + queryBuilder, + avgAggregationBuilder, + starTree, + supportedDimensions, + verifyAggregation(InternalAvg::getValue) + ); } ir.close(); @@ -212,6 +285,7 @@ private void testC QueryBuilder queryBuilder, T aggBuilder, CompositeIndexFieldInfo starTree, + List supportedDimensions, BiConsumer verify ) throws IOException { V starTreeAggregation = searchAndReduceStarTree( @@ -221,6 +295,7 @@ private void testC queryBuilder, aggBuilder, starTree, + supportedDimensions, DEFAULT_MAX_BUCKETS, false, DEFAULT_MAPPED_FIELD @@ -232,6 +307,7 @@ private void testC queryBuilder, aggBuilder, null, + null, DEFAULT_MAX_BUCKETS, false, DEFAULT_MAPPED_FIELD 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 e1069bb0ea519..891e2bcfc0df4 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 @@ -93,14 +93,34 @@ import org.opensearch.index.cache.query.DisabledQueryCache; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.compositeindex.datacube.Dimension; -import org.opensearch.index.compositeindex.datacube.NumericDimension; import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper; import org.opensearch.index.fielddata.IndexFieldData; import org.opensearch.index.fielddata.IndexFieldDataCache; import org.opensearch.index.fielddata.IndexFieldDataService; -import org.opensearch.index.mapper.*; +import org.opensearch.index.mapper.BinaryFieldMapper; +import org.opensearch.index.mapper.CompletionFieldMapper; +import org.opensearch.index.mapper.CompositeMappedFieldType; +import org.opensearch.index.mapper.ConstantKeywordFieldMapper; +import org.opensearch.index.mapper.ContentPath; +import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.mapper.DerivedFieldMapper; +import org.opensearch.index.mapper.FieldAliasMapper; +import org.opensearch.index.mapper.FieldMapper; +import org.opensearch.index.mapper.GeoPointFieldMapper; +import org.opensearch.index.mapper.GeoShapeFieldMapper; +import org.opensearch.index.mapper.KeywordFieldMapper; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.index.mapper.Mapper; import org.opensearch.index.mapper.Mapper.BuilderContext; +import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.MatchOnlyTextFieldMapper; +import org.opensearch.index.mapper.NumberFieldMapper; +import org.opensearch.index.mapper.ObjectMapper; import org.opensearch.index.mapper.ObjectMapper.Nested; +import org.opensearch.index.mapper.RangeFieldMapper; +import org.opensearch.index.mapper.RangeType; +import org.opensearch.index.mapper.StarTreeMapper; +import org.opensearch.index.mapper.TextFieldMapper; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryShardContext; import org.opensearch.index.shard.IndexShard; @@ -139,7 +159,6 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; -import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Objects; @@ -327,19 +346,26 @@ protected CountingAggregator createCountingAggregator( IndexSearcher indexSearcher, IndexSettings indexSettings, CompositeIndexFieldInfo starTree, + List supportedDimensions, MultiBucketConsumer bucketConsumer, MappedFieldType... fieldTypes ) throws IOException { SearchContext searchContext; if (starTree != null) { - searchContext = createSearchContextWithStarTreeContext(indexSearcher, indexSettings, query, queryBuilder, starTree, bucketConsumer, fieldTypes); + searchContext = createSearchContextWithStarTreeContext( + indexSearcher, + indexSettings, + query, + queryBuilder, + starTree, + supportedDimensions, + bucketConsumer, + fieldTypes + ); } else { searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, fieldTypes); } - return new CountingAggregator( - new AtomicInteger(), - createAggregator(aggregationBuilder, searchContext) - ); + return new CountingAggregator(new AtomicInteger(), createAggregator(aggregationBuilder, searchContext)); } /** @@ -361,27 +387,32 @@ protected SearchContext createSearchContextWithStarTreeContext( Query query, QueryBuilder queryBuilder, CompositeIndexFieldInfo starTree, + List supportedDimensions, MultiBucketConsumer bucketConsumer, MappedFieldType... fieldTypes ) throws IOException { - SearchContext searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes); + SearchContext searchContext = createSearchContext( + indexSearcher, + indexSettings, + query, + bucketConsumer, + new NoneCircuitBreakerService(), + fieldTypes + ); // Mock SearchContextAggregations SearchContextAggregations searchContextAggregations = mock(SearchContextAggregations.class); AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class); when(searchContext.aggregations()).thenReturn(searchContextAggregations); when(searchContextAggregations.factories()).thenReturn(aggregatorFactories); - when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[]{}); + when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] {}); StarTreeMapper.StarTreeFieldType compositeMappedFieldType = mock(StarTreeMapper.StarTreeFieldType.class); when(compositeMappedFieldType.name()).thenReturn(starTree.getField()); when(compositeMappedFieldType.getCompositeIndexType()).thenReturn(starTree.getType()); Set compositeFieldTypes = Set.of(compositeMappedFieldType); - List dimensions = new LinkedList<>(); - dimensions.add(new NumericDimension("sndv")); - dimensions.add(new NumericDimension("dv")); - when((compositeMappedFieldType).getDimensions()).thenReturn(dimensions); + when((compositeMappedFieldType).getDimensions()).thenReturn(supportedDimensions); MapperService mapperService = mock(MapperService.class); when(mapperService.getCompositeFieldTypes()).thenReturn(compositeFieldTypes); when(searchContext.mapperService()).thenReturn(mapperService); @@ -391,7 +422,6 @@ protected SearchContext createSearchContextWithStarTreeContext( when(searchContext.getStarTreeQueryContext()).thenReturn(starTreeQueryContext); return searchContext; - } protected SearchContext createSearchContext( @@ -708,6 +738,7 @@ protected A searchAndReduc QueryBuilder queryBuilder, AggregationBuilder builder, CompositeIndexFieldInfo compositeIndexFieldInfo, + List supportedDimensions, int maxBucket, boolean hasNested, MappedFieldType... fieldTypes @@ -724,7 +755,17 @@ protected A searchAndReduc maxBucket, new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST) ); - CountingAggregator countingAggregator = createCountingAggregator(query, queryBuilder, builder, searcher, indexSettings, compositeIndexFieldInfo, bucketConsumer, fieldTypes); + CountingAggregator countingAggregator = createCountingAggregator( + query, + queryBuilder, + builder, + searcher, + indexSettings, + compositeIndexFieldInfo, + supportedDimensions, + bucketConsumer, + fieldTypes + ); countingAggregator.preCollection(); searcher.search(query, countingAggregator);