Refactoring and changelog addition

Signed-off-by: Sandesh Kumar <[email protected]>
sandeshkr419 · Feb 3, 2024 · 2a9cce7 · 2a9cce7
1 parent 6d05716
commit 2a9cce7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -190,6 +190,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Restore support for Java 8 for RestClient ([#11562](https://github.com/opensearch-project/OpenSearch/pull/11562))
 - Add deleted doc count in _cat/shards ([#11678](https://github.com/opensearch-project/OpenSearch/pull/11678))
 - Capture information for additional query types and aggregation types ([#11582](https://github.com/opensearch-project/OpenSearch/pull/11582))
+- Improve string terms aggregation performance using Collector#setWeight ([#11643](https://github.com/opensearch-project/OpenSearch/pull/11643))
 
 ### Deprecated
 

diff --git a/.../org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java b/.../org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java
@@ -40,6 +40,7 @@
 import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.index.Terms;
 import org.apache.lucene.index.TermsEnum;
+import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BytesRef;
@@ -92,7 +93,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
     private final LongPredicate acceptedGlobalOrdinals;
     private final long valueCount;
 
-    private final String fieldName;
+    private String fieldName;
 
     private Weight weight;
     private final GlobalOrdLookupFunction lookupGlobalOrd;
@@ -146,7 +147,8 @@ public GlobalOrdinalsStringTermsAggregator(
                 return new DenseGlobalOrds();
             });
         }
-        this.fieldName = ((ValuesSource.Bytes.WithOrdinals.FieldData) valuesSource).getIndexFieldName();
+            this.fieldName = (valuesSource instanceof ValuesSource.Bytes.WithOrdinals.FieldData) ?
+                ((ValuesSource.Bytes.WithOrdinals.FieldData) valuesSource).getIndexFieldName() : null;
     }
 
     String descriptCollectionStrategy() {
@@ -170,6 +172,10 @@ LeafBucketCollector termDocFreqCollector(
         SortedSetDocValues globalOrds,
         BiConsumer<Long, Integer> ordCountConsumer
     ) throws IOException {
+        if (weight == null) {
+            // Calculate weight if not assigned previously
+            this.weight = context.searcher().createWeight(context.query(), ScoreMode.COMPLETE_NO_SCORES, 1f);
+        }
         if (weight.count(ctx) != ctx.reader().maxDoc()) {
             // Top-level query does not match all docs in this segment.
             return null;
@@ -221,7 +227,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol
                 (ord, docCount) -> incrementBucketDocCount(collectionStrategy.globalOrdToBucketOrd(0, ord), docCount)
             );
             if (termDocFreqCollector != null) {
-                return null;
+                return termDocFreqCollector;
             }
         }