org.elasticsearch.search.aggregations.bucket.sampler.DiversifiedAggregationBuilder 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.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregatorFactories.Builder;
import org.elasticsearch.search.aggregations.AggregatorFactory;
import org.elasticsearch.search.aggregations.support.AggregationContext;
import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
import org.elasticsearch.search.aggregations.support.ValuesSourceType;
import org.elasticsearch.xcontent.ObjectParser;
import org.elasticsearch.xcontent.XContentBuilder;
import java.io.IOException;
import java.util.Map;
import java.util.Objects;
public class DiversifiedAggregationBuilder extends ValuesSourceAggregationBuilder {
public static final String NAME = "diversified_sampler";
public static final ValuesSourceRegistry.RegistryKey REGISTRY_KEY =
new ValuesSourceRegistry.RegistryKey<>(NAME, DiversifiedAggregatorSupplier.class);
public static final int MAX_DOCS_PER_VALUE_DEFAULT = 1;
public static final ObjectParser PARSER = ObjectParser.fromBuilder(
NAME,
DiversifiedAggregationBuilder::new
);
static {
ValuesSourceAggregationBuilder.declareFields(PARSER, true, false, false);
PARSER.declareInt(DiversifiedAggregationBuilder::shardSize, SamplerAggregator.SHARD_SIZE_FIELD);
PARSER.declareInt(DiversifiedAggregationBuilder::maxDocsPerValue, SamplerAggregator.MAX_DOCS_PER_VALUE_FIELD);
PARSER.declareString(DiversifiedAggregationBuilder::executionHint, SamplerAggregator.EXECUTION_HINT_FIELD);
}
public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
DiversifiedAggregatorFactory.registerAggregators(builder);
}
private int shardSize = SamplerAggregationBuilder.DEFAULT_SHARD_SAMPLE_SIZE;
private int maxDocsPerValue = MAX_DOCS_PER_VALUE_DEFAULT;
private String executionHint = null;
public DiversifiedAggregationBuilder(String name) {
super(name);
}
protected DiversifiedAggregationBuilder(DiversifiedAggregationBuilder clone, Builder factoriesBuilder, Map metadata) {
super(clone, factoriesBuilder, metadata);
this.shardSize = clone.shardSize;
this.maxDocsPerValue = clone.maxDocsPerValue;
this.executionHint = clone.executionHint;
}
@Override
protected ValuesSourceType defaultValueSourceType() {
return CoreValuesSourceType.KEYWORD;
}
@Override
protected AggregationBuilder shallowCopy(Builder factoriesBuilder, Map metadata) {
return new DiversifiedAggregationBuilder(this, factoriesBuilder, metadata);
}
/**
* Read from a stream.
*/
public DiversifiedAggregationBuilder(StreamInput in) throws IOException {
super(in);
shardSize = in.readVInt();
maxDocsPerValue = in.readVInt();
executionHint = in.readOptionalString();
}
@Override
protected void innerWriteTo(StreamOutput out) throws IOException {
out.writeVInt(shardSize);
out.writeVInt(maxDocsPerValue);
out.writeOptionalString(executionHint);
}
/**
* Set the max num docs to be returned from each shard.
*/
public DiversifiedAggregationBuilder shardSize(int shardSize) {
if (shardSize < 0) {
throw new IllegalArgumentException(
"[shardSize] must be greater than or equal to 0. Found [" + shardSize + "] in [" + name + "]"
);
}
this.shardSize = shardSize;
return this;
}
/**
* Get the max num docs to be returned from each shard.
*/
public int shardSize() {
return shardSize;
}
/**
* Set the max num docs to be returned per value.
*/
public DiversifiedAggregationBuilder maxDocsPerValue(int maxDocsPerValue) {
if (maxDocsPerValue < 0) {
throw new IllegalArgumentException(
"[maxDocsPerValue] must be greater than or equal to 0. Found [" + maxDocsPerValue + "] in [" + name + "]"
);
}
this.maxDocsPerValue = maxDocsPerValue;
return this;
}
/**
* Get the max num docs to be returned per value.
*/
public int maxDocsPerValue() {
return maxDocsPerValue;
}
/**
* Set the execution hint.
*/
public DiversifiedAggregationBuilder executionHint(String executionHint) {
this.executionHint = executionHint;
return this;
}
/**
* Get the execution hint.
*/
public String executionHint() {
return executionHint;
}
@Override
public BucketCardinality bucketCardinality() {
return BucketCardinality.ONE;
}
@Override
protected ValuesSourceAggregatorFactory innerBuild(
AggregationContext context,
ValuesSourceConfig config,
AggregatorFactory parent,
Builder subFactoriesBuilder
) throws IOException {
DiversifiedAggregatorSupplier aggregatorSupplier = context.getValuesSourceRegistry().getAggregator(REGISTRY_KEY, config);
return new DiversifiedAggregatorFactory(
name,
config,
shardSize,
maxDocsPerValue,
executionHint,
context,
parent,
subFactoriesBuilder,
metadata,
aggregatorSupplier
);
}
@Override
protected XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
builder.field(SamplerAggregator.SHARD_SIZE_FIELD.getPreferredName(), shardSize);
builder.field(SamplerAggregator.MAX_DOCS_PER_VALUE_FIELD.getPreferredName(), maxDocsPerValue);
if (executionHint != null) {
builder.field(SamplerAggregator.EXECUTION_HINT_FIELD.getPreferredName(), executionHint);
}
return builder;
}
@Override
public int hashCode() {
return Objects.hash(super.hashCode(), shardSize, maxDocsPerValue, executionHint);
}
@Override
public boolean equals(Object obj) {
if (this == obj) return true;
if (obj == null || getClass() != obj.getClass()) return false;
if (super.equals(obj) == false) return false;
DiversifiedAggregationBuilder other = (DiversifiedAggregationBuilder) obj;
return Objects.equals(shardSize, other.shardSize)
&& Objects.equals(maxDocsPerValue, other.maxDocsPerValue)
&& Objects.equals(executionHint, other.executionHint);
}
@Override
public String getType() {
return NAME;
}
@Override
protected ValuesSourceRegistry.RegistryKey getRegistryKey() {
return REGISTRY_KEY;
}
}