Skip to content

Commit

Permalink
Cache the score of the parent document in the nested agg (elastic#36019)
Browse files Browse the repository at this point in the history 
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 elastic#35985
Closes elastic#34555
  • Loading branch information
@jimczi
jimczi committed Nov 29, 2018
1 parent 06adb19 commit 8d58334
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 15 deletions.
34 changes: 34 additions & 0 deletions ...r/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregator.java
View file
Expand Up @@ -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);
Expand All @@ -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
Expand All @@ -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);
}
Expand All @@ -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++) {
Expand All @@ -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;
}

}

Expand Down
49 changes: 34 additions & 15 deletions ...rc/test/java/org/elasticsearch/search/aggregations/bucket/terms/TermsAggregatorTests.java
View file
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}
}
}
Expand All @@ -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<>();
Expand Down

0 comments on commit 8d58334

Please sign in to comment.