Cache the score of the parent document in the nested agg (#36019)

The nested agg can defer the collection of children if it is nested
under another aggregation. In such case accessing the score in the children
aggregation throws an error because the scorer has already advanced to the next
parent. This change fixes this error by caching the score of the parent in the
nested aggregation. Children aggregations that work on nested documents will be
able to access the _score. Also note that the _score in this case is always the
parent's score, there is no way to retrieve the score of a nested docs in aggregations.

Closes #35985
Closes #34555

Author: jimczi

server/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregator.java

@@ -140,7 +140,9 @@ class BufferingNestedLeafBucketCollector extends LeafBucketCollectorBase {
         final DocIdSetIterator childDocs;
         final LongArrayList bucketBuffer = new LongArrayList();
 
+        Scorer scorer;
         int currentParentDoc = -1;
+        final CachedScorer cachedScorer = new CachedScorer();
 
         BufferingNestedLeafBucketCollector(LeafBucketCollector sub, BitSet parentDocs, DocIdSetIterator childDocs) {
             super(sub, null);
@@ -149,6 +151,12 @@ class BufferingNestedLeafBucketCollector extends LeafBucketCollectorBase {
             this.childDocs = childDocs;
         }
 
+        @Override
+        public void setScorer(Scorer scorer) throws IOException {
+            this.scorer = scorer;
+            super.setScorer(cachedScorer);
+        }
+
         @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
@@ -159,7 +167,12 @@ public void collect(int parentDoc, long bucket) throws IOException {
 
             if (currentParentDoc != parentDoc) {
                 processBufferedChildBuckets();
+                if (needsScores()) {
+                    // cache the score of the current parent
+                    cachedScorer.score = scorer.score();
+                }
                 currentParentDoc = parentDoc;
+
             }
             bucketBuffer.add(bucket);
         }
@@ -177,6 +190,7 @@ void processBufferedChildBuckets() throws IOException {
             }
 
             for (; childDocId < currentParentDoc; childDocId = childDocs.nextDoc()) {
+                cachedScorer.doc = childDocId;
                 final long[] buffer = bucketBuffer.buffer;
                 final int size = bucketBuffer.size();
                 for (int i = 0; i < size; i++) {
@@ -185,6 +199,26 @@ void processBufferedChildBuckets() throws IOException {
             }
             bucketBuffer.clear();
         }
+    }
+
+    private static class CachedScorer extends Scorer {
+        int doc;
+        float score;
+
+        private CachedScorer() { super(null); }
+
+        @Override
+        public DocIdSetIterator iterator() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public final float score() { return score; }
+
+        @Override
+        public int docID() {
+            return doc;
+        }
 
     }
 

server/src/test/java/org/elasticsearch/search/aggregations/bucket/terms/TermsAggregatorTests.java

@@ -47,6 +47,7 @@
 import org.elasticsearch.index.mapper.SeqNoFieldMapper;
 import org.elasticsearch.index.mapper.TypeFieldMapper;
 import org.elasticsearch.index.mapper.UidFieldMapper;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
 import org.elasticsearch.search.SearchHit;
@@ -60,6 +61,7 @@
 import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
 import org.elasticsearch.search.aggregations.bucket.filter.Filter;
 import org.elasticsearch.search.aggregations.bucket.filter.FilterAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.filter.InternalFilter;
 import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.global.InternalGlobal;
 import org.elasticsearch.search.aggregations.bucket.nested.InternalNested;
@@ -1042,21 +1044,23 @@ public void testWithNestedAggregations() throws IOException {
                         fieldType.setHasDocValues(true);
                         fieldType.setName("nested_value");
                         try (IndexReader indexReader = wrap(DirectoryReader.open(directory))) {
-                            InternalNested result = search(newSearcher(indexReader, false, true),
-                                // match root document only
-                                new DocValuesFieldExistsQuery(PRIMARY_TERM_NAME), nested, fieldType);
-                            InternalMultiBucketAggregation<?, ?> terms = result.getAggregations().get("terms");
-                            assertThat(terms.getBuckets().size(), equalTo(9));
-                            int ptr = 9;
-                            for (MultiBucketsAggregation.Bucket bucket : terms.getBuckets()) {
-                                InternalTopHits topHits = bucket.getAggregations().get("top_hits");
-                                assertThat(topHits.getHits().totalHits, equalTo((long) ptr));
-                                if (withScore) {
-                                    assertThat(topHits.getHits().getMaxScore(), equalTo(1f));
-                                } else {
-                                    assertThat(topHits.getHits().getMaxScore(), equalTo(Float.NaN));
-                                }
-                                --ptr;
+                            {
+                                InternalNested result = search(newSearcher(indexReader, false, true),
+                                    // match root document only
+                                    new DocValuesFieldExistsQuery(PRIMARY_TERM_NAME), nested, fieldType);
+                                InternalMultiBucketAggregation<?, ?> terms = result.getAggregations().get("terms");
+                                assertNestedTopHitsScore(terms, withScore);
+                            }
+
+                            {
+                                FilterAggregationBuilder filter = new FilterAggregationBuilder("filter", new MatchAllQueryBuilder())
+                                    .subAggregation(nested);
+                                InternalFilter result = search(newSearcher(indexReader, false, true),
+                                    // match root document only
+                                    new DocValuesFieldExistsQuery(PRIMARY_TERM_NAME), filter, fieldType);
+                                InternalNested nestedResult = result.getAggregations().get("nested");
+                                InternalMultiBucketAggregation<?, ?> terms = nestedResult.getAggregations().get("terms");
+                                assertNestedTopHitsScore(terms, withScore);
                             }
                         }
                     }
@@ -1065,6 +1069,21 @@ public void testWithNestedAggregations() throws IOException {
         }
     }
 
+    private void assertNestedTopHitsScore(InternalMultiBucketAggregation<?, ?> terms, boolean withScore) {
+        assertThat(terms.getBuckets().size(), equalTo(9));
+        int ptr = 9;
+        for (MultiBucketsAggregation.Bucket bucket : terms.getBuckets()) {
+            InternalTopHits topHits = bucket.getAggregations().get("top_hits");
+            assertThat(topHits.getHits().totalHits, equalTo((long) ptr));
+            if (withScore) {
+                assertThat(topHits.getHits().getMaxScore(), equalTo(1f));
+            } else {
+                assertThat(topHits.getHits().getMaxScore(), equalTo(Float.NaN));
+            }
+            --ptr;
+        }
+    }
+
     private final SeqNoFieldMapper.SequenceIDFields sequenceIDFields = SeqNoFieldMapper.SequenceIDFields.emptySeqID();
     private List<Document> generateDocsWithNested(String id, int value, int[] nestedValues) {
         List<Document> documents = new ArrayList<>();