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 53c5e2bfd13e4..3fe982fc0091a 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 @@ -11,6 +11,7 @@ import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.SegmentReader; import org.apache.lucene.search.CollectionTerminatedException; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.util.FixedBitSet; import org.opensearch.common.lucene.Lucene; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; @@ -78,19 +79,21 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context } } + // need to cache star tree values only for multiple aggregations boolean cacheStarTreeValues = context.aggregations().factories().getFactories().length > 1; + int cacheSize = cacheStarTreeValues ? context.indexShard().segments(false).size() : -1; - return StarTreeQueryHelper.toStarTreeQueryContext(starTree, compositeMappedFieldType, source.query(), cacheStarTreeValues); + return StarTreeQueryHelper.tryCreateStarTreeQueryContext(starTree, compositeMappedFieldType, source.query(), cacheSize); } /** * Uses query builder and composite index info to form star-tree query context */ - private static StarTreeQueryContext toStarTreeQueryContext( + private static StarTreeQueryContext tryCreateStarTreeQueryContext( CompositeIndexFieldInfo compositeIndexFieldInfo, CompositeDataCubeFieldType compositeFieldType, QueryBuilder queryBuilder, - boolean cacheStarTreeValues + int cacheStarTreeValuesSize ) { Map queryMap; if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) { @@ -107,7 +110,7 @@ private static StarTreeQueryContext toStarTreeQueryContext( } else { return null; } - return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap, cacheStarTreeValues); + return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap, cacheStarTreeValuesSize); } /** @@ -193,7 +196,9 @@ public static LeafBucketCollector getStarTreeLeafCollector( int numBits = filteredValues.length(); // Get the number of the filtered values (matching docs) if (numBits > 0) { // Iterate over the filtered values - for (int bit = filteredValues.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? filteredValues.nextSetBit(bit + 1) : -1) { + for (int bit = filteredValues.nextSetBit(0); bit != DocIdSetIterator.NO_MORE_DOCS; bit = (bit + 1 < numBits) + ? filteredValues.nextSetBit(bit + 1) + : DocIdSetIterator.NO_MORE_DOCS) { // Advance to the entryId in the valuesIterator if (valuesIterator.advanceExact(bit) == false) { continue; // Skip if no more entries @@ -226,16 +231,11 @@ public void collect(int doc, long bucket) { */ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { - Map valueCache = context.getStarTreeQueryContext().getStarTreeValuesMap(); - if (valueCache != null && valueCache.containsKey(ctx)) { - return valueCache.get(ctx); - } - - StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap()); - FixedBitSet result = filter.getStarTreeResult(); - - if (valueCache != null) { - valueCache.put(ctx, result); + FixedBitSet result = context.getStarTreeQueryContext().getStarTreeValues(ctx); + if (result == null) { + StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap()); + result = filter.getStarTreeResult(); + context.getStarTreeQueryContext().setStarTreeValues(ctx, result); } return result; } 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 4299ebc475058..0a3a514701e15 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 @@ -33,6 +33,7 @@ import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.CollectionTerminatedException; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.util.FixedBitSet; import org.apache.lucene.util.NumericUtils; @@ -174,7 +175,9 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet if (numBits > 0) { // Iterate over the FixedBitSet - for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) { + for (int bit = matchedDocIds.nextSetBit(0); bit != DocIdSetIterator.NO_MORE_DOCS; bit = bit + 1 < numBits + ? matchedDocIds.nextSetBit(bit + 1) + : DocIdSetIterator.NO_MORE_DOCS) { // Advance to the bit (entryId) in the valuesIterator if ((sumValuesIterator.advanceExact(bit) && countValueIterator.advanceExact(bit)) == false) { continue; // Skip if no more entries diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java index 8fb29675eac1c..3d237a94c5699 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java @@ -62,13 +62,13 @@ */ public class SumAggregator extends NumericMetricsAggregator.SingleValue { - protected final ValuesSource.Numeric valuesSource; - protected final DocValueFormat format; + private final ValuesSource.Numeric valuesSource; + private final DocValueFormat format; - protected DoubleArray sums; - protected DoubleArray compensations; + private DoubleArray sums; + private DoubleArray compensations; - public SumAggregator( + SumAggregator( String name, ValuesSourceConfig valuesSourceConfig, SearchContext context, 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 39fb3d3a68dda..165de139ec0d4 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java @@ -14,7 +14,6 @@ import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; /** * Query class for querying star tree data structure. @@ -34,20 +33,20 @@ public class StarTreeQueryContext { * Map of field name to a value to be queried for that field * This is used to filter the data based on the query */ - private volatile Map queryMap; + private final Map queryMap; /** * Cache for leaf results * 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 */ - protected volatile Map starTreeValuesMap; + private FixedBitSet[] starTreeValues; - public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map queryMap, boolean cacheStarTreeValues) { + public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map queryMap, int cacheStarTreeValuesSize) { this.starTree = starTree; this.queryMap = queryMap; - if (cacheStarTreeValues) { - starTreeValuesMap = new ConcurrentHashMap<>(); + if (cacheStarTreeValuesSize > -1) { + starTreeValues = new FixedBitSet[cacheStarTreeValuesSize]; } } @@ -59,7 +58,21 @@ public Map getQueryMap() { return queryMap; } - public Map getStarTreeValuesMap() { - return starTreeValuesMap; + public FixedBitSet[] getStarTreeValues() { + return starTreeValues; } + + public FixedBitSet getStarTreeValues(LeafReaderContext ctx) { + if (starTreeValues != null) { + return starTreeValues[ctx.ord]; + } + return null; + } + + public void setStarTreeValues(LeafReaderContext ctx, FixedBitSet values) { + if (starTreeValues != null) { + starTreeValues[ctx.ord] = values; + } + } + } diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java index b2953f5482c9e..88d5265dfda78 100644 --- a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java +++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java @@ -75,11 +75,11 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException { // Case 1: No query or aggregations, should not use star tree SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); - assertStarTreeContext(request, sourceBuilder, null, false); + assertStarTreeContext(request, sourceBuilder, null, -1); // Case 2: MatchAllQuery present but no aggregations, should not use star tree sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder()); - assertStarTreeContext(request, sourceBuilder, null, false); + assertStarTreeContext(request, sourceBuilder, null, -1); // Case 3: MatchAllQuery and aggregations present, should use star tree sourceBuilder = new SearchSourceBuilder().size(0) @@ -90,21 +90,21 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException { CompositeMappedFieldType.CompositeFieldType.STAR_TREE ); Map expectedQueryMap = null; - assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, false), false); + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, -1), -1); // Case 4: MatchAllQuery and aggregations present, but postFilter specified, should not use star tree sourceBuilder = new SearchSourceBuilder().size(0) .query(new MatchAllQueryBuilder()) .aggregation(AggregationBuilders.max("test").field("field")) .postFilter(new MatchAllQueryBuilder()); - assertStarTreeContext(request, sourceBuilder, null, false); + assertStarTreeContext(request, sourceBuilder, null, -1); // 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), false); + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, -1), -1); // Case 6: TermQuery and multiple aggregations present, should use star tree & initialize cache sourceBuilder = new SearchSourceBuilder().size(0) @@ -112,11 +112,11 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException { .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), true); + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, 0), 0); // 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), false); + assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, null, -1), -1); setStarTreeIndexSetting(null); } @@ -133,7 +133,7 @@ private void assertStarTreeContext( ShardSearchRequest request, SearchSourceBuilder sourceBuilder, StarTreeQueryContext expectedContext, - boolean expectedCacheUsage + int expectedCacheUsage ) throws IOException { request.source(sourceBuilder); SearchService searchService = getInstanceFromNode(SearchService.class); @@ -148,7 +148,11 @@ private void assertStarTreeContext( assertEquals(expectedContext.getStarTree().getType(), actualContext.getStarTree().getType()); assertEquals(expectedContext.getStarTree().getField(), actualContext.getStarTree().getField()); assertEquals(expectedContext.getQueryMap(), actualContext.getQueryMap()); - assertThat(context.getStarTreeQueryContext().getStarTreeValuesMap(), expectedCacheUsage ? notNullValue() : nullValue()); + if (expectedCacheUsage > -1) { + assertEquals(expectedCacheUsage, actualContext.getStarTreeValues().length); + } else { + assertNull(actualContext.getStarTreeValues()); + } } searchService.doStop(); }