From a81441bbf5cf1e43b120f1ba35ba4f37dd5a4617 Mon Sep 17 00:00:00 2001 From: Sandesh Kumar Date: Mon, 30 Sep 2024 21:56:20 -0700 Subject: [PATCH] max/min agg fix, spotless, test fixes Signed-off-by: Sandesh Kumar --- .../startree/utils/StarTreeQueryHelper.java | 16 ++++++++++------ .../SortedNumericStarTreeValuesIterator.java | 4 ++++ .../startree/MetricAggregatorTests.java | 6 +++--- .../search/aggregations/AggregatorTestCase.java | 7 +------ 4 files changed, 18 insertions(+), 15 deletions(-) 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 ccaaaa597ec0f..a87f6e5f94406 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 @@ -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/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java index 44f9545ce4f7f..7d65e9eb1da8e 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java +++ b/server/src/main/java/org/opensearch/index/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 index 4d7a8a9f9c644..06fcae5642e78 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 @@ -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 index 322bac6cafb2e..e1069bb0ea519 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 @@ -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,7 +94,6 @@ 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; @@ -103,7 +101,6 @@ 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);