Asserting single segment in StreamStringTermsAggregator

Signed-off-by: Mikhail Khludnev <[email protected]>

Author: mkhludnev

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

@@ -47,6 +47,7 @@ public class StreamStringTermsAggregator extends AbstractStringTermsAggregator {
     protected int segmentsWithSingleValuedOrds = 0;
     protected int segmentsWithMultiValuedOrds = 0;
     protected final ResultStrategy<?, ?, ?> resultStrategy;
+    private boolean madeLeafOnce = false;
 
     public StreamStringTermsAggregator(
         String name,
@@ -72,6 +73,7 @@ public void doReset() {
         super.doReset();
         valueCount = 0;
         sortedDocValuesPerBatch = null;
+        this.madeLeafOnce = false;
     }
 
     @Override
@@ -91,6 +93,11 @@ public InternalAggregation buildEmptyAggregation() {
 
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
+        if (this.madeLeafOnce) {
+            throw new IllegalStateException("pardon you might already collected " + docCounts);
+        } else {
+            this.madeLeafOnce = true;
+        }
         this.sortedDocValuesPerBatch = valuesSource.ordinalsValues(ctx);
         this.valueCount = sortedDocValuesPerBatch.getValueCount(); // for streaming case, the value count is reset to per batch
         // cardinality

server/src/test/java/org/opensearch/search/aggregations/bucket/terms/StreamStringTermsAggregatorTests.java

@@ -344,6 +344,9 @@ public void testBuildAggregationsBatchReset() throws Exception {
                 Document document = new Document();
                 document.add(new SortedSetDocValuesField("field", new BytesRef("test")));
                 indexWriter.addDocument(document);
+                document = new Document();
+                document.add(new SortedSetDocValuesField("field", new BytesRef("best")));
+                indexWriter.addDocument(document);
 
                 try (IndexReader indexReader = maybeWrapReaderEs(indexWriter.getReader())) {
                     IndexSearcher indexSearcher = newIndexSearcher(indexReader);
@@ -369,7 +372,7 @@ public void testBuildAggregationsBatchReset() throws Exception {
                     aggregator.postCollection();
 
                     StringTerms firstResult = (StringTerms) aggregator.buildAggregations(new long[] { 0 })[0];
-                    assertThat(firstResult.getBuckets().size(), equalTo(1));
+                    assertThat(firstResult.getBuckets().size(), equalTo(2));
 
                     aggregator.doReset();
 
@@ -379,7 +382,7 @@ public void testBuildAggregationsBatchReset() throws Exception {
                     aggregator.postCollection();
 
                     StringTerms secondResult = (StringTerms) aggregator.buildAggregations(new long[] { 0 })[0];
-                    assertThat(secondResult.getBuckets().size(), equalTo(1));
+                    assertThat(secondResult.getBuckets().size(), equalTo(2));
                     assertThat(secondResult.getBuckets().get(0).getDocCount(), equalTo(1L));
                 }
             }