org.elasticsearch.search.aggregations.bucket.sampler.DiversifiedMapSamplerAggregator Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch Show documentation
Show all versions of elasticsearch Show documentation
Elasticsearch subproject :server
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0 and the Server Side Public License, v 1; you may not use this file except
* in compliance with, at your election, the Elastic License 2.0 or the Server
* Side Public License, v 1.
*/
package org.elasticsearch.search.aggregations.bucket.sampler;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.misc.search.DiversifiedTopDocsCollector;
import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.util.BytesRefHash;
import org.elasticsearch.core.Releasables;
import org.elasticsearch.index.fielddata.AbstractNumericDocValues;
import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
import org.elasticsearch.search.aggregations.Aggregator;
import org.elasticsearch.search.aggregations.AggregatorFactories;
import org.elasticsearch.search.aggregations.bucket.DeferringBucketCollector;
import org.elasticsearch.search.aggregations.support.AggregationContext;
import org.elasticsearch.search.aggregations.support.ValuesSource;
import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
import java.io.IOException;
import java.util.Map;
import java.util.function.Consumer;
import static org.apache.lucene.misc.search.DiversifiedTopDocsCollector.ScoreDocKey;
public class DiversifiedMapSamplerAggregator extends SamplerAggregator {
private ValuesSource valuesSource;
private int maxDocsPerValue;
private BytesRefHash bucketOrds;
DiversifiedMapSamplerAggregator(
String name,
int shardSize,
AggregatorFactories factories,
AggregationContext context,
Aggregator parent,
Map metadata,
ValuesSourceConfig valuesSourceConfig,
int maxDocsPerValue
) throws IOException {
super(name, shardSize, factories, context, parent, metadata);
assert valuesSourceConfig.hasValues();
this.valuesSource = valuesSourceConfig.getValuesSource();
this.maxDocsPerValue = maxDocsPerValue;
// Need to use super class shardSize since it is limited to maxDoc
bucketOrds = new BytesRefHash(this.shardSize, context.bigArrays());
}
@Override
protected void doClose() {
Releasables.close(bucketOrds);
super.doClose();
}
@Override
public DeferringBucketCollector buildDeferringCollector() {
bdd = new DiverseDocsDeferringCollector(this::addRequestCircuitBreakerBytes);
return bdd;
}
/**
* A {@link DeferringBucketCollector} that identifies top scoring documents
* but de-duped by a key then passes only these on to nested collectors.
* This implementation is only for use with a single bucket aggregation.
*/
class DiverseDocsDeferringCollector extends BestDocsDeferringCollector {
DiverseDocsDeferringCollector(Consumer circuitBreakerConsumer) {
super(shardSize, bigArrays(), circuitBreakerConsumer);
}
@Override
protected TopDocsCollector createTopDocsCollector(int size) {
// Make sure we do not allow size > maxDoc, to prevent accidental OOM
int minMaxDocsPerValue = Math.min(maxDocsPerValue, searcher().getIndexReader().maxDoc());
return new ValuesDiversifiedTopDocsCollector(size, minMaxDocsPerValue);
}
@Override
protected long getPriorityQueueSlotSize() {
return SCOREDOCKEY_SIZE;
}
// This class extends the DiversifiedTopDocsCollector and provides
// a lookup from elasticsearch's ValuesSource
class ValuesDiversifiedTopDocsCollector extends DiversifiedTopDocsCollector {
private SortedBinaryDocValues values;
ValuesDiversifiedTopDocsCollector(int numHits, int maxHitsPerKey) {
super(numHits, maxHitsPerKey);
}
@Override
protected NumericDocValues getKeys(LeafReaderContext context) {
try {
values = valuesSource.bytesValues(context);
} catch (IOException e) {
throw new ElasticsearchException("Error reading values", e);
}
return new AbstractNumericDocValues() {
private int docID = -1;
@Override
public boolean advanceExact(int target) throws IOException {
docID = target;
if (values.advanceExact(target)) {
if (values.docValueCount() > 1) {
throw new IllegalArgumentException("Sample diversifying key must be a single valued-field");
}
return true;
} else {
return false;
}
}
@Override
public int docID() {
return docID;
}
@Override
public long longValue() throws IOException {
final BytesRef bytes = values.nextValue();
long bucketOrdinal = bucketOrds.add(bytes);
if (bucketOrdinal < 0) { // already seen
bucketOrdinal = -1 - bucketOrdinal;
}
return bucketOrdinal;
}
};
}
}
}
}