refactor cache map to cache array

Signed-off-by: Sandesh Kumar <[email protected]>

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<String, Long> 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<LeafReaderContext, FixedBitSet> 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;
     }

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

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,

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<String, Long> queryMap;
+    private final Map<String, Long> 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<LeafReaderContext, FixedBitSet> starTreeValuesMap;
+    private FixedBitSet[] starTreeValues;
 
-    public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap, boolean cacheStarTreeValues) {
+    public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> 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<String, Long> getQueryMap() {
         return queryMap;
     }
 
-    public Map<LeafReaderContext, FixedBitSet> 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;
+        }
+    }
+
 }

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,33 +90,33 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
             CompositeMappedFieldType.CompositeFieldType.STAR_TREE
         );
         Map<String, Long> 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)
             .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), 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();
         }