Allow composite aggregation to be a sub agg of the random sampler (#8…

…5064) composite aggregations are supported within a sampling context. Consequently, composite aggs need to allow themselves to be a sub-agg of a random sampler aggregation.
elastic · Mar 18, 2022 · c0db855 · c0db855
1 parent 3a5a4ea
commit c0db855
Show file tree

Hide file tree

Showing 3 changed files with 355 additions and 118 deletions.
diff --git a/...a/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java b/...a/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java
@@ -18,6 +18,7 @@
 import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.bucket.filter.FilterAggregatorFactory;
 import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregatorFactory;
+import org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplerAggregatorFactory;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
@@ -176,11 +177,13 @@ public BucketCardinality bucketCardinality() {
     private AggregatorFactory validateParentAggregations(AggregatorFactory factory) {
         if (factory == null) {
             return null;
-        } else if (factory instanceof NestedAggregatorFactory || factory instanceof FilterAggregatorFactory) {
-            return validateParentAggregations(factory.getParent());
-        } else {
-            return factory;
-        }
+        } else if (factory instanceof NestedAggregatorFactory
+            || factory instanceof FilterAggregatorFactory
+            || factory instanceof RandomSamplerAggregatorFactory) {
+                return validateParentAggregations(factory.getParent());
+            } else {
+                return factory;
+            }
     }
 
     private static void validateSources(List<CompositeValuesSourceBuilder<?>> sources) {

diff --git a/.../org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplerAggregator.java b/.../org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplerAggregator.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.search.aggregations.bucket.sampler.random;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.Weight;
@@ -106,12 +107,18 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
         }
         final DocIdSetIterator docIt = scorer.iterator();
         final Bits liveDocs = ctx.reader().getLiveDocs();
-        // Iterate every document provided by the scorer iterator
-        for (int docId = docIt.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = docIt.nextDoc()) {
-            // If liveDocs is null, that means that every doc is a live doc, no need to check if it has been deleted or not
-            if (liveDocs == null || liveDocs.get(docIt.docID())) {
-                collectBucket(sub, docIt.docID(), 0);
+        try {
+            // Iterate every document provided by the scorer iterator
+            for (int docId = docIt.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = docIt.nextDoc()) {
+                // If liveDocs is null, that means that every doc is a live doc, no need to check if it has been deleted or not
+                if (liveDocs == null || liveDocs.get(docIt.docID())) {
+                    collectBucket(sub, docIt.docID(), 0);
+                }
             }
+            // This collector could throw `CollectionTerminatedException` if the last leaf collector has stopped collecting
+            // So, catch here and indicate no-op
+        } catch (CollectionTerminatedException e) {
+            return LeafBucketCollector.NO_OP_COLLECTOR;
         }
         // Since we have done our own collection, there is nothing for the leaf collector to do
         return LeafBucketCollector.NO_OP_COLLECTOR;