Skip to content

Commit e417086

Browse files
sandeshkr419sandeshkr419
authored
Support for star-tree keyword aggregation dfs mode (#18354)
Signed-off-by: Sandesh Kumar <[email protected]>
1 parent 4d6f8ba commit e417086

File tree

4 files changed

+107
-95
lines changed

4 files changed

+107
-95
lines changed

‎server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -109,15 +109,6 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
109109
protected int segmentsWithMultiValuedOrds = 0;
110110
LongUnaryOperator globalOperator;
111111

112-
/**
113-
* Lookup global ordinals
114-
*
115-
* @opensearch.internal
116-
*/
117-
public interface GlobalOrdLookupFunction {
118-
BytesRef apply(long ord) throws IOException;
119-
}
120-
121112
public GlobalOrdinalsStringTermsAggregator(
122113
String name,
123114
AggregatorFactories factories,
@@ -375,7 +366,7 @@ public void collectStarTreeEntry(int starTreeEntry, long owningBucketOrd) throws
375366

376367
if (docCountsIterator.advanceExact(starTreeEntry)) {
377368
long metricValue = docCountsIterator.nextValue();
378-
long bucketOrd = collectionStrategy.globalOrdToBucketOrd(0, ord);
369+
long bucketOrd = collectionStrategy.getOrAddBucketOrd(owningBucketOrd, ord);
379370
collectStarTreeBucket(this, metricValue, bucketOrd, starTreeEntry);
380371
}
381372
}
@@ -638,7 +629,7 @@ abstract class CollectionStrategy implements Releasable {
638629
abstract long globalOrdToBucketOrd(long owningBucketOrd, long globalOrd);
639630

640631
/**
641-
* Iterate all of the buckets. Implementations take into account
632+
* Iterate all the buckets. Implementations take into account
642633
* the {@link BucketCountThresholds}. In particular,
643634
* if the {@link BucketCountThresholds#getMinDocCount()} is 0 then
644635
* they'll make sure to iterate a bucket even if it was never
@@ -647,6 +638,12 @@ abstract class CollectionStrategy implements Releasable {
647638
* they'll skip all global ords that weren't collected.
648639
*/
649640
abstract void forEach(long owningBucketOrd, BucketInfoConsumer consumer) throws IOException;
641+
642+
/**
643+
* Add a global ordinal if it hasn't been seen before.
644+
* Convert the global ordinal into a bucket ordinal.
645+
*/
646+
abstract long getOrAddBucketOrd(long owningBucketOrd, long globalOrd) throws IOException;
650647
}
651648

652649
interface BucketInfoConsumer {
@@ -697,6 +694,11 @@ void forEach(long owningBucketOrd, BucketInfoConsumer consumer) throws IOExcepti
697694
}
698695
}
699696

697+
@Override
698+
long getOrAddBucketOrd(long owningBucketOrd, long globalOrd) {
699+
return globalOrdToBucketOrd(owningBucketOrd, globalOrd);
700+
}
701+
700702
@Override
701703
public void close() {}
702704
}
@@ -778,6 +780,11 @@ void forEach(long owningBucketOrd, BucketInfoConsumer consumer) throws IOExcepti
778780
}
779781
}
780782

783+
@Override
784+
long getOrAddBucketOrd(long owningBucketOrd, long globalOrd) {
785+
return bucketOrds.add(owningBucketOrd, globalOrd);
786+
}
787+
781788
@Override
782789
public void close() {
783790
bucketOrds.close();

‎server/src/main/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregator.java

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -291,11 +291,7 @@ private boolean subAggsNeedScore() {
291291

292292
@Override
293293
protected boolean shouldDefer(Aggregator aggregator) {
294-
if (context.getQueryShardContext().getStarTreeQueryContext() == null) {
295-
return collectMode == SubAggCollectionMode.BREADTH_FIRST && !aggsUsedForSorting.contains(aggregator);
296-
} else {
297-
// when pre-computing using star-tree - return false (don't defer) for BREADTH_FIRST case
298-
return collectMode != SubAggCollectionMode.BREADTH_FIRST;
299-
}
294+
return context.getQueryShardContext().getStarTreeQueryContext() == null
295+
&& (collectMode == SubAggCollectionMode.BREADTH_FIRST && !aggsUsedForSorting.contains(aggregator));
300296
}
301297
}

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

Lines changed: 47 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
import org.opensearch.index.mapper.NumberFieldMapper;
4646
import org.opensearch.index.query.QueryBuilder;
4747
import org.opensearch.index.query.TermQueryBuilder;
48-
import org.opensearch.search.aggregations.Aggregator;
48+
import org.opensearch.search.aggregations.Aggregator.SubAggCollectionMode;
4949
import org.opensearch.search.aggregations.AggregatorTestCase;
5050
import org.opensearch.search.aggregations.bucket.terms.InternalTerms;
5151
import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
@@ -66,6 +66,8 @@
6666
import static org.opensearch.search.aggregations.AggregationBuilders.min;
6767
import static org.opensearch.search.aggregations.AggregationBuilders.sum;
6868
import static org.opensearch.search.aggregations.AggregationBuilders.terms;
69+
import static org.opensearch.search.aggregations.Aggregator.SubAggCollectionMode.BREADTH_FIRST;
70+
import static org.opensearch.search.aggregations.Aggregator.SubAggCollectionMode.DEPTH_FIRST;
6971
import static org.opensearch.test.InternalAggregationTestCase.DEFAULT_MAX_BUCKETS;
7072

7173
public class KeywordTermsAggregatorTests extends AggregatorTestCase {
@@ -155,8 +157,7 @@ public void testStarTreeKeywordTerms() throws IOException {
155157

156158
Query query = new MatchAllDocsQuery();
157159
QueryBuilder queryBuilder = null;
158-
TermsAggregationBuilder termsAggregationBuilder = terms("terms_agg").field(CLIENTIP)
159-
.collectMode(Aggregator.SubAggCollectionMode.BREADTH_FIRST);
160+
TermsAggregationBuilder termsAggregationBuilder = terms("terms_agg").field(CLIENTIP);
160161
testCase(indexSearcher, query, queryBuilder, termsAggregationBuilder, starTree, supportedDimensions);
161162

162163
ValuesSourceAggregationBuilder[] aggBuilders = {
@@ -170,9 +171,7 @@ public void testStarTreeKeywordTerms() throws IOException {
170171
query = new MatchAllDocsQuery();
171172
queryBuilder = null;
172173

173-
termsAggregationBuilder = terms("terms_agg").field(CLIENTIP)
174-
.subAggregation(aggregationBuilder)
175-
.collectMode(Aggregator.SubAggCollectionMode.BREADTH_FIRST);
174+
termsAggregationBuilder = terms("terms_agg").field(CLIENTIP).subAggregation(aggregationBuilder);
176175
testCase(indexSearcher, query, queryBuilder, termsAggregationBuilder, starTree, supportedDimensions);
177176

178177
// Numeric-terms query with keyword terms aggregation
@@ -204,43 +203,47 @@ private void testCase(
204203
CompositeIndexFieldInfo starTree,
205204
LinkedHashMap<Dimension, MappedFieldType> supportedDimensions
206205
) throws IOException {
207-
InternalTerms starTreeAggregation = searchAndReduceStarTree(
208-
createIndexSettings(),
209-
indexSearcher,
210-
query,
211-
queryBuilder,
212-
termsAggregationBuilder,
213-
starTree,
214-
supportedDimensions,
215-
null,
216-
DEFAULT_MAX_BUCKETS,
217-
false,
218-
null,
219-
true,
220-
STATUS_FIELD_TYPE,
221-
SIZE_FIELD_NAME,
222-
CLIENTIP_FIELD_NAME
223-
);
224-
225-
InternalTerms defaultAggregation = searchAndReduceStarTree(
226-
createIndexSettings(),
227-
indexSearcher,
228-
query,
229-
queryBuilder,
230-
termsAggregationBuilder,
231-
null,
232-
null,
233-
null,
234-
DEFAULT_MAX_BUCKETS,
235-
false,
236-
null,
237-
false,
238-
STATUS_FIELD_TYPE,
239-
SIZE_FIELD_NAME,
240-
CLIENTIP_FIELD_NAME
241-
);
242-
243-
assertEquals(defaultAggregation.getBuckets().size(), starTreeAggregation.getBuckets().size());
244-
assertEquals(defaultAggregation.getBuckets(), starTreeAggregation.getBuckets());
206+
for (SubAggCollectionMode collectionMode : List.of(DEPTH_FIRST, BREADTH_FIRST)) {
207+
termsAggregationBuilder.collectMode(collectionMode);
208+
209+
InternalTerms starTreeAggregation = searchAndReduceStarTree(
210+
createIndexSettings(),
211+
indexSearcher,
212+
query,
213+
queryBuilder,
214+
termsAggregationBuilder,
215+
starTree,
216+
supportedDimensions,
217+
null,
218+
DEFAULT_MAX_BUCKETS,
219+
false,
220+
null,
221+
true,
222+
STATUS_FIELD_TYPE,
223+
SIZE_FIELD_NAME,
224+
CLIENTIP_FIELD_NAME
225+
);
226+
227+
InternalTerms defaultAggregation = searchAndReduceStarTree(
228+
createIndexSettings(),
229+
indexSearcher,
230+
query,
231+
queryBuilder,
232+
termsAggregationBuilder,
233+
null,
234+
null,
235+
null,
236+
DEFAULT_MAX_BUCKETS,
237+
false,
238+
null,
239+
false,
240+
STATUS_FIELD_TYPE,
241+
SIZE_FIELD_NAME,
242+
CLIENTIP_FIELD_NAME
243+
);
244+
245+
assertEquals(defaultAggregation.getBuckets().size(), starTreeAggregation.getBuckets().size());
246+
assertEquals(defaultAggregation.getBuckets(), starTreeAggregation.getBuckets());
247+
}
245248
}
246249
}

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

Lines changed: 40 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import org.opensearch.index.mapper.NumberFieldMapper;
4141
import org.opensearch.index.query.QueryBuilder;
4242
import org.opensearch.index.query.TermQueryBuilder;
43+
import org.opensearch.search.aggregations.Aggregator;
4344
import org.opensearch.search.aggregations.AggregatorTestCase;
4445
import org.opensearch.search.aggregations.bucket.terms.InternalTerms;
4546
import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
@@ -61,6 +62,8 @@
6162
import static org.opensearch.search.aggregations.AggregationBuilders.min;
6263
import static org.opensearch.search.aggregations.AggregationBuilders.sum;
6364
import static org.opensearch.search.aggregations.AggregationBuilders.terms;
65+
import static org.opensearch.search.aggregations.Aggregator.SubAggCollectionMode.BREADTH_FIRST;
66+
import static org.opensearch.search.aggregations.Aggregator.SubAggCollectionMode.DEPTH_FIRST;
6467
import static org.opensearch.test.InternalAggregationTestCase.DEFAULT_MAX_BUCKETS;
6568

6669
public class NumericTermsAggregatorTests extends AggregatorTestCase {
@@ -205,42 +208,45 @@ private void testCase(
205208
CompositeIndexFieldInfo starTree,
206209
LinkedHashMap<Dimension, MappedFieldType> supportedDimensions
207210
) throws IOException {
208-
InternalTerms starTreeAggregation = searchAndReduceStarTree(
209-
createIndexSettings(),
210-
indexSearcher,
211-
query,
212-
queryBuilder,
213-
termsAggregationBuilder,
214-
starTree,
215-
supportedDimensions,
216-
null,
217-
DEFAULT_MAX_BUCKETS,
218-
false,
219-
null,
220-
true,
221-
STATUS_FIELD_TYPE,
222-
SIZE_FIELD_NAME
223-
);
211+
for (Aggregator.SubAggCollectionMode collectionMode : List.of(DEPTH_FIRST, BREADTH_FIRST)) {
212+
termsAggregationBuilder.collectMode(collectionMode);
213+
InternalTerms starTreeAggregation = searchAndReduceStarTree(
214+
createIndexSettings(),
215+
indexSearcher,
216+
query,
217+
queryBuilder,
218+
termsAggregationBuilder,
219+
starTree,
220+
supportedDimensions,
221+
null,
222+
DEFAULT_MAX_BUCKETS,
223+
false,
224+
null,
225+
true,
226+
STATUS_FIELD_TYPE,
227+
SIZE_FIELD_NAME
228+
);
224229

225-
InternalTerms defaultAggregation = searchAndReduceStarTree(
226-
createIndexSettings(),
227-
indexSearcher,
228-
query,
229-
queryBuilder,
230-
termsAggregationBuilder,
231-
null,
232-
null,
233-
null,
234-
DEFAULT_MAX_BUCKETS,
235-
false,
236-
null,
237-
false,
238-
STATUS_FIELD_TYPE,
239-
SIZE_FIELD_NAME
240-
);
230+
InternalTerms defaultAggregation = searchAndReduceStarTree(
231+
createIndexSettings(),
232+
indexSearcher,
233+
query,
234+
queryBuilder,
235+
termsAggregationBuilder,
236+
null,
237+
null,
238+
null,
239+
DEFAULT_MAX_BUCKETS,
240+
false,
241+
null,
242+
false,
243+
STATUS_FIELD_TYPE,
244+
SIZE_FIELD_NAME
245+
);
241246

242-
assertEquals(defaultAggregation.getBuckets().size(), starTreeAggregation.getBuckets().size());
243-
assertEquals(defaultAggregation.getBuckets(), starTreeAggregation.getBuckets());
247+
assertEquals(defaultAggregation.getBuckets().size(), starTreeAggregation.getBuckets().size());
248+
assertEquals(defaultAggregation.getBuckets(), starTreeAggregation.getBuckets());
249+
}
244250
}
245251

246252
public static XContentBuilder getExpandedMapping(int maxLeafDocs, boolean skipStarNodeCreationForStatusDimension) throws IOException {

0 commit comments

Comments
 (0)