diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java index bbee8329a4821..a5acea004c3b2 100644 --- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java +++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java @@ -107,7 +107,7 @@ public class FeatureFlags { * aggregations. */ public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled"; - public static final Setting STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope); + public static final Setting STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope); /** * Gates the functionality of application based configuration templates. diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java index 35b1f1997bcfe..ccaaaa597ec0f 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java @@ -54,24 +54,19 @@ public class StarTreeQueryHelper { /** * Checks if the search context can be supported by star-tree */ - public static boolean isStarTreeSupported(SearchContext context, boolean trackTotalHits) { - boolean canUseStarTree = context.aggregations() != null - && context.size() == 0 + public static boolean isStarTreeSupported(SearchContext context) { + return context.aggregations() != null && context.mapperService().isCompositeIndexPresent() && context.parsedPostFilter() == null && context.innerHits().getInnerHits().isEmpty() && context.sort() == null - && (!trackTotalHits || context.trackTotalHitsUpTo() == SearchContext.TRACK_TOTAL_HITS_DISABLED) && context.trackScores() == false - && context.minimumScore() == null - && context.terminateAfter() == 0; - return canUseStarTree; + && context.minimumScore() == null; } - /** - * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder. - * Returns null if the query cannot be supported. + * Gets StarTreeQueryContext from the search context and source builder. + * Returns null if the query & aggregation cannot be supported. */ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException { // Current implementation assumes only single star-tree is supported @@ -93,7 +88,6 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context return null; } - for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) { MetricStat metricStat = validateStarTreeMetricSupport(compositeMappedFieldType, aggregatorFactory); if (metricStat == null) { @@ -104,20 +98,22 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context if (context.aggregations().factories().getFactories().length > 1) { context.initializeStarTreeValuesMap(); } - return starTreeQueryContext; } + /** + * Uses query builder & composite index info to form star-tree query context + */ private static StarTreeQueryContext toStarTreeQueryContext( - CompositeIndexFieldInfo starTree, - CompositeDataCubeFieldType compositeIndexFieldInfo, + CompositeIndexFieldInfo compositeIndexFieldInfo, + CompositeDataCubeFieldType compositeFieldType, QueryBuilder queryBuilder ) { Map queryMap; if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) { queryMap = null; } else if (queryBuilder instanceof TermQueryBuilder) { - List supportedDimensions = compositeIndexFieldInfo.getDimensions() + List supportedDimensions = compositeFieldType.getDimensions() .stream() .map(Dimension::getField) .collect(Collectors.toList()); @@ -128,13 +124,12 @@ private static StarTreeQueryContext toStarTreeQueryContext( } else { return null; } - - return new StarTreeQueryContext(starTree, queryMap); + return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap); } /** * Parse query body to star-tree predicates - * @param queryBuilder to match supported query shape + * @param queryBuilder to match star-tree supported query shape * @return predicates to match */ private static Map getStarTreePredicates(QueryBuilder queryBuilder, List supportedDimensions) { @@ -185,6 +180,10 @@ public static StarTreeValues getStarTreeValues(LeafReaderContext context, Compos return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree); } + /** + * Get the star-tree leaf collector + * This collector computes the aggregation prematurely and invokes an early termination collector + */ public static LeafBucketCollector getStarTreeLeafCollector( SearchContext context, ValuesSource.Numeric valuesSource, @@ -225,7 +224,6 @@ public static LeafBucketCollector getStarTreeLeafCollector( } } - // Call the final consumer after processing all entries finalConsumer.run(); @@ -238,7 +236,11 @@ public void collect(int doc, long bucket) { }; } - public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { + /** + * Get the filtered values for the star-tree query + */ + public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) + throws IOException { if (context.getStarTreeValuesMap() != null && context.getStarTreeValuesMap().containsKey(ctx)) { return context.getStarTreeValuesMap().get(ctx); } diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java index feaafd9bbbc6c..74a7482d975df 100644 --- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java +++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java @@ -34,7 +34,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.BooleanClause.Occur; import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.BoostQuery; @@ -43,7 +42,6 @@ import org.apache.lucene.search.FieldDoc; import org.apache.lucene.search.MatchNoDocsQuery; import org.apache.lucene.search.Query; -import org.apache.lucene.util.FixedBitSet; import org.opensearch.Version; import org.opensearch.action.search.SearchShardTask; import org.opensearch.action.search.SearchType; @@ -58,7 +56,6 @@ import org.opensearch.index.IndexService; import org.opensearch.index.IndexSettings; import org.opensearch.index.cache.bitset.BitsetFilterCache; -import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.engine.Engine; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; @@ -101,8 +98,6 @@ import org.opensearch.search.rescore.RescoreContext; import org.opensearch.search.slice.SliceBuilder; import org.opensearch.search.sort.SortAndFormats; -import org.opensearch.search.startree.StarTreeFilter; -import org.opensearch.search.startree.StarTreeQueryContext; import org.opensearch.search.suggest.SuggestionSearchContext; import java.io.IOException; @@ -275,7 +270,6 @@ final class DefaultSearchContext extends SearchContext { this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold(); this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories; this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled(); - this.starTreeValuesMap = new HashMap<>(); } @Override diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index d38d1280c592e..eb111f7d43008 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -1544,7 +1544,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc if (this.indicesService.getCompositeIndexSettings() != null && this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled() - && StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) { + && StarTreeQueryHelper.isStarTreeSupported(context)) { try { StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source); if (starTreeQueryContext != null) { diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java index 4a8fe891ae066..de0fde8df58da 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java @@ -108,7 +108,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc } CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); if (supportedStarTree != null) { - return getStarTreeLeafCollector(ctx, sub, supportedStarTree); + return getStarTreeLeafCollector(ctx, sub, supportedStarTree); } return getDefaultLeafCollector(ctx, sub); } @@ -164,7 +164,6 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB MetricStat.VALUE_COUNT.getTypeName() ); - final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), 0); SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues .getMetricValuesIterator(sumMetricName); diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java index 13340fa10f849..41eb0b21390f6 100644 --- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java +++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java @@ -46,7 +46,6 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.BigArrays; import org.opensearch.index.cache.bitset.BitsetFilterCache; -import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.ObjectMapper; @@ -79,11 +78,9 @@ import org.opensearch.search.query.ReduceableSearchResult; import org.opensearch.search.rescore.RescoreContext; import org.opensearch.search.sort.SortAndFormats; -import org.opensearch.search.startree.StarTreeFilter; import org.opensearch.search.startree.StarTreeQueryContext; import org.opensearch.search.suggest.SuggestionSearchContext; -import java.io.IOException; import java.util.Collection; import java.util.HashMap; import java.util.List; @@ -538,7 +535,6 @@ public boolean keywordIndexOrDocValuesEnabled() { return false; } - public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) { this.starTreeQueryContext = starTreeQueryContext; return this; diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java index d49111609f95e..b80f1419be761 100644 --- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java +++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java @@ -35,8 +35,6 @@ import org.apache.lucene.index.FilterDirectoryReader; import org.apache.lucene.index.LeafReader; import org.apache.lucene.search.FieldDoc; -import org.apache.lucene.search.IndexOrDocValuesQuery; -import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.Query; import org.apache.lucene.store.AlreadyClosedException; import org.opensearch.OpenSearchException; @@ -78,10 +76,12 @@ import org.opensearch.index.IndexNotFoundException; import org.opensearch.index.IndexService; import org.opensearch.index.IndexSettings; +import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests; import org.opensearch.index.compositeindex.CompositeIndexSettings; import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings; import org.opensearch.index.engine.Engine; +import org.opensearch.index.mapper.CompositeMappedFieldType; import org.opensearch.index.mapper.DerivedFieldType; import org.opensearch.index.query.AbstractQueryBuilder; import org.opensearch.index.query.MatchAllQueryBuilder; @@ -121,6 +121,7 @@ import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.MinAndMax; import org.opensearch.search.sort.SortOrder; +import org.opensearch.search.startree.StarTreeQueryContext; import org.opensearch.search.suggest.SuggestBuilder; import org.opensearch.test.OpenSearchSingleNodeTestCase; import org.opensearch.threadpool.ThreadPool; @@ -151,6 +152,7 @@ import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.CoreMatchers.startsWith; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -2268,75 +2270,84 @@ public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IO assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, 1000), true); } -// public void testParseQueryToOriginalOrStarTreeQuery() throws IOException { -// FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build()); -// setStarTreeIndexSetting("true"); -// -// Settings settings = Settings.builder() -// .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) -// .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1) -// .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true) -// .build(); -// CreateIndexRequestBuilder builder = client().admin() -// .indices() -// .prepareCreate("test") -// .setSettings(settings) -// .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping()); -// createIndex("test", builder); -// -// IndicesService indicesService = getInstanceFromNode(IndicesService.class); -// IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test")); -// IndexShard indexShard = indexService.getShard(0); -// ShardSearchRequest request = new ShardSearchRequest( -// OriginalIndices.NONE, -// new SearchRequest().allowPartialSearchResults(true), -// indexShard.shardId(), -// 1, -// new AliasFilter(null, Strings.EMPTY_ARRAY), -// 1.0f, -// -1, -// null, -// null -// ); -// -// // Case 1: No query or aggregations, should not use star tree -// SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); -// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class); -// -// // Case 2: MatchAllQuery present but no aggregations, should not use star tree -// sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder()); -// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class); -// -// // Case 3: MatchAllQuery and aggregations present, should use star tree -// sourceBuilder = new SearchSourceBuilder().size(0) -// .query(new MatchAllQueryBuilder()) -// .aggregation(AggregationBuilders.max("test").field("field")); -// assertQueryType(request, sourceBuilder, StarTreeQuery.class); -// -// // Case 4: MatchAllQuery and aggregations present, but trackTotalHitsUpTo specified, should not use star tree -// sourceBuilder = new SearchSourceBuilder().size(0) -// .query(new MatchAllQueryBuilder()) -// .aggregation(AggregationBuilders.max("test").field("field")) -// .trackTotalHitsUpTo(1000); -// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class); -// -// // Case 5: TermQuery and aggregations present, should use star tree -// sourceBuilder = new SearchSourceBuilder().size(0) -// .query(new TermQueryBuilder("sndv", 1)) -// .aggregation(AggregationBuilders.max("test").field("field")); -// assertQueryType(request, sourceBuilder, StarTreeQuery.class); -// -// // Case 6: No query, metric aggregations present, should use star tree -// sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field")); -// assertQueryType(request, sourceBuilder, StarTreeQuery.class); -// -// // Case 7: TermQuery and aggregations present, size != 0, should not use star tree -// sourceBuilder = new SearchSourceBuilder().query(new TermQueryBuilder("sndv", 1)) -// .aggregation(AggregationBuilders.max("test").field("field")); -// assertQueryType(request, sourceBuilder, IndexOrDocValuesQuery.class); -// -// setStarTreeIndexSetting(null); -// } + public void testParseQueryToOriginalOrStarTreeQuery() throws IOException { + FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build()); + setStarTreeIndexSetting("true"); + + Settings settings = Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1) + .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true) + .build(); + CreateIndexRequestBuilder builder = client().admin() + .indices() + .prepareCreate("test") + .setSettings(settings) + .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping()); + createIndex("test", builder); + + IndicesService indicesService = getInstanceFromNode(IndicesService.class); + IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test")); + IndexShard indexShard = indexService.getShard(0); + ShardSearchRequest request = new ShardSearchRequest( + OriginalIndices.NONE, + new SearchRequest().allowPartialSearchResults(true), + indexShard.shardId(), + 1, + new AliasFilter(null, Strings.EMPTY_ARRAY), + 1.0f, + -1, + null, + null + ); + + // Case 1: No query or aggregations, should not use star tree + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + assertStarTreeContext(request, sourceBuilder, null, false); + + // Case 2: MatchAllQuery present but no aggregations, should not use star tree + sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder()); + assertStarTreeContext(request, sourceBuilder, null, false); + + // Case 3: MatchAllQuery and aggregations present, should use star tree + sourceBuilder = new SearchSourceBuilder().size(0) + .query(new MatchAllQueryBuilder()) + .aggregation(AggregationBuilders.max("test").field("field")); + CompositeIndexFieldInfo expectedStarTree = new CompositeIndexFieldInfo( + "startree", + CompositeMappedFieldType.CompositeFieldType.STAR_TREE + ); + Map expectedQueryMap = null; + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), false); + + // Case 4: MatchAllQuery and aggregations present, but trackScores true, should not use star tree + sourceBuilder = new SearchSourceBuilder().size(0) + .query(new MatchAllQueryBuilder()) + .aggregation(AggregationBuilders.max("test").field("field")) + .trackScores(true); + assertStarTreeContext(request, sourceBuilder, null, false); + + // Case 5: TermQuery and single aggregation, should use star tree, but not initialize query cache + sourceBuilder = new SearchSourceBuilder().size(0) + .query(new TermQueryBuilder("sndv", 1)) + .aggregation(AggregationBuilders.max("test").field("field")); + expectedQueryMap = Map.of("sndv", 1L); + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), false); + + // Case 6: TermQuery and multiple aggregations present, should use star tree & initialize cache + sourceBuilder = new SearchSourceBuilder().size(0) + .query(new TermQueryBuilder("sndv", 1)) + .aggregation(AggregationBuilders.max("test").field("field")) + .aggregation(AggregationBuilders.sum("test2").field("field")); + expectedQueryMap = Map.of("sndv", 1L); + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), true); + + // Case 7: No query, metric aggregations present, should use star tree + sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field")); + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, null), false); + + setStarTreeIndexSetting(null); + } private void setStarTreeIndexSetting(String value) throws IOException { client().admin() @@ -2346,14 +2357,28 @@ private void setStarTreeIndexSetting(String value) throws IOException { .execute(); } -// private void assertQueryType(ShardSearchRequest request, SearchSourceBuilder sourceBuilder, Class expectedQueryClass) -// throws IOException { -// request.source(sourceBuilder); -// SearchService searchService = getInstanceFromNode(SearchService.class); -// try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) { -// SearchContext context = searchService.createContext(reader, request, null, true); -// assertThat(context.query(), instanceOf(expectedQueryClass)); -// searchService.doStop(); -// } -// } + private void assertStarTreeContext( + ShardSearchRequest request, + SearchSourceBuilder sourceBuilder, + StarTreeQueryContext expectedContext, + boolean expectedCacheUsage + ) throws IOException { + request.source(sourceBuilder); + SearchService searchService = getInstanceFromNode(SearchService.class); + try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) { + SearchContext context = searchService.createContext(reader, request, null, true); + StarTreeQueryContext actualContext = context.getStarTreeQueryContext(); + + if (expectedContext == null) { + assertThat(context.getStarTreeQueryContext(), nullValue()); + } else { + assertThat(actualContext, notNullValue()); + assertEquals(expectedContext.getStarTree().getType(), actualContext.getStarTree().getType()); + assertEquals(expectedContext.getStarTree().getField(), actualContext.getStarTree().getField()); + assertEquals(expectedContext.getQueryMap(), actualContext.getQueryMap()); + } + assertThat(context.getStarTreeValuesMap(), expectedCacheUsage ? notNullValue() : nullValue()); + searchService.doStop(); + } + } } 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 eba7036550987..4d7a8a9f9c644 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,11 +32,9 @@ 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.mapper.CompositeDataCubeFieldType; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.NumberFieldMapper; -import org.opensearch.index.mapper.StarTreeMapper; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.search.aggregations.AggregationBuilder; @@ -58,9 +56,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.List; -import java.util.Map; import java.util.Random; -import java.util.Set; import java.util.function.BiConsumer; import java.util.function.Function; @@ -98,136 +94,6 @@ protected Codec getCodec() { return new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger); } -// public void testStarTreeDocValues1() throws IOException { -// Directory directory = newDirectory(); -// IndexWriterConfig conf = newIndexWriterConfig(null); -// conf.setCodec(getCodec()); -// conf.setMergePolicy(newLogMergePolicy()); -// RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf); -// -// Random random = RandomizedTest.getRandom(); -// int totalDocs = 100; -// final String SNDV = "sndv"; -// final String DV = "dv"; -// int val; -// -// // Index 100 random documents -// for (int i = 0; i < totalDocs; i++) { -// Document doc = new Document(); -// if (random.nextBoolean()) { -// val = random.nextInt(10) - 5; // Random long between -5 and 4 -// doc.add(new SortedNumericDocValuesField(SNDV, val)); -// } -// if (random.nextBoolean()) { -// val = random.nextInt(20) - 10; // Random long between -10 and 9 -// doc.add(new SortedNumericDocValuesField(DV, val)); -// } -// if (random.nextBoolean()) { -// val = random.nextInt(50); // Random long between 0 and 49 -// doc.add(new SortedNumericDocValuesField(FIELD_NAME, val)); -// } -// iw.addDocument(doc); -// } -// -// if (randomBoolean()) { -// iw.forceMerge(1); -// } -// iw.close(); -// -// DirectoryReader ir = DirectoryReader.open(directory); -// initValuesSourceRegistry(); -// LeafReaderContext context = ir.leaves().get(0); -// -// SegmentReader reader = Lucene.segmentReader(context.reader()); -// IndexSearcher indexSearcher = newSearcher(reader, false, false); -// CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader(); -// List compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields(); -// -// CompositeIndexFieldInfo starTree = compositeIndexFields.get(0); -// -// SumAggregationBuilder sumAggregationBuilder = sum("_name").field(FIELD_NAME); -// MaxAggregationBuilder maxAggregationBuilder = max("_name").field(FIELD_NAME); -// MinAggregationBuilder minAggregationBuilder = min("_name").field(FIELD_NAME); -// ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME); -// AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME); -// -// // match-all query -// Query defaultQuery = new MatchAllDocsQuery(); -// StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, null); // no predicates -// testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue)); -// testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue)); -// testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue)); -// testCase(indexSearcher, defaultQuery, starTreeQuery, valueCountAggregationBuilder, verifyAggregation(InternalValueCount::getValue)); -// testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue)); -// // numeric-terms query -// for (int cases = 0; cases < 100; cases++) { -// Map queryMap; -// String queryField; -// long queryValue; -// if (randomBoolean()) { -// queryField = SNDV; -// queryValue = random.nextInt(10); -// } else { -// queryField = DV; -// queryValue = random.nextInt(20) - 15; -// } -// defaultQuery = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue); -// queryMap = Map.of(queryField, queryValue); -// starTreeQuery = new StarTreeQuery(starTree, queryMap); -// -// testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue)); -// testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue)); -// testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue)); -// testCase( -// indexSearcher, -// defaultQuery, -// starTreeQuery, -// valueCountAggregationBuilder, -// verifyAggregation(InternalValueCount::getValue) -// ); -// testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue)); -// } -// ir.close(); -// directory.close(); -// } - -// private void testCase( -// IndexSearcher searcher, -// Query defaultQuery, -// T builder, -// BiConsumer verify -// ) throws IOException { -//// OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, defaultQuery); -// V starTreeAggregation = searchAndReduceStarTree( -// createIndexSettings(), -// searcher, -// defaultQuery, -// builder, -// DEFAULT_MAX_BUCKETS, -// false, -// DEFAULT_MAPPED_FIELD -// ); -// V expectedAggregation = searchAndReduceStarTree( -// createIndexSettings(), -// searcher, -// defaultQuery, -// builder, -// DEFAULT_MAX_BUCKETS, -// false, -// DEFAULT_MAPPED_FIELD -// ); -// verify.accept(expectedAggregation, starTreeAggregation); -// } - - BiConsumer verifyAggregation(Function valueExtractor) { - return (expectedAggregation, actualAggregation) -> assertEquals( - valueExtractor.apply(expectedAggregation).doubleValue(), - valueExtractor.apply(actualAggregation).doubleValue(), - 0.0f - ); - } - - public void testStarTreeDocValues() throws IOException { Directory directory = newDirectory(); IndexWriterConfig conf = newIndexWriterConfig(null); @@ -286,11 +152,18 @@ public void testStarTreeDocValues() throws IOException { 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, valueCountAggregationBuilder, starTree, verifyAggregation(InternalValueCount::getValue)); -// testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue)); + 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, + valueCountAggregationBuilder, + starTree, + verifyAggregation(InternalValueCount::getValue) + ); + testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue)); // Numeric-terms query for (int cases = 0; cases < 1; cases++) { @@ -303,30 +176,35 @@ public void testStarTreeDocValues() throws IOException { queryField = DV; queryValue = random.nextInt(20) - 15; } - queryField = DV; - queryValue = 1; 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, -// valueCountAggregationBuilder, -// starTree, -// verifyAggregation(InternalValueCount::getValue) -// ); -// testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue)); + 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, + valueCountAggregationBuilder, + starTree, + verifyAggregation(InternalValueCount::getValue) + ); + testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue)); } ir.close(); directory.close(); } + BiConsumer verifyAggregation(Function valueExtractor) { + return (expectedAggregation, actualAggregation) -> assertEquals( + valueExtractor.apply(expectedAggregation).doubleValue(), + valueExtractor.apply(actualAggregation).doubleValue(), + 0.0f + ); + } private void testCase( IndexSearcher searcher,