Overriding bulk collection API for few other aggregations

jainankitk · jainankitk · commit c967b02afd1b · 2025-11-11T17:51:55.000-08:00
Signed-off-by: Ankit Jain &lt;jainankitk@apache.org&gt;
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java
@@ -42,6 +42,7 @@
 import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.DocIdSet;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.FieldComparator;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.LeafFieldComparator;
@@ -636,6 +637,11 @@ public void collect(int doc, long zeroBucket) throws IOException {
                         assert zeroBucket == 0L;
                         inner.collect(doc);
                     }
+
+                    @Override
+                    public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                        super.collect(stream, owningBucketOrd);
+                    }
                 };
             }
         }
@@ -665,6 +671,11 @@ public void collect(int doc, long bucket) throws IOException {
                     throw exc;
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 
@@ -724,6 +735,11 @@ public void collect(int doc, long zeroBucket) throws IOException {
                     subCollector.collect(doc, slot);
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FilterAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FilterAggregator.java
@@ -32,6 +32,7 @@
 package org.opensearch.search.aggregations.bucket.filter;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.opensearch.common.lucene.Lucene;
@@ -82,6 +83,11 @@ public void collect(int doc, long bucket) throws IOException {
                     collectBucket(sub, doc, bucket);
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FiltersAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FiltersAggregator.java
@@ -33,6 +33,7 @@
 package org.opensearch.search.aggregations.bucket.filter;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.opensearch.common.lucene.Lucene;
@@ -191,6 +192,11 @@ public void collect(int doc, long bucket) throws IOException {
                     collectBucket(sub, doc, bucketOrd(bucket, bits.length));
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java
@@ -41,6 +41,7 @@
 import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.index.Terms;
 import org.apache.lucene.index.TermsEnum;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BytesRef;
@@ -307,6 +308,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                         collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);
                     }
                 }
+
+                @Override
+                public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                    super.collect(stream, owningBucketOrd);
+                }
             });
         }
         return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) {
@@ -324,6 +330,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                     collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         });
     }
 
@@ -552,6 +563,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                         long docCount = docCountProvider.getDocCount(doc);
                         segmentDocCounts.increment(ord + 1, docCount);
                     }
+
+                    @Override
+                    public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                        super.collect(stream, owningBucketOrd);
+                    }
                 });
             }
             segmentsWithMultiValuedOrds++;
@@ -569,6 +585,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                         segmentDocCounts.increment(segmentOrd + 1, docCount);
                     }
                 }
+
+                @Override
+                public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                    super.collect(stream, owningBucketOrd);
+                }
             });
         }
 
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java
@@ -33,6 +33,7 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.DocIdStream;
 import org.opensearch.common.lease.Releasables;
 import org.opensearch.common.util.LongHash;
 import org.opensearch.common.util.SetBackedScalingCuckooFilter;
@@ -116,6 +117,11 @@ public void collect(int docId, long owningBucketOrd) throws IOException {
                     }
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java
@@ -35,6 +35,7 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.NumericUtils;
 import org.apache.lucene.util.PriorityQueue;
@@ -157,6 +158,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                     }
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         });
     }
 
@@ -715,6 +721,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                     subsetSizes = context.bigArrays().grow(subsetSizes, owningBucketOrd + 1);
                     subsetSizes.increment(owningBucketOrd, 1);
                 }
+
+                @Override
+                public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                    super.collect(stream, owningBucketOrd);
+                }
             };
         }
 
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregator.java
@@ -44,6 +44,7 @@
 import org.apache.lucene.search.DisiPriorityQueue;
 import org.apache.lucene.search.DisiWrapper;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.TermQuery;
@@ -593,6 +594,11 @@ public void collect(int doc, long bucketOrd) throws IOException {
             }
         }
 
+        @Override
+        public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+            super.collect(stream, owningBucketOrd);
+        }
+
         @Override
         public void postCollect() throws IOException {
             try (BitArray allVisitedOrds = new BitArray(maxOrd, bigArrays)) {

Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,7 @@`
`42`	`42`	`import org.apache.lucene.search.CollectionTerminatedException;`
`43`	`43`	`import org.apache.lucene.search.DocIdSet;`
`44`	`44`	`import org.apache.lucene.search.DocIdSetIterator;`
	`45`	`+import org.apache.lucene.search.DocIdStream;`
`45`	`46`	`import org.apache.lucene.search.FieldComparator;`
`46`	`47`	`import org.apache.lucene.search.FieldDoc;`
`47`	`48`	`import org.apache.lucene.search.LeafFieldComparator;`
`@@ -636,6 +637,11 @@ public void collect(int doc, long zeroBucket) throws IOException {`
`636`	`637`	`assert zeroBucket == 0L;`
`637`	`638`	`inner.collect(doc);`
`638`	`639`	`}`
	`640`	`+`
	`641`	`+ @Override`
	`642`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`643`	`+ super.collect(stream, owningBucketOrd);`
	`644`	`+ }`
`639`	`645`	`};`
`640`	`646`	`}`
`641`	`647`	`}`
`@@ -665,6 +671,11 @@ public void collect(int doc, long bucket) throws IOException {`
`665`	`671`	`throw exc;`
`666`	`672`	`}`
`667`	`673`	`}`
	`674`	`+`
	`675`	`+ @Override`
	`676`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`677`	`+ super.collect(stream, owningBucketOrd);`
	`678`	`+ }`
`668`	`679`	`};`
`669`	`680`	`}`
`670`	`681`
`@@ -724,6 +735,11 @@ public void collect(int doc, long zeroBucket) throws IOException {`
`724`	`735`	`subCollector.collect(doc, slot);`
`725`	`736`	`}`
`726`	`737`	`}`
	`738`	`+`
	`739`	`+ @Override`
	`740`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`741`	`+ super.collect(stream, owningBucketOrd);`
	`742`	`+ }`
`727`	`743`	`};`
`728`	`744`	`}`
`729`	`745`
Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@`
`32`	`32`	`package org.opensearch.search.aggregations.bucket.filter;`
`33`	`33`
`34`	`34`	`import org.apache.lucene.index.LeafReaderContext;`
	`35`	`+import org.apache.lucene.search.DocIdStream;`
`35`	`36`	`import org.apache.lucene.search.Weight;`
`36`	`37`	`import org.apache.lucene.util.Bits;`
`37`	`38`	`import org.opensearch.common.lucene.Lucene;`
`@@ -82,6 +83,11 @@ public void collect(int doc, long bucket) throws IOException {`
`82`	`83`	`collectBucket(sub, doc, bucket);`
`83`	`84`	`}`
`84`	`85`	`}`
	`86`	`+`
	`87`	`+ @Override`
	`88`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`89`	`+ super.collect(stream, owningBucketOrd);`
	`90`	`+ }`
`85`	`91`	`};`
`86`	`92`	`}`
`87`	`93`
Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@`
`33`	`33`	`package org.opensearch.search.aggregations.bucket.filter;`
`34`	`34`
`35`	`35`	`import org.apache.lucene.index.LeafReaderContext;`
	`36`	`+import org.apache.lucene.search.DocIdStream;`
`36`	`37`	`import org.apache.lucene.search.Weight;`
`37`	`38`	`import org.apache.lucene.util.Bits;`
`38`	`39`	`import org.opensearch.common.lucene.Lucene;`
`@@ -191,6 +192,11 @@ public void collect(int doc, long bucket) throws IOException {`
`191`	`192`	`collectBucket(sub, doc, bucketOrd(bucket, bits.length));`
`192`	`193`	`}`
`193`	`194`	`}`
	`195`	`+`
	`196`	`+ @Override`
	`197`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`198`	`+ super.collect(stream, owningBucketOrd);`
	`199`	`+ }`
`194`	`200`	`};`
`195`	`201`	`}`
`196`	`202`
Original file line number	Diff line number	Diff line change
`@@ -41,6 +41,7 @@`
`41`	`41`	`import org.apache.lucene.index.SortedSetDocValues;`
`42`	`42`	`import org.apache.lucene.index.Terms;`
`43`	`43`	`import org.apache.lucene.index.TermsEnum;`
	`44`	`+import org.apache.lucene.search.DocIdStream;`
`44`	`45`	`import org.apache.lucene.search.Weight;`
`45`	`46`	`import org.apache.lucene.util.ArrayUtil;`
`46`	`47`	`import org.apache.lucene.util.BytesRef;`
`@@ -307,6 +308,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {`
`307`	`308`	`collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);`
`308`	`309`	`}`
`309`	`310`	`}`
	`311`	`+`
	`312`	`+ @Override`
	`313`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`314`	`+ super.collect(stream, owningBucketOrd);`
	`315`	`+ }`
`310`	`316`	`});`
`311`	`317`	`}`
`312`	`318`	`return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) {`
`@@ -324,6 +330,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {`
`324`	`330`	`collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub);`
`325`	`331`	`}`
`326`	`332`	`}`
	`333`	`+`
	`334`	`+ @Override`
	`335`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`336`	`+ super.collect(stream, owningBucketOrd);`
	`337`	`+ }`
`327`	`338`	`});`
`328`	`339`	`}`
`329`	`340`
`@@ -552,6 +563,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {`
`552`	`563`	`long docCount = docCountProvider.getDocCount(doc);`
`553`	`564`	`segmentDocCounts.increment(ord + 1, docCount);`
`554`	`565`	`}`
	`566`	`+`
	`567`	`+ @Override`
	`568`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`569`	`+ super.collect(stream, owningBucketOrd);`
	`570`	`+ }`
`555`	`571`	`});`
`556`	`572`	`}`
`557`	`573`	`segmentsWithMultiValuedOrds++;`
`@@ -569,6 +585,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {`
`569`	`585`	`segmentDocCounts.increment(segmentOrd + 1, docCount);`
`570`	`586`	`}`
`571`	`587`	`}`
	`588`	`+`
	`589`	`+ @Override`
	`590`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`591`	`+ super.collect(stream, owningBucketOrd);`
	`592`	`+ }`
`572`	`593`	`});`
`573`	`594`	`}`
`574`	`595`
Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@`
`35`	`35`	`import org.apache.lucene.index.IndexReader;`
`36`	`36`	`import org.apache.lucene.index.LeafReaderContext;`
`37`	`37`	`import org.apache.lucene.index.SortedNumericDocValues;`
	`38`	`+import org.apache.lucene.search.DocIdStream;`
`38`	`39`	`import org.apache.lucene.search.ScoreMode;`
`39`	`40`	`import org.apache.lucene.util.NumericUtils;`
`40`	`41`	`import org.apache.lucene.util.PriorityQueue;`
`@@ -157,6 +158,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {`
`157`	`158`	`}`
`158`	`159`	`}`
`159`	`160`	`}`
	`161`	`+`
	`162`	`+ @Override`
	`163`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`164`	`+ super.collect(stream, owningBucketOrd);`
	`165`	`+ }`
`160`	`166`	`});`
`161`	`167`	`}`
`162`	`168`
`@@ -715,6 +721,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {`
`715`	`721`	`subsetSizes = context.bigArrays().grow(subsetSizes, owningBucketOrd + 1);`
`716`	`722`	`subsetSizes.increment(owningBucketOrd, 1);`
`717`	`723`	`}`
	`724`	`+`
	`725`	`+ @Override`
	`726`	`+ public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {`
	`727`	`+ super.collect(stream, owningBucketOrd);`
	`728`	`+ }`
`718`	`729`	`};`
`719`	`730`	`}`
`720`	`731`