aggs: Allow aggregation sorting via nested aggregation.

martijnvg · martijnvg · commit 872004a57100 · 2017-09-20T07:46:03.000+02:00
The nested aggregator now buffers all bucket ords per parent document and emits all bucket ords for a parent document's nested document once. This way the nested documents document DocIdSetIterator gets used once per bucket instead of wrapping the nested aggregator inside a multi bucket aggregator, which was the current solution upto now. This allows sorting by buckets under a nested bucket. Closes #16838
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregator.java
@@ -18,6 +18,7 @@
  */
 package org.elasticsearch.search.aggregations.bucket.nested;
 
+import com.carrotsearch.hppc.LongArrayList;
 import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
@@ -51,14 +52,19 @@ class NestedAggregator extends BucketsAggregator implements SingleBucketAggregat
 
     private final BitSetProducer parentFilter;
     private final Query childFilter;
+    private final boolean collectsFromSingleBucket;
+
+    private BufferingNestedLeafBucketCollector bufferingNestedLeafBucketCollector;
 
     NestedAggregator(String name, AggregatorFactories factories, ObjectMapper parentObjectMapper, ObjectMapper childObjectMapper,
-            SearchContext context, Aggregator parentAggregator,
-                            List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
+                     SearchContext context, Aggregator parentAggregator,
+                     List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData,
+                     boolean collectsFromSingleBucket) throws IOException {
         super(name, factories, context, parentAggregator, pipelineAggregators, metaData);
         Query parentFilter = parentObjectMapper != null ? parentObjectMapper.nestedTypeFilter() : Queries.newNonNestedFilter();
         this.parentFilter = context.bitsetFilterCache().getBitSetProducer(parentFilter);
         this.childFilter = childObjectMapper.nestedTypeFilter();
+        this.collectsFromSingleBucket = collectsFromSingleBucket;
     }
 
     @Override
@@ -71,26 +77,38 @@ public LeafBucketCollector getLeafCollector(final LeafReaderContext ctx, final L
 
         final BitSet parentDocs = parentFilter.getBitSet(ctx);
         final DocIdSetIterator childDocs = childDocsScorer != null ? childDocsScorer.iterator() : null;
-        return new LeafBucketCollectorBase(sub, null) {
-            @Override
-            public void collect(int parentDoc, long bucket) throws IOException {
-                // if parentDoc is 0 then this means that this parent doesn't have child docs (b/c these appear always before the parent
-                // doc), so we can skip:
-                if (parentDoc == 0 || parentDocs == null || childDocs == null) {
-                    return;
-                }
+        if (collectsFromSingleBucket) {
+            return new LeafBucketCollectorBase(sub, null) {
+                @Override
+                public void collect(int parentDoc, long bucket) throws IOException {
+                    // if parentDoc is 0 then this means that this parent doesn't have child docs (b/c these appear always before the parent
+                    // doc), so we can skip:
+                    if (parentDoc == 0 || parentDocs == null || childDocs == null) {
+                        return;
+                    }
 
-                final int prevParentDoc = parentDocs.prevSetBit(parentDoc - 1);
-                int childDocId = childDocs.docID();
-                if (childDocId <= prevParentDoc) {
-                    childDocId = childDocs.advance(prevParentDoc + 1);
-                }
+                    final int prevParentDoc = parentDocs.prevSetBit(parentDoc - 1);
+                    int childDocId = childDocs.docID();
+                    if (childDocId <= prevParentDoc) {
+                        childDocId = childDocs.advance(prevParentDoc + 1);
+                    }
 
-                for (; childDocId < parentDoc; childDocId = childDocs.nextDoc()) {
-                    collectBucket(sub, childDocId, bucket);
+                    for (; childDocId < parentDoc; childDocId = childDocs.nextDoc()) {
+                        collectBucket(sub, childDocId, bucket);
+                    }
                 }
-            }
-        };
+            };
+        } else {
+            doPostCollection();
+            return bufferingNestedLeafBucketCollector = new BufferingNestedLeafBucketCollector(sub, parentDocs, childDocs);
+        }
+    }
+
+    @Override
+    protected void doPostCollection() throws IOException {
+        if (bufferingNestedLeafBucketCollector != null) {
+            bufferingNestedLeafBucketCollector.postCollect();
+        }
     }
 
     @Override
@@ -104,4 +122,63 @@ public InternalAggregation buildEmptyAggregation() {
         return new InternalNested(name, 0, buildEmptySubAggregations(), pipelineAggregators(), metaData());
     }
 
+    class BufferingNestedLeafBucketCollector extends LeafBucketCollectorBase {
+
+        final BitSet parentDocs;
+        final LeafBucketCollector sub;
+        final DocIdSetIterator childDocs;
+        final LongArrayList bucketBuffer = new LongArrayList();
+
+        int currentParentDoc = -1;
+
+        BufferingNestedLeafBucketCollector(LeafBucketCollector sub, BitSet parentDocs, DocIdSetIterator childDocs) {
+            super(sub, null);
+            this.sub = sub;
+            this.parentDocs = parentDocs;
+            this.childDocs = childDocs;
+        }
+
+        @Override
+        public void collect(int parentDoc, long bucket) throws IOException {
+            // if parentDoc is 0 then this means that this parent doesn't have child docs (b/c these appear always before the parent
+            // doc), so we can skip:
+            if (parentDoc == 0 || parentDocs == null || childDocs == null) {
+                return;
+            }
+
+            if (currentParentDoc != parentDoc) {
+                processChildBuckets(currentParentDoc, bucketBuffer);
+                currentParentDoc = parentDoc;
+            }
+            bucketBuffer.add(bucket);
+        }
+
+        void processChildBuckets(int parentDoc, LongArrayList buckets) throws IOException {
+            if (bucketBuffer.isEmpty()) {
+                return;
+            }
+
+
+            final int prevParentDoc = parentDocs.prevSetBit(parentDoc - 1);
+            int childDocId = childDocs.docID();
+            if (childDocId <= prevParentDoc) {
+                childDocId = childDocs.advance(prevParentDoc + 1);
+            }
+
+            for (; childDocId < parentDoc; childDocId = childDocs.nextDoc()) {
+                final long[] buffer = buckets.buffer;
+                final int size = buckets.size();
+                for (int i = 0; i < size; i++) {
+                    collectBucket(sub, childDocId, buffer[i]);
+                }
+            }
+            bucketBuffer.clear();
+        }
+
+        void postCollect() throws IOException {
+            processChildBuckets(currentParentDoc, bucketBuffer);
+        }
+
+    }
+
 }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregatorFactory.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregatorFactory.java
@@ -48,13 +48,11 @@ class NestedAggregatorFactory extends AggregatorFactory<NestedAggregatorFactory>
     @Override
     public Aggregator createInternal(Aggregator parent, boolean collectsFromSingleBucket, List<PipelineAggregator> pipelineAggregators,
             Map<String, Object> metaData) throws IOException {
-        if (collectsFromSingleBucket == false) {
-            return asMultiBucketAggregator(this, context, parent);
-        }
         if (childObjectMapper == null) {
             return new Unmapped(name, context, parent, pipelineAggregators, metaData);
         }
-        return new NestedAggregator(name, factories, parentObjectMapper, childObjectMapper, context, parent, pipelineAggregators, metaData);
+        return new NestedAggregator(name, factories, parentObjectMapper, childObjectMapper, context, parent,
+            pipelineAggregators, metaData, collectsFromSingleBucket);
     }
 
     private static final class Unmapped extends NonCollectingAggregator {
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregator.java
@@ -91,10 +91,6 @@ public boolean needsScores() {
     public LeafBucketCollector getLeafCollector(final LeafReaderContext ctx,
             final LeafBucketCollector sub) throws IOException {
 
-        for (LongObjectPagedHashMap.Cursor<TopDocsAndLeafCollector> cursor : topDocsCollectors) {
-            cursor.value.leafCollector = cursor.value.topLevelCollector.getLeafCollector(ctx);
-        }
-
         return new LeafBucketCollectorBase(sub, null) {
 
             Scorer scorer;
@@ -103,6 +99,11 @@ public LeafBucketCollector getLeafCollector(final LeafReaderContext ctx,
             public void setScorer(Scorer scorer) throws IOException {
                 this.scorer = scorer;
                 for (LongObjectPagedHashMap.Cursor<TopDocsAndLeafCollector> cursor : topDocsCollectors) {
+                    // Instantiate the leaf collector not in the getLeafCollector(...) method or in the constructor of this
+                    // anonymous class. Otherwise in the case this leaf bucket collector gets invoked with post collection
+                    // then we already have moved on to the next reader and then we may encounter assertion errors or
+                    // incorrect results.
+                    cursor.value.leafCollector = cursor.value.topLevelCollector.getLeafCollector(ctx);
                     cursor.value.leafCollector.setScorer(scorer);
                 }
                 super.setScorer(scorer);
diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregatorTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregatorTests.java
