Skip to content

Commit e4a270a

Browse files
sandeshkr419sandeshkr419
committed
Adding test cases and max,min aggregator support
Signed-off-by: Sandesh Kumar <[email protected]>
1 parent e16e01f commit e4a270a

File tree

10 files changed

+366
-24
lines changed

10 files changed

+366
-24
lines changed

‎server/src/main/java/org/opensearch/index/query/QueryShardContext.java‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@
7979
import org.opensearch.script.ScriptFactory;
8080
import org.opensearch.script.ScriptService;
8181
import org.opensearch.search.aggregations.AggregatorFactory;
82+
import org.opensearch.search.aggregations.metrics.MaxAggregatorFactory;
83+
import org.opensearch.search.aggregations.metrics.MinAggregatorFactory;
8284
import org.opensearch.search.aggregations.metrics.SumAggregatorFactory;
8385
import org.opensearch.search.aggregations.support.AggregationUsageService;
8486
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
@@ -585,7 +587,7 @@ private Map<String, List<Predicate<Long>>> getStarTreePredicates(QueryBuilder qu
585587
}
586588

587589
public boolean validateStarTreeMetricSuport(CompositeDataCubeFieldType compositeIndexFieldInfo, AggregatorFactory aggregatorFactory) {
588-
String field = null;
590+
String field;
589591
Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
590592
.stream()
591593
.collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
@@ -595,14 +597,16 @@ public boolean validateStarTreeMetricSuport(CompositeDataCubeFieldType composite
595597
return false;
596598
}
597599

598-
// TODO: increment supported aggregation type
599600
if (aggregatorFactory instanceof SumAggregatorFactory) {
600601
field = ((SumAggregatorFactory) aggregatorFactory).getField();
601-
if (supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(MetricStat.SUM)) {
602-
return true;
603-
}
602+
return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(MetricStat.SUM);
603+
} else if (aggregatorFactory instanceof MaxAggregatorFactory) {
604+
field = ((MaxAggregatorFactory) aggregatorFactory).getField();
605+
return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(MetricStat.MAX);
606+
} else if (aggregatorFactory instanceof MinAggregatorFactory) {
607+
field = ((MinAggregatorFactory) aggregatorFactory).getField();
608+
return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(MetricStat.MIN);
604609
}
605-
606610
return false;
607611
}
608612

‎server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,16 @@
3434
import org.apache.lucene.index.LeafReader;
3535
import org.apache.lucene.index.LeafReaderContext;
3636
import org.apache.lucene.index.PointValues;
37+
import org.apache.lucene.index.SortedNumericDocValues;
3738
import org.apache.lucene.search.CollectionTerminatedException;
3839
import org.apache.lucene.search.ScoreMode;
3940
import org.apache.lucene.util.Bits;
4041
import org.opensearch.common.lease.Releasables;
4142
import org.opensearch.common.util.BigArrays;
4243
import org.opensearch.common.util.DoubleArray;
44+
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
45+
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
46+
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
4347
import org.opensearch.index.fielddata.NumericDoubleValues;
4448
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
4549
import org.opensearch.search.DocValueFormat;
@@ -51,6 +55,8 @@
5155
import org.opensearch.search.aggregations.support.ValuesSource;
5256
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
5357
import org.opensearch.search.internal.SearchContext;
58+
import org.opensearch.search.startree.OriginalOrStarTreeQuery;
59+
import org.opensearch.search.startree.StarTreeQuery;
5460

5561
import java.io.IOException;
5662
import java.util.Arrays;
@@ -120,6 +126,16 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
120126
throw new CollectionTerminatedException();
121127
}
122128
}
129+
130+
if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
131+
StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
132+
return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
133+
}
134+
return getDefaultLeafCollector(ctx, sub);
135+
}
136+
137+
private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
138+
123139
final BigArrays bigArrays = context.bigArrays();
124140
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
125141
final NumericDoubleValues values = MultiValueMode.MAX.select(allValues);
@@ -143,6 +159,34 @@ public void collect(int doc, long bucket) throws IOException {
143159
};
144160
}
145161

162+
private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
163+
throws IOException {
164+
StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
165+
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
166+
String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "max");
167+
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
168+
169+
final BigArrays bigArrays = context.bigArrays();
170+
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
171+
return new LeafBucketCollectorBase(sub, allValues) {
172+
173+
@Override
174+
public void collect(int doc, long bucket) throws IOException {
175+
if (bucket >= maxes.size()) {
176+
long from = maxes.size();
177+
maxes = bigArrays.grow(maxes, bucket + 1);
178+
maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY);
179+
}
180+
if (values.advanceExact(doc)) {
181+
final double value = Double.longBitsToDouble(values.nextValue());
182+
double max = maxes.get(bucket);
183+
max = Math.max(max, value);
184+
maxes.set(bucket, max);
185+
}
186+
}
187+
};
188+
}
189+
146190
@Override
147191
public double metric(long owningBucketOrd) {
148192
if (valuesSource == null || owningBucketOrd >= maxes.size()) {

‎server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@
5252
*
5353
* @opensearch.internal
5454
*/
55-
class MaxAggregatorFactory extends ValuesSourceAggregatorFactory {
55+
public class MaxAggregatorFactory extends ValuesSourceAggregatorFactory {
5656

5757
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
5858
builder.register(

‎server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,16 @@
3434
import org.apache.lucene.index.LeafReader;
3535
import org.apache.lucene.index.LeafReaderContext;
3636
import org.apache.lucene.index.PointValues;
37+
import org.apache.lucene.index.SortedNumericDocValues;
3738
import org.apache.lucene.search.CollectionTerminatedException;
3839
import org.apache.lucene.search.ScoreMode;
3940
import org.apache.lucene.util.Bits;
4041
import org.opensearch.common.lease.Releasables;
4142
import org.opensearch.common.util.BigArrays;
4243
import org.opensearch.common.util.DoubleArray;
44+
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
45+
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
46+
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
4347
import org.opensearch.index.fielddata.NumericDoubleValues;
4448
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
4549
import org.opensearch.search.DocValueFormat;
@@ -51,6 +55,8 @@
5155
import org.opensearch.search.aggregations.support.ValuesSource;
5256
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
5357
import org.opensearch.search.internal.SearchContext;
58+
import org.opensearch.search.startree.OriginalOrStarTreeQuery;
59+
import org.opensearch.search.startree.StarTreeQuery;
5460

5561
import java.io.IOException;
5662
import java.util.Map;
@@ -119,6 +125,16 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
119125
throw new CollectionTerminatedException();
120126
}
121127
}
128+
129+
if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
130+
StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
131+
return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
132+
}
133+
return getDefaultLeafCollector(ctx, sub);
134+
}
135+
136+
private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub)
137+
throws IOException {
122138
final BigArrays bigArrays = context.bigArrays();
123139
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
124140
final NumericDoubleValues values = MultiValueMode.MIN.select(allValues);
@@ -138,10 +154,38 @@ public void collect(int doc, long bucket) throws IOException {
138154
mins.set(bucket, min);
139155
}
140156
}
157+
};
158+
}
141159

160+
private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
161+
throws IOException {
162+
StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
163+
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
164+
String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "min");
165+
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
166+
167+
final BigArrays bigArrays = context.bigArrays();
168+
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
169+
return new LeafBucketCollectorBase(sub, allValues) {
170+
171+
@Override
172+
public void collect(int doc, long bucket) throws IOException {
173+
if (bucket >= mins.size()) {
174+
long from = mins.size();
175+
mins = bigArrays.grow(mins, bucket + 1);
176+
mins.fill(from, mins.size(), Double.POSITIVE_INFINITY);
177+
}
178+
if (values.advanceExact(doc)) {
179+
final double value = Double.longBitsToDouble(values.nextValue());
180+
double min = mins.get(bucket);
181+
min = Math.min(min, value);
182+
mins.set(bucket, min);
183+
}
184+
}
142185
};
143186
}
144187

188+
145189
@Override
146190
public double metric(long owningBucketOrd) {
147191
if (valuesSource == null || owningBucketOrd >= mins.size()) {

‎server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@
5252
*
5353
* @opensearch.internal
5454
*/
55-
class MinAggregatorFactory extends ValuesSourceAggregatorFactory {
55+
public class MinAggregatorFactory extends ValuesSourceAggregatorFactory {
5656

5757
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
5858
builder.register(

‎server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,10 @@ public ScoreMode scoreMode() {
9292

9393
@Override
9494
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
95+
if (valuesSource == null) {
96+
return LeafBucketCollector.NO_OP_COLLECTOR;
97+
}
98+
9599
if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
96100
StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
97101
return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
@@ -100,9 +104,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
100104
}
101105

102106
private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
103-
if (valuesSource == null) {
104-
return LeafBucketCollector.NO_OP_COLLECTOR;
105-
}
106107
final BigArrays bigArrays = context.bigArrays();
107108
final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
108109
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
@@ -134,15 +135,14 @@ public void collect(int doc, long bucket) throws IOException {
134135

135136
private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
136137
throws IOException {
137-
final BigArrays bigArrays = context.bigArrays();
138-
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
139-
140138
StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
141139
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
142140
String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "sum");
143-
144141
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
145142

143+
final BigArrays bigArrays = context.bigArrays();
144+
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
145+
146146
return new LeafBucketCollectorBase(sub, values) {
147147
@Override
148148
public void collect(int doc, long bucket) throws IOException {

‎server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,12 +45,12 @@ public void visit(QueryVisitor queryVisitor) {
4545

4646
@Override
4747
public boolean equals(Object o) {
48-
return false;
48+
return true;
4949
}
5050

5151
@Override
5252
public int hashCode() {
53-
return 0;
53+
return originalQuery.hashCode();
5454
}
5555

5656
@Override

‎server/src/test/java/org/opensearch/index/codec/composite99/datacube/startree/StarTreeDocValuesFormatTests.java‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,9 @@
7272
*/
7373
@LuceneTestCase.SuppressSysoutChecks(bugUrl = "we log a lot on purpose")
7474
public class StarTreeDocValuesFormatTests extends BaseDocValuesFormatTestCase {
75-
MapperService mapperService = null;
75+
7676
StarTreeFieldConfiguration.StarTreeBuildMode buildMode;
77+
MapperService mapperService;
7778

7879
public StarTreeDocValuesFormatTests(StarTreeFieldConfiguration.StarTreeBuildMode buildMode) {
7980
this.buildMode = buildMode;
@@ -105,13 +106,14 @@ public void teardown() throws IOException {
105106
@Override
106107
protected Codec getCodec() {
107108
final Logger testLogger = LogManager.getLogger(StarTreeDocValuesFormatTests.class);
108-
109+
Codec codec;
109110
try {
110-
createMapperService(getExpandedMapping());
111+
mapperService = createMapperService(getExpandedMapping());
112+
codec = new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger);
111113
} catch (IOException e) {
112114
throw new RuntimeException(e);
113115
}
114-
Codec codec = new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger);
116+
115117
return codec;
116118
}
117119

@@ -195,7 +197,7 @@ public void testStarTreeDocValues() throws IOException {
195197
directory.close();
196198
}
197199

198-
private XContentBuilder getExpandedMapping() throws IOException {
200+
public static XContentBuilder getExpandedMapping() throws IOException {
199201
return topMapping(b -> {
200202
b.startObject("composite");
201203
b.startObject("startree");
@@ -215,6 +217,8 @@ private XContentBuilder getExpandedMapping() throws IOException {
215217
b.field("name", "field");
216218
b.startArray("stats");
217219
b.value("sum");
220+
b.value("max");
221+
b.value("min");
218222
b.value("count");
219223
b.endArray();
220224
b.endObject();
@@ -236,13 +240,14 @@ private XContentBuilder getExpandedMapping() throws IOException {
236240
});
237241
}
238242

239-
private XContentBuilder topMapping(CheckedConsumer<XContentBuilder, IOException> buildFields) throws IOException {
243+
private static XContentBuilder topMapping(CheckedConsumer<XContentBuilder, IOException> buildFields) throws IOException {
240244
XContentBuilder builder = XContentFactory.jsonBuilder().startObject().startObject("_doc");
241245
buildFields.accept(builder);
242246
return builder.endObject().endObject();
243247
}
244248

245-
private void createMapperService(XContentBuilder builder) throws IOException {
249+
public static MapperService createMapperService(XContentBuilder builder) throws IOException {
250+
MapperService mapperService = null;
246251
IndexMetadata indexMetadata = IndexMetadata.builder("test")
247252
.settings(
248253
Settings.builder()
@@ -261,5 +266,6 @@ private void createMapperService(XContentBuilder builder) throws IOException {
261266
"test"
262267
);
263268
mapperService.merge(indexMetadata, MapperService.MergeReason.INDEX_TEMPLATE);
269+
return mapperService;
264270
}
265271
}

0 commit comments

Comments
 (0)