Skip to content

Commit

Permalink
spotless fix, test refactoring
Browse files Browse the repository at this point in the history 
Signed-off-by: Sandesh Kumar <[email protected]>
  • Loading branch information
@sandeshkr419
sandeshkr419 committed Oct 1, 2024
1 parent de3a242 commit 2ac5915
Show file tree
Hide file tree
Showing 2 changed files with 141 additions and 24 deletions.
92 changes: 84 additions & 8 deletions server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
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.compositeindex.datacube.Dimension;
import org.opensearch.index.compositeindex.datacube.NumericDimension;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.NumberFieldMapper;
Expand All @@ -55,6 +57,7 @@

import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.function.BiConsumer;
Expand Down Expand Up @@ -149,21 +152,58 @@ public void testStarTreeDocValues() throws IOException {
ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME);
AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);

List<Dimension> supportedDimensions = new LinkedList<>();
supportedDimensions.add(new NumericDimension(SNDV));
supportedDimensions.add(new NumericDimension(DV));

Query query = new MatchAllDocsQuery();
// match-all query
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,
sumAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalSum::getValue)
);
testCase(
indexSearcher,
query,
queryBuilder,
maxAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalMax::getValue)
);
testCase(
indexSearcher,
query,
queryBuilder,
minAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalMin::getValue)
);
testCase(
indexSearcher,
query,
queryBuilder,
valueCountAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalValueCount::getValue)
);
testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
testCase(
indexSearcher,
query,
queryBuilder,
avgAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalAvg::getValue)
);

// Numeric-terms query
for (int cases = 0; cases < 100; cases++) {
Expand All @@ -180,18 +220,51 @@ public void testStarTreeDocValues() throws IOException {
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,
sumAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalSum::getValue)
);
testCase(
indexSearcher,
query,
queryBuilder,
maxAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalMax::getValue)
);
testCase(
indexSearcher,
query,
queryBuilder,
minAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalMin::getValue)
);
testCase(
indexSearcher,
query,
queryBuilder,
valueCountAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalValueCount::getValue)
);
testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
testCase(
indexSearcher,
query,
queryBuilder,
avgAggregationBuilder,
starTree,
supportedDimensions,
verifyAggregation(InternalAvg::getValue)
);
}

ir.close();
Expand All @@ -212,6 +285,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
QueryBuilder queryBuilder,
T aggBuilder,
CompositeIndexFieldInfo starTree,
List<Dimension> supportedDimensions,
BiConsumer<V, V> verify
) throws IOException {
V starTreeAggregation = searchAndReduceStarTree(
Expand All @@ -221,6 +295,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
queryBuilder,
aggBuilder,
starTree,
supportedDimensions,
DEFAULT_MAX_BUCKETS,
false,
DEFAULT_MAPPED_FIELD
Expand All @@ -232,6 +307,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
queryBuilder,
aggBuilder,
null,
null,
DEFAULT_MAX_BUCKETS,
false,
DEFAULT_MAPPED_FIELD
Expand Down
73 changes: 57 additions & 16 deletions test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,14 +93,34 @@
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.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.BinaryFieldMapper;
import org.opensearch.index.mapper.CompletionFieldMapper;
import org.opensearch.index.mapper.CompositeMappedFieldType;
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.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.QueryBuilder;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.index.shard.IndexShard;
Expand Down Expand Up @@ -139,7 +159,6 @@
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
Expand Down Expand Up @@ -327,19 +346,26 @@ protected CountingAggregator createCountingAggregator(
IndexSearcher indexSearcher,
IndexSettings indexSettings,
CompositeIndexFieldInfo starTree,
List<Dimension> supportedDimensions,
MultiBucketConsumer bucketConsumer,
MappedFieldType... fieldTypes
) throws IOException {
SearchContext searchContext;
if (starTree != null) {
searchContext = createSearchContextWithStarTreeContext(indexSearcher, indexSettings, query, queryBuilder, starTree, bucketConsumer, fieldTypes);
searchContext = createSearchContextWithStarTreeContext(
indexSearcher,
indexSettings,
query,
queryBuilder,
starTree,
supportedDimensions,
bucketConsumer,
fieldTypes
);
} else {
searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, fieldTypes);
}
return new CountingAggregator(
new AtomicInteger(),
createAggregator(aggregationBuilder, searchContext)
);
return new CountingAggregator(new AtomicInteger(), createAggregator(aggregationBuilder, searchContext));
}

/**
Expand All @@ -361,27 +387,32 @@ protected SearchContext createSearchContextWithStarTreeContext(
Query query,
QueryBuilder queryBuilder,
CompositeIndexFieldInfo starTree,
List<Dimension> supportedDimensions,
MultiBucketConsumer bucketConsumer,
MappedFieldType... fieldTypes
) throws IOException {
SearchContext searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes);
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[]{});
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<CompositeMappedFieldType> compositeFieldTypes = Set.of(compositeMappedFieldType);

List<Dimension> dimensions = new LinkedList<>();
dimensions.add(new NumericDimension("sndv"));
dimensions.add(new NumericDimension("dv"));
when((compositeMappedFieldType).getDimensions()).thenReturn(dimensions);
when((compositeMappedFieldType).getDimensions()).thenReturn(supportedDimensions);
MapperService mapperService = mock(MapperService.class);
when(mapperService.getCompositeFieldTypes()).thenReturn(compositeFieldTypes);
when(searchContext.mapperService()).thenReturn(mapperService);
Expand All @@ -391,7 +422,6 @@ protected SearchContext createSearchContextWithStarTreeContext(

when(searchContext.getStarTreeQueryContext()).thenReturn(starTreeQueryContext);
return searchContext;

}

protected SearchContext createSearchContext(
Expand Down Expand Up @@ -708,6 +738,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
QueryBuilder queryBuilder,
AggregationBuilder builder,
CompositeIndexFieldInfo compositeIndexFieldInfo,
List<Dimension> supportedDimensions,
int maxBucket,
boolean hasNested,
MappedFieldType... fieldTypes
Expand All @@ -724,7 +755,17 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
maxBucket,
new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
);
CountingAggregator countingAggregator = createCountingAggregator(query, queryBuilder, builder, searcher, indexSettings, compositeIndexFieldInfo, bucketConsumer, fieldTypes);
CountingAggregator countingAggregator = createCountingAggregator(
query,
queryBuilder,
builder,
searcher,
indexSettings,
compositeIndexFieldInfo,
supportedDimensions,
bucketConsumer,
fieldTypes
);

countingAggregator.preCollection();
searcher.search(query, countingAggregator);
Expand Down

0 comments on commit 2ac5915

Please sign in to comment.