From a18249d69d2872190ced7bf29a17546d5ea681e0 Mon Sep 17 00:00:00 2001 From: Sandesh Kumar Date: Mon, 30 Sep 2024 15:10:11 -0700 Subject: [PATCH] adding tests Signed-off-by: Sandesh Kumar --- .../startree/utils/StarTreeQueryHelper.java | 42 ++-- .../index/mapper/StarTreeMapper.java | 2 +- .../search/DefaultSearchContext.java | 28 --- .../aggregations/metrics/AvgAggregator.java | 144 ++--------- .../aggregations/metrics/SumAggregator.java | 1 - .../search/internal/SearchContext.java | 20 +- .../search/startree/StarTreeFilter.java | 1 - .../search/startree/StarTreeQueryContext.java | 39 +-- .../opensearch/search/SearchServiceTests.java | 159 ++++++------- .../startree/MetricAggregatorTests.java | 225 ++++++++++++++---- .../aggregations/AggregatorTestCase.java | 138 +++++++---- 11 files changed, 418 insertions(+), 381 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 179c236775a21..35b1f1997bcfe 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 @@ -12,7 +12,6 @@ import org.apache.lucene.index.SegmentReader; import org.apache.lucene.search.CollectionTerminatedException; import org.apache.lucene.util.FixedBitSet; -import org.apache.lucene.util.NumericUtils; import org.opensearch.common.lucene.Lucene; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.codec.composite.CompositeIndexReader; @@ -30,15 +29,14 @@ import org.opensearch.search.aggregations.AggregatorFactory; import org.opensearch.search.aggregations.LeafBucketCollector; import org.opensearch.search.aggregations.LeafBucketCollectorBase; -import org.opensearch.search.aggregations.metrics.CompensatedSum; import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory; import org.opensearch.search.aggregations.support.ValuesSource; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.internal.SearchContext; +import org.opensearch.search.startree.StarTreeFilter; import org.opensearch.search.startree.StarTreeQueryContext; import java.io.IOException; -import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -95,14 +93,16 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context return null; } - boolean needCaching = context.aggregations().factories().getFactories().length > 1; -// List metricInfos = new ArrayList<>(); + for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) { - MetricStat metricStat = validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory); + MetricStat metricStat = validateStarTreeMetricSupport(compositeMappedFieldType, aggregatorFactory); if (metricStat == null) { return null; } -// metricInfos.add(new ) + } + + if (context.aggregations().factories().getFactories().length > 1) { + context.initializeStarTreeValuesMap(); } return starTreeQueryContext; @@ -134,7 +134,7 @@ private static StarTreeQueryContext toStarTreeQueryContext( /** * Parse query body to star-tree predicates - * @param queryBuilder + * @param queryBuilder to match supported query shape * @return predicates to match */ private static Map getStarTreePredicates(QueryBuilder queryBuilder, List supportedDimensions) { @@ -151,11 +151,10 @@ private static Map getStarTreePredicates(QueryBuilder queryBuilder return predicateMap; } - private static MetricStat validateStarTreeMetricSuport( + private static MetricStat validateStarTreeMetricSupport( CompositeDataCubeFieldType compositeIndexFieldInfo, AggregatorFactory aggregatorFactory ) { -// List metricStats = new ArrayList<>(); if (aggregatorFactory instanceof MetricAggregatorFactory && aggregatorFactory.getSubFactories().getFactories().length == 0) { String field; Map> supportedMetrics = compositeIndexFieldInfo.getMetrics() @@ -197,6 +196,7 @@ public static LeafBucketCollector getStarTreeLeafCollector( Runnable finalConsumer ) throws IOException { StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); + assert starTreeValues != null; String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric); @@ -205,15 +205,10 @@ public static LeafBucketCollector getStarTreeLeafCollector( metricName ); // Obtain a FixedBitSet of matched document IDs - FixedBitSet matchedDocIds = context.getStarTreeFilteredValues(ctx, starTreeValues); // Assuming this method gives a FixedBitSet - - // Safety check: make sure the FixedBitSet is non-null and valid - if (matchedDocIds == null) { - throw new IllegalStateException("FixedBitSet is null"); - } + FixedBitSet matchedDocIds = getStarTreeFilteredValues(context, ctx, starTreeValues); // Assuming this method gives a FixedBitSet + assert matchedDocIds != null; 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) { @@ -230,6 +225,7 @@ public static LeafBucketCollector getStarTreeLeafCollector( } } + // Call the final consumer after processing all entries finalConsumer.run(); @@ -241,4 +237,16 @@ public void collect(int doc, long bucket) { } }; } + + 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); + } + + StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap()); + FixedBitSet result = filter.getStarTreeResult(); + + context.getStarTreeValuesMap().put(ctx, result); + return result; + } } diff --git a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java index 52dab17e0b0bb..55456c9cc07e8 100644 --- a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java @@ -463,7 +463,7 @@ protected void parseCreateField(ParseContext context) { * @opensearch.experimental */ @ExperimentalApi - public static final class StarTreeFieldType extends CompositeDataCubeFieldType { + public static class StarTreeFieldType extends CompositeDataCubeFieldType { private final StarTreeFieldConfiguration starTreeConfig; diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java index cfe2ad7681a58..feaafd9bbbc6c 100644 --- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java +++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java @@ -43,7 +43,6 @@ import org.apache.lucene.search.FieldDoc; import org.apache.lucene.search.MatchNoDocsQuery; import org.apache.lucene.search.Query; -import org.apache.lucene.util.BitSet; import org.apache.lucene.util.FixedBitSet; import org.opensearch.Version; import org.opensearch.action.search.SearchShardTask; @@ -59,10 +58,7 @@ import org.opensearch.index.IndexService; import org.opensearch.index.IndexSettings; import org.opensearch.index.cache.bitset.BitsetFilterCache; -import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; -import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper; -import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import org.opensearch.index.engine.Engine; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; @@ -185,7 +181,6 @@ final class DefaultSearchContext extends SearchContext { private SliceBuilder sliceBuilder; private SearchShardTask task; private final Version minNodeVersion; - private StarTreeQueryContext starTreeQueryContext; /** * The original query as sent by the user without the types and aliases @@ -1158,27 +1153,4 @@ public boolean evaluateKeywordIndexOrDocValuesEnabled() { } return false; } - - @Override - public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) { - this.starTreeQueryContext = starTreeQueryContext; - return this; - } - - @Override - public StarTreeQueryContext getStarTreeQueryContext() { - return this.starTreeQueryContext; - } - - @Override - public FixedBitSet getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { - if (this.starTreeValuesMap.containsKey(ctx)) { - return starTreeValuesMap.get(ctx); - } - StarTreeFilter filter = new StarTreeFilter(starTreeValues, this.getStarTreeQueryContext().getQueryMap()); - FixedBitSet result = filter.getStarTreeResult(); - - starTreeValuesMap.put(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 b375239e94e5b..4a8fe891ae066 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 @@ -56,12 +56,11 @@ import org.opensearch.search.aggregations.support.ValuesSource; import org.opensearch.search.aggregations.support.ValuesSourceConfig; import org.opensearch.search.internal.SearchContext; -import org.opensearch.search.startree.StarTreeFilter; import java.io.IOException; import java.util.Map; -import java.util.function.Consumer; +import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeFilteredValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree; /** @@ -109,13 +108,12 @@ 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); } private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException { - final BigArrays bigArrays = context.bigArrays(); final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx); final CompensatedSum kahanSummation = new CompensatedSum(0, 0); @@ -149,105 +147,11 @@ public void collect(int doc, long bucket) throws IOException { }; } - public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) - throws IOException { - final CompensatedSum kahanSummation = new CompensatedSum(0, 0); - return StarTreeQueryHelper.getStarTreeLeafCollector( - context, - valuesSource, - ctx, - sub, - starTree, - MetricStat.SUM.getTypeName(), - value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)), - () -> sums.set(0, kahanSummation.value()) - ); - } - -// private LeafBucketCollector getStarTreeLeafCollector1(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) -// throws IOException { -// final BigArrays bigArrays = context.bigArrays(); -// final CompensatedSum kahanSummation = new CompensatedSum(0, 0); -// -// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); -// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); -// String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( -// starTree.getField(), -// fieldName, -// MetricStat.SUM.getTypeName() -// ); -// assert starTreeValues != null; -// SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName); -// -// String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( -// starTree.getField(), -// fieldName, -// MetricStat.VALUE_COUNT.getTypeName() -// ); -// SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName); -// -// return new LeafBucketCollectorBase(sub, values) { -// @Override -// public void collect(int doc, long bucket) throws IOException { -// counts = bigArrays.grow(counts, bucket + 1); -// sums = bigArrays.grow(sums, bucket + 1); -// compensations = bigArrays.grow(compensations, bucket + 1); -// -// if (values.advanceExact(doc) && countValues.advanceExact(doc)) { -// final long valueCount = values.docValueCount(); -// counts.increment(bucket, countValues.nextValue()); -// // Compute the sum of double values with Kahan summation algorithm which is more -// // accurate than naive summation. -// double sum = sums.get(bucket); -// double compensation = compensations.get(bucket); -// -// kahanSummation.reset(sum, compensation); -// -// for (int i = 0; i < valueCount; i++) { -// double value = NumericUtils.sortableLongToDouble(values.nextValue()); -// kahanSummation.add(value); -// } -// -// sums.set(bucket, kahanSummation.value()); -// compensations.set(bucket, kahanSummation.delta()); -// } -// } -// }; -// } - - - public LeafBucketCollector getStarTreeCollector2(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) throws IOException { - final CompensatedSum kahanSummation = new CompensatedSum(0, 0); - StarTreeQueryHelper.getStarTreeLeafCollector( - context, - valuesSource, - ctx, - sub, - starTree, - MetricStat.SUM.getTypeName(), - value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)), - () -> sums.set(0, kahanSummation.value()) - ); - StarTreeQueryHelper.getStarTreeLeafCollector( - context, - valuesSource, - ctx, - sub, - starTree, - MetricStat.VALUE_COUNT.getTypeName(), - value -> counts.increment(0, value), - () -> {} - ); - return LeafBucketCollector.NO_OP_COLLECTOR; - } - - public LeafBucketCollector getStarTreeLeafCollector( - LeafReaderContext ctx, - LeafBucketCollector sub, - CompositeIndexFieldInfo starTree - ) throws IOException { StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree); + assert starTreeValues != null; + String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( starTree.getField(), @@ -260,35 +164,31 @@ public LeafBucketCollector getStarTreeLeafCollector( MetricStat.VALUE_COUNT.getTypeName() ); - assert starTreeValues != null; final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), 0); - SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator( - sumMetricName - ); - SortedNumericStarTreeValuesIterator countValueIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator( - countMetricName - ); - FixedBitSet matchedDocIds = context.getStarTreeFilteredValues(ctx, starTreeValues); - - - // Safety check: make sure the FixedBitSet is non-null and valid - if (matchedDocIds == null) { - throw new IllegalStateException("FixedBitSet is null"); - } + SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues + .getMetricValuesIterator(sumMetricName); + SortedNumericStarTreeValuesIterator countValueIterator = (SortedNumericStarTreeValuesIterator) starTreeValues + .getMetricValuesIterator(countMetricName); + FixedBitSet matchedDocIds = getStarTreeFilteredValues(context, ctx, starTreeValues); + assert matchedDocIds != null; 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) { + // Advance to the bit (entryId) in the valuesIterator + if (sumValuesIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES + || countValueIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES) { + continue; // Skip if no more entries + } - // 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 (sumValuesIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES && - countValueIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES) { - int count = sumValuesIterator.valuesCount(); - for (int i = 0; i < count; i++) { + // Iterate over the values for the current entryId + for (int i = 0; i < sumValuesIterator.valuesCount(); i++) { kahanSummation.add(NumericUtils.sortableLongToDouble(sumValuesIterator.nextValue())); counts.increment(0, countValueIterator.nextValue()); // Apply the consumer operation (e.g., max, sum) } + } } 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 c3c7ad299430c..8fb29675eac1c 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 @@ -100,7 +100,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc if (supportedStarTree != null) { return getStarTreeCollector(ctx, sub, supportedStarTree); } - System.out.println("nopes nopes"); return getDefaultLeafCollector(ctx, sub); } 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 ae78e8ca76845..13340fa10f849 100644 --- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java +++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java @@ -46,9 +46,7 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.BigArrays; import org.opensearch.index.cache.bitset.BitsetFilterCache; -import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; -import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.ObjectMapper; @@ -81,6 +79,7 @@ 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; @@ -133,6 +132,7 @@ public List toInternalAggregations(Collection co private InnerHitsContext innerHitsContext; protected volatile Map starTreeValuesMap; private volatile boolean searchTimedOut; + private StarTreeQueryContext starTreeQueryContext; protected SearchContext() {} @@ -538,15 +538,21 @@ public boolean keywordIndexOrDocValuesEnabled() { return false; } - public StarTreeQueryContext getStarTreeQueryContext() { - return null; - } public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) { + this.starTreeQueryContext = starTreeQueryContext; return this; } - public FixedBitSet getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { - return null; + public StarTreeQueryContext getStarTreeQueryContext() { + return this.starTreeQueryContext; + } + + public void initializeStarTreeValuesMap() { + this.starTreeValuesMap = new HashMap<>(); + } + + public Map getStarTreeValuesMap() { + return starTreeValuesMap; } } diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java index d56ce506ebddb..dcc5bff8a7887 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java @@ -10,7 +10,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.apache.lucene.util.BitSet; import org.apache.lucene.util.DocIdSetBuilder; import org.apache.lucene.util.FixedBitSet; import org.opensearch.index.compositeindex.datacube.Dimension; 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 6d184938ccfb5..74f05c0b3e7fa 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java @@ -8,15 +8,10 @@ package org.opensearch.search.startree; -import org.apache.lucene.index.LeafReaderContext; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; -import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper; -import java.util.ArrayList; -import java.util.List; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; /** * Query class for querying star tree data structure. @@ -38,17 +33,17 @@ public class StarTreeQueryContext { */ private final Map queryMap; -// /** -// * Cache for leaf results -// * This is used to cache the results for each leaf reader context -// * to avoid reading the data from the leaf reader context multiple times -// */ -// private volatile Map> leafResultsCache; + // /** + // * Cache for leaf results + // * This is used to cache the results for each leaf reader context + // * to avoid reading the data from the leaf reader context multiple times + // */ + // private volatile Map> leafResultsCache; -// /** -// * List of metrics to be computed & cached -// */ -// private List metrics; + // /** + // * List of metrics to be computed & cached + // */ + // private List metrics; public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map queryMap) { this.starTree = starTree; @@ -62,18 +57,4 @@ public CompositeIndexFieldInfo getStarTree() { public Map getQueryMap() { return queryMap; } - -// public void initializeLeafResultsCache() { -// this.leafResultsCache = new ConcurrentHashMap<>(); -// this.metrics = new ArrayList<>(); -// } -// -// public Map> getLeafResultsCache() { -// return leafResultsCache; -// } -// -// public void addMetric(StarTreeQueryHelper.MetricInfo metric) { -// metrics.add(metric); -// } - } diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java index 8b8e038a41360..d49111609f95e 100644 --- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java +++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java @@ -121,7 +121,6 @@ import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.MinAndMax; import org.opensearch.search.sort.SortOrder; -import org.opensearch.search.startree.StarTreeQuery; import org.opensearch.search.suggest.SuggestBuilder; import org.opensearch.test.OpenSearchSingleNodeTestCase; import org.opensearch.threadpool.ThreadPool; @@ -2269,75 +2268,75 @@ 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(); +// 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); +// } private void setStarTreeIndexSetting(String value) throws IOException { client().admin() @@ -2347,14 +2346,14 @@ 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 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(); +// } +// } } 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 41eba6b548b45..eba7036550987 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,9 +32,13 @@ 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; import org.opensearch.search.aggregations.AggregatorTestCase; import org.opensearch.search.aggregations.InternalAggregation; @@ -48,15 +52,15 @@ import org.opensearch.search.aggregations.metrics.MinAggregationBuilder; import org.opensearch.search.aggregations.metrics.SumAggregationBuilder; import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder; -import org.opensearch.search.startree.OriginalOrStarTreeQuery; -import org.opensearch.search.startree.StarTreeQuery; import org.junit.After; import org.junit.Before; 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; @@ -94,6 +98,136 @@ 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); @@ -107,6 +241,7 @@ public void testStarTreeDocValues() throws IOException { final String DV = "dv"; int val; + List docs = new ArrayList<>(); // Index 100 random documents for (int i = 0; i < totalDocs; i++) { Document doc = new Document(); @@ -123,6 +258,7 @@ public void testStarTreeDocValues() throws IOException { doc.add(new SortedNumericDocValuesField(FIELD_NAME, val)); } iw.addDocument(doc); + docs.add(doc); } if (randomBoolean()) { @@ -137,8 +273,8 @@ public void testStarTreeDocValues() throws IOException { SegmentReader reader = Lucene.segmentReader(context.reader()); IndexSearcher indexSearcher = newSearcher(reader, false, false); CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader(); - List compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields(); + List compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields(); CompositeIndexFieldInfo starTree = compositeIndexFields.get(0); SumAggregationBuilder sumAggregationBuilder = sum("_name").field(FIELD_NAME); @@ -147,17 +283,17 @@ public void testStarTreeDocValues() throws IOException { ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME); AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME); + Query query = new MatchAllDocsQuery(); // 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; + 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)); + + // Numeric-terms query + for (int cases = 0; cases < 1; cases++) { String queryField; long queryValue; if (randomBoolean()) { @@ -167,39 +303,46 @@ public void testStarTreeDocValues() throws IOException { 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)); + 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)); } + ir.close(); directory.close(); } + private void testCase( IndexSearcher searcher, - Query defaultQuery, - StarTreeQuery starTreeQuery, - T builder, + Query query, + QueryBuilder queryBuilder, + T aggBuilder, + CompositeIndexFieldInfo starTree, BiConsumer verify ) throws IOException { - OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, defaultQuery); V starTreeAggregation = searchAndReduceStarTree( createIndexSettings(), searcher, - originalOrStarTreeQuery, - builder, + query, + queryBuilder, + aggBuilder, + starTree, DEFAULT_MAX_BUCKETS, false, DEFAULT_MAPPED_FIELD @@ -207,20 +350,14 @@ private void testC V expectedAggregation = searchAndReduceStarTree( createIndexSettings(), searcher, - defaultQuery, - builder, + query, + queryBuilder, + aggBuilder, + null, 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 - ); - } } 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 63288602abb7e..322bac6cafb2e 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,6 +60,7 @@ 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; @@ -91,32 +92,19 @@ import org.opensearch.index.cache.bitset.BitsetFilterCache; import org.opensearch.index.cache.bitset.BitsetFilterCache.Listener; import org.opensearch.index.cache.query.DisabledQueryCache; +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.BinaryFieldMapper; -import org.opensearch.index.mapper.CompletionFieldMapper; -import org.opensearch.index.mapper.ConstantKeywordFieldMapper; -import org.opensearch.index.mapper.ContentPath; -import org.opensearch.index.mapper.DateFieldMapper; -import org.opensearch.index.mapper.DerivedFieldMapper; -import org.opensearch.index.mapper.FieldAliasMapper; -import org.opensearch.index.mapper.FieldMapper; -import org.opensearch.index.mapper.GeoPointFieldMapper; -import org.opensearch.index.mapper.GeoShapeFieldMapper; -import org.opensearch.index.mapper.KeywordFieldMapper; -import org.opensearch.index.mapper.MappedFieldType; -import org.opensearch.index.mapper.Mapper; +import org.opensearch.index.mapper.*; import org.opensearch.index.mapper.Mapper.BuilderContext; -import org.opensearch.index.mapper.MapperService; -import org.opensearch.index.mapper.MatchOnlyTextFieldMapper; -import org.opensearch.index.mapper.NumberFieldMapper; -import org.opensearch.index.mapper.ObjectMapper; import org.opensearch.index.mapper.ObjectMapper.Nested; -import org.opensearch.index.mapper.RangeFieldMapper; -import org.opensearch.index.mapper.RangeType; -import org.opensearch.index.mapper.StarTreeMapper; -import org.opensearch.index.mapper.TextFieldMapper; +import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryShardContext; import org.opensearch.index.shard.IndexShard; import org.opensearch.index.shard.SearchOperationListener; @@ -135,12 +123,15 @@ import org.opensearch.search.aggregations.support.CoreValuesSourceType; import org.opensearch.search.aggregations.support.ValuesSourceRegistry; import org.opensearch.search.aggregations.support.ValuesSourceType; +import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.fetch.FetchPhase; import org.opensearch.search.fetch.subphase.FetchDocValuesPhase; import org.opensearch.search.fetch.subphase.FetchSourcePhase; 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; import org.junit.After; @@ -152,9 +143,12 @@ 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; import java.util.Objects; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.function.Function; @@ -331,6 +325,28 @@ protected A createAggregator(AggregationBuilder aggregati return aggregator; } + protected CountingAggregator createCountingAggregator( + Query query, + QueryBuilder queryBuilder, + AggregationBuilder aggregationBuilder, + IndexSearcher indexSearcher, + IndexSettings indexSettings, + CompositeIndexFieldInfo starTree, + MultiBucketConsumer bucketConsumer, + MappedFieldType... fieldTypes + ) throws IOException { + SearchContext searchContext; + if (starTree != null) { + searchContext = createSearchContextWithStarTreeContext(indexSearcher, indexSettings, query, queryBuilder, starTree, bucketConsumer, fieldTypes); + } else { + searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, fieldTypes); + } + return new CountingAggregator( + new AtomicInteger(), + createAggregator(aggregationBuilder, searchContext) + ); + } + /** * Create a {@linkplain SearchContext} for testing an {@link Aggregator}. */ @@ -344,6 +360,45 @@ protected SearchContext createSearchContext( return createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes); } + protected SearchContext createSearchContextWithStarTreeContext( + IndexSearcher indexSearcher, + IndexSettings indexSettings, + Query query, + QueryBuilder queryBuilder, + CompositeIndexFieldInfo starTree, + MultiBucketConsumer bucketConsumer, + MappedFieldType... fieldTypes + ) throws IOException { + SearchContext searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes); + +// Mock SearchContextAggregations + SearchContextAggregations searchContextAggregations = mock(SearchContextAggregations.class); + AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class); + when(searchContext.aggregations()).thenReturn(searchContextAggregations); + when(searchContextAggregations.factories()).thenReturn(aggregatorFactories); + when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[]{}); + + StarTreeMapper.StarTreeFieldType compositeMappedFieldType = mock(StarTreeMapper.StarTreeFieldType.class); + when(compositeMappedFieldType.name()).thenReturn(starTree.getField()); + when(compositeMappedFieldType.getCompositeIndexType()).thenReturn(starTree.getType()); + Set compositeFieldTypes = Set.of(compositeMappedFieldType); + + List dimensions = new LinkedList<>(); + dimensions.add(new NumericDimension("sndv")); + dimensions.add(new NumericDimension("dv")); + when((compositeMappedFieldType).getDimensions()).thenReturn(dimensions); + MapperService mapperService = mock(MapperService.class); + when(mapperService.getCompositeFieldTypes()).thenReturn(compositeFieldTypes); + when(searchContext.mapperService()).thenReturn(mapperService); + + SearchSourceBuilder sb = new SearchSourceBuilder().query(queryBuilder); + StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(searchContext, sb); + + when(searchContext.getStarTreeQueryContext()).thenReturn(starTreeQueryContext); + return searchContext; + + } + protected SearchContext createSearchContext( IndexSearcher indexSearcher, IndexSettings indexSettings, @@ -655,7 +710,9 @@ protected A searchAndReduc IndexSettings indexSettings, IndexSearcher searcher, Query query, + QueryBuilder queryBuilder, AggregationBuilder builder, + CompositeIndexFieldInfo compositeIndexFieldInfo, int maxBucket, boolean hasNested, MappedFieldType... fieldTypes @@ -672,35 +729,14 @@ protected A searchAndReduc maxBucket, new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST) ); - C root = createAggregator(query, builder, searcher, bucketConsumer, fieldTypes); - - if (randomBoolean() && searcher.getIndexReader().leaves().size() > 0) { - assertTrue(ctx instanceof LeafReaderContext); - final LeafReaderContext compCTX = (LeafReaderContext) ctx; - final int size = compCTX.leaves().size(); - final ShardSearcher[] subSearchers = new ShardSearcher[size]; - for (int searcherIDX = 0; searcherIDX < subSearchers.length; searcherIDX++) { - final LeafReaderContext leave = compCTX.leaves().get(searcherIDX); - subSearchers[searcherIDX] = new ShardSearcher(leave, compCTX); - } - for (ShardSearcher subSearcher : subSearchers) { - MultiBucketConsumer shardBucketConsumer = new MultiBucketConsumer( - maxBucket, - new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST) - ); - C a = createAggregator(query, builder, subSearcher, indexSettings, shardBucketConsumer, fieldTypes); - a.preCollection(); - Weight weight = subSearcher.createWeight(query, ScoreMode.COMPLETE, 1f); - - subSearcher.search(weight, a); - a.postCollection(); - aggs.add(a.buildTopLevel()); - } - } else { - root.preCollection(); - searcher.search(query, root); - root.postCollection(); - aggs.add(root.buildTopLevel()); + CountingAggregator countingAggregator = createCountingAggregator(query, queryBuilder, builder, searcher, indexSettings, compositeIndexFieldInfo, bucketConsumer, fieldTypes); + + countingAggregator.preCollection(); + searcher.search(query, countingAggregator); + countingAggregator.postCollection(); + aggs.add(countingAggregator.buildTopLevel()); + if (compositeIndexFieldInfo != null) { + assertEquals(0, countingAggregator.collectCounter.get()); } MultiBucketConsumer reduceBucketConsumer = new MultiBucketConsumer( @@ -708,7 +744,7 @@ protected A searchAndReduc new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST) ); InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forFinalReduction( - root.context().bigArrays(), + countingAggregator.context().bigArrays(), getMockScriptService(), reduceBucketConsumer, pipelines