max/min agg fix, spotless, test fixes

Signed-off-by: Sandesh Kumar <[email protected]>
sandeshkr419 · Oct 1, 2024 · a81441b · a81441b
1 parent 4c35a89
commit a81441b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 15 deletions.
diff --git a/...java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/...java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -20,7 +20,6 @@
 import org.opensearch.index.compositeindex.datacube.MetricStat;
 import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
 import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
-import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
 import org.opensearch.index.mapper.CompositeDataCubeFieldType;
 import org.opensearch.index.mapper.StarTreeMapper;
 import org.opensearch.index.query.MatchAllQueryBuilder;
@@ -66,7 +65,7 @@ public static boolean isStarTreeSupported(SearchContext context) {
 
     /**
      * Gets StarTreeQueryContext from the search context and source builder.
-     * Returns null if the query & aggregation cannot be supported.
+     * Returns null if the query and aggregation cannot be supported.
      */
     public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException {
         // Current implementation assumes only single star-tree is supported
@@ -102,7 +101,7 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
     }
 
     /**
-     * Uses query builder & composite index info to form star-tree query context
+     * Uses query builder and composite index info to form star-tree query context
      */
     private static StarTreeQueryContext toStarTreeQueryContext(
         CompositeIndexFieldInfo compositeIndexFieldInfo,
@@ -211,8 +210,8 @@ public static LeafBucketCollector getStarTreeLeafCollector(
         if (numBits > 0) {
             // Iterate over the FixedBitSet
             for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? matchedDocIds.nextSetBit(bit + 1) : -1) {
-                // Advance to the bit (entryId) in the valuesIterator
-                if (valuesIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES) {
+                // Advance to the entryId in the valuesIterator
+                if (!valuesIterator.advanceExact(bit)) {
                     continue;  // Skip if no more entries
                 }
 
@@ -238,6 +237,8 @@ public void collect(int doc, long bucket) {
 
     /**
      * Get the filtered values for the star-tree query
+     * Cache the results in case of multiple aggregations (if cache is initialized)
+     * @return FixedBitSet of matched document IDs
      */
     public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
         throws IOException {
@@ -248,7 +249,10 @@ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafR
         StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
         FixedBitSet result = filter.getStarTreeResult();
 
-        context.getStarTreeValuesMap().put(ctx, result);
+        if (context.getStarTreeValuesMap() != null) {
+            context.getStarTreeValuesMap().put(ctx, result);
+        }
         return result;
+
     }
 }
diff --git a/.../compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java b/.../compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
@@ -33,4 +33,8 @@ public long nextValue() throws IOException {
     public int valuesCount() throws IOException {
         return ((SortedNumericDocValues) docIdSetIterator).docValueCount();
     }
+
+    public boolean advanceExact(int target) throws IOException {
+        return ((SortedNumericDocValues) docIdSetIterator).advanceExact(target);
+    }
 }
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
@@ -166,7 +166,7 @@ public void testStarTreeDocValues() throws IOException {
         testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
 
         // Numeric-terms query
-        for (int cases = 0; cases < 1; cases++) {
+        for (int cases = 0; cases < 100; cases++) {
             String queryField;
             long queryValue;
             if (randomBoolean()) {
@@ -181,8 +181,8 @@ public void testStarTreeDocValues() throws IOException {
             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, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
+            testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
             testCase(
                 indexSearcher,
                 query,

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
@@ -60,7 +60,6 @@
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.apache.lucene.tests.search.AssertingIndexSearcher;
 import org.apache.lucene.util.BytesRef;
-import org.apache.lucene.util.FixedBitSet;
 import org.apache.lucene.util.NumericUtils;
 import org.opensearch.Version;
 import org.opensearch.cluster.metadata.IndexMetadata;
@@ -95,15 +94,13 @@
 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.index.StarTreeValues;
 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.Mapper.BuilderContext;
 import org.opensearch.index.mapper.ObjectMapper.Nested;
-import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.shard.IndexShard;
@@ -130,7 +127,6 @@
 import org.opensearch.search.internal.ContextIndexSearcher;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.lookup.SearchLookup;
-import org.opensearch.search.startree.StarTreeFilter;
 import org.opensearch.search.startree.StarTreeQueryContext;
 import org.opensearch.test.InternalAggregationTestCase;
 import org.opensearch.test.OpenSearchTestCase;
@@ -143,7 +139,6 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -371,7 +366,7 @@ protected SearchContext createSearchContextWithStarTreeContext(
     ) throws IOException {
         SearchContext searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes);
 
-// Mock SearchContextAggregations
+        // Mock SearchContextAggregations
         SearchContextAggregations searchContextAggregations = mock(SearchContextAggregations.class);
         AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class);
         when(searchContext.aggregations()).thenReturn(searchContextAggregations);