Skip to content

Commit e7e19f7

Browse files
expaniexpani
authored
Changes to support unmapped fields in metric aggregation (opensearch-project#16481)
Avoids exception when querying unmapped field when star tree experimental feature is enables. --------- Signed-off-by: expani <anijainc@amazon.com>
1 parent 0b36599 commit e7e19f7

File tree

5 files changed

+172
-6
lines changed

5 files changed

+172
-6
lines changed

‎server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ private static MetricStat validateStarTreeMetricSupport(
152152
MetricStat metricStat = ((MetricAggregatorFactory) aggregatorFactory).getMetricStat();
153153
field = ((MetricAggregatorFactory) aggregatorFactory).getField();
154154

155-
if (supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat)) {
155+
if (field != null && supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat)) {
156156
return metricStat;
157157
}
158158
}

‎server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,6 @@ public String getStatsSubtype() {
104104
}
105105

106106
public String getField() {
107-
return config.fieldContext().field();
107+
return config.fieldContext() != null ? config.fieldContext().field() : null;
108108
}
109109
}

‎server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java

+139
Original file line numberDiff line numberDiff line change
@@ -28,18 +28,27 @@
2828
import org.opensearch.common.lucene.Lucene;
2929
import org.opensearch.common.settings.Settings;
3030
import org.opensearch.common.util.FeatureFlags;
31+
import org.opensearch.common.util.MockBigArrays;
32+
import org.opensearch.common.util.MockPageCacheRecycler;
33+
import org.opensearch.core.indices.breaker.CircuitBreakerService;
34+
import org.opensearch.core.indices.breaker.NoneCircuitBreakerService;
3135
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
3236
import org.opensearch.index.codec.composite.CompositeIndexReader;
3337
import org.opensearch.index.codec.composite.composite912.Composite912Codec;
3438
import org.opensearch.index.codec.composite912.datacube.startree.StarTreeDocValuesFormatTests;
3539
import org.opensearch.index.compositeindex.datacube.Dimension;
40+
import org.opensearch.index.compositeindex.datacube.Metric;
41+
import org.opensearch.index.compositeindex.datacube.MetricStat;
3642
import org.opensearch.index.compositeindex.datacube.NumericDimension;
3743
import org.opensearch.index.mapper.MappedFieldType;
3844
import org.opensearch.index.mapper.MapperService;
3945
import org.opensearch.index.mapper.NumberFieldMapper;
4046
import org.opensearch.index.query.QueryBuilder;
47+
import org.opensearch.index.query.QueryShardContext;
4148
import org.opensearch.index.query.TermQueryBuilder;
4249
import org.opensearch.search.aggregations.AggregationBuilder;
50+
import org.opensearch.search.aggregations.AggregatorFactories;
51+
import org.opensearch.search.aggregations.AggregatorFactory;
4352
import org.opensearch.search.aggregations.AggregatorTestCase;
4453
import org.opensearch.search.aggregations.InternalAggregation;
4554
import org.opensearch.search.aggregations.metrics.AvgAggregationBuilder;
@@ -49,14 +58,17 @@
4958
import org.opensearch.search.aggregations.metrics.InternalSum;
5059
import org.opensearch.search.aggregations.metrics.InternalValueCount;
5160
import org.opensearch.search.aggregations.metrics.MaxAggregationBuilder;
61+
import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory;
5262
import org.opensearch.search.aggregations.metrics.MinAggregationBuilder;
5363
import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
5464
import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
65+
import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
5566
import org.junit.After;
5667
import org.junit.Before;
5768

5869
import java.io.IOException;
5970
import java.util.ArrayList;
71+
import java.util.Collections;
6072
import java.util.LinkedList;
6173
import java.util.List;
6274
import java.util.Random;
@@ -69,6 +81,8 @@
6981
import static org.opensearch.search.aggregations.AggregationBuilders.min;
7082
import static org.opensearch.search.aggregations.AggregationBuilders.sum;
7183
import static org.opensearch.test.InternalAggregationTestCase.DEFAULT_MAX_BUCKETS;
84+
import static org.mockito.Mockito.mock;
85+
import static org.mockito.Mockito.when;
7286

7387
public class MetricAggregatorTests extends AggregatorTestCase {
7488

@@ -267,6 +281,110 @@ public void testStarTreeDocValues() throws IOException {
267281
);
268282
}
269283

284+
CircuitBreakerService circuitBreakerService = new NoneCircuitBreakerService();
285+
286+
QueryShardContext queryShardContext = queryShardContextMock(
287+
indexSearcher,
288+
mapperServiceMock(),
289+
createIndexSettings(),
290+
circuitBreakerService,
291+
new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), circuitBreakerService).withCircuitBreaking()
292+
);
293+
294+
MetricAggregatorFactory aggregatorFactory = mock(MetricAggregatorFactory.class);
295+
when(aggregatorFactory.getSubFactories()).thenReturn(AggregatorFactories.EMPTY);
296+
when(aggregatorFactory.getField()).thenReturn(FIELD_NAME);
297+
when(aggregatorFactory.getMetricStat()).thenReturn(MetricStat.SUM);
298+
299+
// Case when field and metric type in aggregation are fully supported by star tree.
300+
testCase(
301+
indexSearcher,
302+
query,
303+
queryBuilder,
304+
sumAggregationBuilder,
305+
starTree,
306+
supportedDimensions,
307+
List.of(new Metric(FIELD_NAME, List.of(MetricStat.SUM, MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
308+
verifyAggregation(InternalSum::getValue),
309+
aggregatorFactory,
310+
true
311+
);
312+
313+
// Case when the field is not supported by star tree
314+
SumAggregationBuilder invalidFieldSumAggBuilder = sum("_name").field("hello");
315+
testCase(
316+
indexSearcher,
317+
query,
318+
queryBuilder,
319+
invalidFieldSumAggBuilder,
320+
starTree,
321+
supportedDimensions,
322+
Collections.emptyList(),
323+
verifyAggregation(InternalSum::getValue),
324+
invalidFieldSumAggBuilder.build(queryShardContext, null),
325+
false // Invalid fields will return null StarTreeQueryContext which will not cause early termination by leaf collector
326+
);
327+
328+
// Case when metric type in aggregation is not supported by star tree but the field is supported.
329+
testCase(
330+
indexSearcher,
331+
query,
332+
queryBuilder,
333+
sumAggregationBuilder,
334+
starTree,
335+
supportedDimensions,
336+
List.of(new Metric(FIELD_NAME, List.of(MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
337+
verifyAggregation(InternalSum::getValue),
338+
aggregatorFactory,
339+
false
340+
);
341+
342+
// Case when field is not present in supported metrics
343+
testCase(
344+
indexSearcher,
345+
query,
346+
queryBuilder,
347+
sumAggregationBuilder,
348+
starTree,
349+
supportedDimensions,
350+
List.of(new Metric("hello", List.of(MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
351+
verifyAggregation(InternalSum::getValue),
352+
aggregatorFactory,
353+
false
354+
);
355+
356+
AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class);
357+
when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] { mock(MetricAggregatorFactory.class) });
358+
when(aggregatorFactory.getSubFactories()).thenReturn(aggregatorFactories);
359+
360+
// Case when sub aggregations are present
361+
testCase(
362+
indexSearcher,
363+
query,
364+
queryBuilder,
365+
sumAggregationBuilder,
366+
starTree,
367+
supportedDimensions,
368+
List.of(new Metric("hello", List.of(MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
369+
verifyAggregation(InternalSum::getValue),
370+
aggregatorFactory,
371+
false
372+
);
373+
374+
// Case when aggregation factory is not metric aggregation
375+
testCase(
376+
indexSearcher,
377+
query,
378+
queryBuilder,
379+
sumAggregationBuilder,
380+
starTree,
381+
supportedDimensions,
382+
List.of(new Metric("hello", List.of(MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
383+
verifyAggregation(InternalSum::getValue),
384+
mock(ValuesSourceAggregatorFactory.class),
385+
false
386+
);
387+
270388
ir.close();
271389
directory.close();
272390
}
@@ -287,6 +405,21 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
287405
CompositeIndexFieldInfo starTree,
288406
List<Dimension> supportedDimensions,
289407
BiConsumer<V, V> verify
408+
) throws IOException {
409+
testCase(searcher, query, queryBuilder, aggBuilder, starTree, supportedDimensions, Collections.emptyList(), verify, null, true);
410+
}
411+
412+
private <T extends AggregationBuilder, V extends InternalAggregation> void testCase(
413+
IndexSearcher searcher,
414+
Query query,
415+
QueryBuilder queryBuilder,
416+
T aggBuilder,
417+
CompositeIndexFieldInfo starTree,
418+
List<Dimension> supportedDimensions,
419+
List<Metric> supportedMetrics,
420+
BiConsumer<V, V> verify,
421+
AggregatorFactory aggregatorFactory,
422+
boolean assertCollectorEarlyTermination
290423
) throws IOException {
291424
V starTreeAggregation = searchAndReduceStarTree(
292425
createIndexSettings(),
@@ -296,8 +429,11 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
296429
aggBuilder,
297430
starTree,
298431
supportedDimensions,
432+
supportedMetrics,
299433
DEFAULT_MAX_BUCKETS,
300434
false,
435+
aggregatorFactory,
436+
assertCollectorEarlyTermination,
301437
DEFAULT_MAPPED_FIELD
302438
);
303439
V expectedAggregation = searchAndReduceStarTree(
@@ -308,8 +444,11 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
308444
aggBuilder,
309445
null,
310446
null,
447+
null,
311448
DEFAULT_MAX_BUCKETS,
312449
false,
450+
aggregatorFactory,
451+
assertCollectorEarlyTermination,
313452
DEFAULT_MAPPED_FIELD
314453
);
315454
verify.accept(expectedAggregation, starTreeAggregation);

‎server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java

+11-2
Original file line numberDiff line numberDiff line change
@@ -87,15 +87,15 @@ public void testStarTreeFilterWithDocsInSVDFieldButNoStarNode() throws IOExcepti
8787
testStarTreeFilter(10, false);
8888
}
8989

90-
private void testStarTreeFilter(int maxLeafDoc, boolean skipStarNodeCreationForSDVDimension) throws IOException {
90+
private Directory createStarTreeIndex(int maxLeafDoc, boolean skipStarNodeCreationForSDVDimension, List<Document> docs)
91+
throws IOException {
9192
Directory directory = newDirectory();
9293
IndexWriterConfig conf = newIndexWriterConfig(null);
9394
conf.setCodec(getCodec(maxLeafDoc, skipStarNodeCreationForSDVDimension));
9495
conf.setMergePolicy(newLogMergePolicy());
9596
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
9697
int totalDocs = 100;
9798

98-
List<Document> docs = new ArrayList<>();
9999
for (int i = 0; i < totalDocs; i++) {
100100
Document doc = new Document();
101101
doc.add(new SortedNumericDocValuesField(SNDV, i));
@@ -110,6 +110,15 @@ private void testStarTreeFilter(int maxLeafDoc, boolean skipStarNodeCreationForS
110110
}
111111
iw.forceMerge(1);
112112
iw.close();
113+
return directory;
114+
}
115+
116+
private void testStarTreeFilter(int maxLeafDoc, boolean skipStarNodeCreationForSDVDimension) throws IOException {
117+
List<Document> docs = new ArrayList<>();
118+
119+
Directory directory = createStarTreeIndex(maxLeafDoc, skipStarNodeCreationForSDVDimension, docs);
120+
121+
int totalDocs = docs.size();
113122

114123
DirectoryReader ir = DirectoryReader.open(directory);
115124
initValuesSourceRegistry();

‎test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java

+20-2
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@
9393
import org.opensearch.index.cache.query.DisabledQueryCache;
9494
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
9595
import org.opensearch.index.compositeindex.datacube.Dimension;
96+
import org.opensearch.index.compositeindex.datacube.Metric;
9697
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
9798
import org.opensearch.index.fielddata.IndexFieldData;
9899
import org.opensearch.index.fielddata.IndexFieldDataCache;
@@ -348,7 +349,9 @@ protected CountingAggregator createCountingAggregator(
348349
IndexSettings indexSettings,
349350
CompositeIndexFieldInfo starTree,
350351
List<Dimension> supportedDimensions,
352+
List<Metric> supportedMetrics,
351353
MultiBucketConsumer bucketConsumer,
354+
AggregatorFactory aggregatorFactory,
352355
MappedFieldType... fieldTypes
353356
) throws IOException {
354357
SearchContext searchContext;
@@ -360,7 +363,9 @@ protected CountingAggregator createCountingAggregator(
360363
queryBuilder,
361364
starTree,
362365
supportedDimensions,
366+
supportedMetrics,
363367
bucketConsumer,
368+
aggregatorFactory,
364369
fieldTypes
365370
);
366371
} else {
@@ -389,7 +394,9 @@ protected SearchContext createSearchContextWithStarTreeContext(
389394
QueryBuilder queryBuilder,
390395
CompositeIndexFieldInfo starTree,
391396
List<Dimension> supportedDimensions,
397+
List<Metric> supportedMetrics,
392398
MultiBucketConsumer bucketConsumer,
399+
AggregatorFactory aggregatorFactory,
393400
MappedFieldType... fieldTypes
394401
) throws IOException {
395402
SearchContext searchContext = createSearchContext(
@@ -406,14 +413,20 @@ protected SearchContext createSearchContextWithStarTreeContext(
406413
AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class);
407414
when(searchContext.aggregations()).thenReturn(searchContextAggregations);
408415
when(searchContextAggregations.factories()).thenReturn(aggregatorFactories);
409-
when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] {});
416+
417+
if (aggregatorFactory != null) {
418+
when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] { aggregatorFactory });
419+
} else {
420+
when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] {});
421+
}
410422

411423
CompositeDataCubeFieldType compositeMappedFieldType = mock(CompositeDataCubeFieldType.class);
412424
when(compositeMappedFieldType.name()).thenReturn(starTree.getField());
413425
when(compositeMappedFieldType.getCompositeIndexType()).thenReturn(starTree.getType());
414426
Set<CompositeMappedFieldType> compositeFieldTypes = Set.of(compositeMappedFieldType);
415427

416428
when((compositeMappedFieldType).getDimensions()).thenReturn(supportedDimensions);
429+
when((compositeMappedFieldType).getMetrics()).thenReturn(supportedMetrics);
417430
MapperService mapperService = mock(MapperService.class);
418431
when(mapperService.getCompositeFieldTypes()).thenReturn(compositeFieldTypes);
419432
when(searchContext.mapperService()).thenReturn(mapperService);
@@ -740,8 +753,11 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
740753
AggregationBuilder builder,
741754
CompositeIndexFieldInfo compositeIndexFieldInfo,
742755
List<Dimension> supportedDimensions,
756+
List<Metric> supportedMetrics,
743757
int maxBucket,
744758
boolean hasNested,
759+
AggregatorFactory aggregatorFactory,
760+
boolean assertCollectorEarlyTermination,
745761
MappedFieldType... fieldTypes
746762
) throws IOException {
747763
query = query.rewrite(searcher);
@@ -764,15 +780,17 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
764780
indexSettings,
765781
compositeIndexFieldInfo,
766782
supportedDimensions,
783+
supportedMetrics,
767784
bucketConsumer,
785+
aggregatorFactory,
768786
fieldTypes
769787
);
770788

771789
countingAggregator.preCollection();
772790
searcher.search(query, countingAggregator);
773791
countingAggregator.postCollection();
774792
aggs.add(countingAggregator.buildTopLevel());
775-
if (compositeIndexFieldInfo != null) {
793+
if (compositeIndexFieldInfo != null && assertCollectorEarlyTermination) {
776794
assertEquals(0, countingAggregator.collectCounter.get());
777795
}
778796

0 commit comments

Comments
 (0)