org.elasticsearch.search.aggregations.bucket.range.InternalBinaryRange 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.range;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.TransportVersions;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.aggregations.AggregationReduceContext;
import org.elasticsearch.search.aggregations.AggregatorReducer;
import org.elasticsearch.search.aggregations.InternalAggregation;
import org.elasticsearch.search.aggregations.InternalAggregations;
import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
import org.elasticsearch.search.aggregations.support.SamplingContext;
import org.elasticsearch.xcontent.XContentBuilder;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import static java.util.Collections.unmodifiableList;
/** A range aggregation for data that is encoded in doc values using a binary representation. */
public final class InternalBinaryRange extends InternalMultiBucketAggregation
implements
Range {
public static class Bucket extends InternalMultiBucketAggregation.InternalBucket implements Range.Bucket {
private final transient DocValueFormat format;
private final transient boolean keyed;
private final String key;
private final BytesRef from, to;
private final long docCount;
private final InternalAggregations aggregations;
public Bucket(
DocValueFormat format,
boolean keyed,
String key,
BytesRef from,
BytesRef to,
long docCount,
InternalAggregations aggregations
) {
this.format = format;
this.keyed = keyed;
this.key = key;
this.from = from;
this.to = to;
this.docCount = docCount;
this.aggregations = aggregations;
}
private static String generateKey(BytesRef from, BytesRef to, DocValueFormat format) {
return (from == null ? "*" : format.format(from)) + "-" + (to == null ? "*" : format.format(to));
}
private static Bucket createFromStream(StreamInput in, DocValueFormat format, boolean keyed) throws IOException {
// NOTE: the key is required in version == 8.0.0 and version <= 7.17.0,
// while it is optional for all subsequent versions.
String key = in.getTransportVersion().equals(TransportVersions.V_8_0_0) ? in.readString()
: in.getTransportVersion().onOrAfter(TransportVersions.V_7_17_1) ? in.readOptionalString()
: in.readString();
BytesRef from = in.readBoolean() ? in.readBytesRef() : null;
BytesRef to = in.readBoolean() ? in.readBytesRef() : null;
long docCount = in.readLong();
InternalAggregations aggregations = InternalAggregations.readFrom(in);
return new Bucket(format, keyed, key, from, to, docCount, aggregations);
}
@Override
public void writeTo(StreamOutput out) throws IOException {
if (out.getTransportVersion().equals(TransportVersions.V_8_0_0)) {
out.writeString(key == null ? generateKey(from, to, format) : key);
} else if (out.getTransportVersion().onOrAfter(TransportVersions.V_7_17_1)) {
out.writeOptionalString(key);
} else {
out.writeString(key == null ? generateKey(from, to, format) : key);
}
out.writeBoolean(from != null);
if (from != null) {
out.writeBytesRef(from);
}
out.writeBoolean(to != null);
if (to != null) {
out.writeBytesRef(to);
}
out.writeLong(docCount);
aggregations.writeTo(out);
}
@Override
public Object getKey() {
return getKeyAsString();
}
@Override
public String getKeyAsString() {
return this.key == null ? generateKey(this.from, this.to, this.format) : this.key;
}
@Override
public long getDocCount() {
return docCount;
}
@Override
public InternalAggregations getAggregations() {
return aggregations;
}
@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
final String key = getKeyAsString();
if (keyed) {
builder.startObject(key);
} else {
builder.startObject();
builder.field(CommonFields.KEY.getPreferredName(), key);
}
if (from != null) {
builder.field(CommonFields.FROM.getPreferredName(), getFrom());
}
if (to != null) {
builder.field(CommonFields.TO.getPreferredName(), getTo());
}
builder.field(CommonFields.DOC_COUNT.getPreferredName(), docCount);
aggregations.toXContentInternal(builder, params);
builder.endObject();
return builder;
}
@Override
public Object getFrom() {
return getFromAsString();
}
@Override
public String getFromAsString() {
return from == null ? null : format.format(from).toString();
}
@Override
public Object getTo() {
return getToAsString();
}
@Override
public String getToAsString() {
return to == null ? null : format.format(to).toString();
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
Bucket bucket = (Bucket) o;
if (docCount != bucket.docCount) return false;
// keyed and format are ignored since they are already tested on the InternalBinaryRange object
return Objects.equals(key, bucket.key)
&& Objects.equals(from, bucket.from)
&& Objects.equals(to, bucket.to)
&& Objects.equals(aggregations, bucket.aggregations);
}
@Override
public int hashCode() {
return Objects.hash(getClass(), docCount, key, from, to, aggregations);
}
Bucket finalizeSampling(SamplingContext samplingContext) {
return new Bucket(
format,
keyed,
key,
from,
to,
samplingContext.scaleUp(docCount),
InternalAggregations.finalizeSampling(aggregations, samplingContext)
);
}
}
protected final DocValueFormat format;
protected final boolean keyed;
private final List buckets;
public InternalBinaryRange(String name, DocValueFormat format, boolean keyed, List buckets, Map metadata) {
super(name, metadata);
this.format = format;
this.keyed = keyed;
this.buckets = buckets;
}
/**
* Read from a stream.
*/
public InternalBinaryRange(StreamInput in) throws IOException {
super(in);
format = in.readNamedWriteable(DocValueFormat.class);
keyed = in.readBoolean();
buckets = in.readCollectionAsList(stream -> Bucket.createFromStream(stream, format, keyed));
}
@Override
protected void doWriteTo(StreamOutput out) throws IOException {
out.writeNamedWriteable(format);
out.writeBoolean(keyed);
out.writeCollection(buckets);
}
@Override
public String getWriteableName() {
return IpRangeAggregationBuilder.NAME;
}
@Override
public List getBuckets() {
return unmodifiableList(buckets);
}
@Override
public InternalBinaryRange create(List buckets) {
return new InternalBinaryRange(name, format, keyed, buckets, metadata);
}
@Override
public Bucket createBucket(InternalAggregations aggregations, Bucket prototype) {
return new Bucket(format, keyed, prototype.key, prototype.from, prototype.to, prototype.docCount, aggregations);
}
@Override
protected AggregatorReducer getLeaderReducer(AggregationReduceContext reduceContext, int size) {
return new AggregatorReducer() {
final List aggregations = new ArrayList<>(size);
@Override
public void accept(InternalAggregation aggregation) {
aggregations.add((InternalBinaryRange) aggregation);
}
@Override
public InternalAggregation get() {
reduceContext.consumeBucketsAndMaybeBreak(buckets.size());
long[] docCounts = new long[buckets.size()];
InternalAggregations[][] aggs = new InternalAggregations[buckets.size()][];
for (int i = 0; i < aggs.length; ++i) {
aggs[i] = new InternalAggregations[aggregations.size()];
}
for (int i = 0; i < aggregations.size(); ++i) {
InternalBinaryRange range = aggregations.get(i);
if (range.buckets.size() != buckets.size()) {
throw new IllegalStateException(
"Expected [" + buckets.size() + "] buckets, but got [" + range.buckets.size() + "]"
);
}
for (int j = 0; j < buckets.size(); ++j) {
Bucket bucket = range.buckets.get(j);
docCounts[j] += bucket.docCount;
aggs[j][i] = bucket.aggregations;
}
}
List buckets = new ArrayList<>(getBuckets().size());
for (int i = 0; i < getBuckets().size(); ++i) {
Bucket b = getBuckets().get(i);
buckets.add(
new Bucket(
format,
keyed,
b.key,
b.from,
b.to,
docCounts[i],
InternalAggregations.reduce(Arrays.asList(aggs[i]), reduceContext)
)
);
}
return new InternalBinaryRange(name, format, keyed, buckets, metadata);
}
};
}
@Override
public InternalAggregation finalizeSampling(SamplingContext samplingContext) {
return new InternalBinaryRange(
name,
format,
keyed,
buckets.stream().map(b -> b.finalizeSampling(samplingContext)).toList(),
metadata
);
}
private Bucket reduceBucket(List buckets, AggregationReduceContext context) {
assert buckets.isEmpty() == false;
final List aggregations = new BucketAggregationList<>(buckets);
final InternalAggregations aggs = InternalAggregations.reduce(aggregations, context);
return createBucket(aggs, buckets.get(0));
}
@Override
public XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
if (keyed) {
builder.startObject(CommonFields.BUCKETS.getPreferredName());
} else {
builder.startArray(CommonFields.BUCKETS.getPreferredName());
}
for (Bucket range : buckets) {
range.toXContent(builder, params);
}
if (keyed) {
builder.endObject();
} else {
builder.endArray();
}
return builder;
}
@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;
InternalBinaryRange that = (InternalBinaryRange) obj;
return Objects.equals(buckets, that.buckets) && Objects.equals(format, that.format) && Objects.equals(keyed, that.keyed);
}
public int hashCode() {
return Objects.hash(super.hashCode(), buckets, format, keyed);
}
}