Please wait. This can take some minutes ...
Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
org.elasticsearch.search.aggregations.bucket.composite.InternalComposite Maven / Gradle / Ivy
/*
* 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.composite;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.PriorityQueue;
import org.elasticsearch.Version;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.InternalAggregation;
import org.elasticsearch.search.aggregations.InternalAggregations;
import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
import org.elasticsearch.search.aggregations.KeyComparable;
import org.elasticsearch.xcontent.XContentBuilder;
import java.io.IOException;
import java.util.AbstractMap;
import java.util.AbstractSet;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
public class InternalComposite extends InternalMultiBucketAggregation
implements
CompositeAggregation {
private final int size;
private final List buckets;
private final CompositeKey afterKey;
private final int[] reverseMuls;
private final MissingOrder[] missingOrders;
private final List sourceNames;
private final List formats;
private final boolean earlyTerminated;
InternalComposite(
String name,
int size,
List sourceNames,
List formats,
List buckets,
CompositeKey afterKey,
int[] reverseMuls,
MissingOrder[] missingOrders,
boolean earlyTerminated,
Map metadata
) {
super(name, metadata);
this.sourceNames = sourceNames;
this.formats = formats;
this.buckets = buckets;
this.afterKey = afterKey;
this.size = size;
this.reverseMuls = reverseMuls;
this.missingOrders = missingOrders;
this.earlyTerminated = earlyTerminated;
}
public InternalComposite(StreamInput in) throws IOException {
super(in);
this.size = in.readVInt();
this.sourceNames = in.readStringList();
this.formats = new ArrayList<>(sourceNames.size());
for (int i = 0; i < sourceNames.size(); i++) {
if (in.getVersion().onOrAfter(Version.V_6_3_0)) {
formats.add(in.readNamedWriteable(DocValueFormat.class));
} else {
formats.add(DocValueFormat.RAW);
}
}
this.reverseMuls = in.readIntArray();
if (in.getVersion().onOrAfter(Version.V_7_16_0)) {
this.missingOrders = in.readArray(MissingOrder::readFromStream, MissingOrder[]::new);
} else {
this.missingOrders = new MissingOrder[reverseMuls.length];
Arrays.fill(missingOrders, MissingOrder.DEFAULT);
}
this.buckets = in.readList((input) -> new InternalBucket(input, sourceNames, formats, reverseMuls, missingOrders));
if (in.getVersion().onOrAfter(Version.V_6_3_0)) {
this.afterKey = in.readBoolean() ? new CompositeKey(in) : null;
} else {
this.afterKey = buckets.size() > 0 ? buckets.get(buckets.size() - 1).key : null;
}
this.earlyTerminated = in.getVersion().onOrAfter(Version.V_7_6_0) ? in.readBoolean() : false;
}
@Override
protected void doWriteTo(StreamOutput out) throws IOException {
out.writeVInt(size);
out.writeStringCollection(sourceNames);
if (out.getVersion().onOrAfter(Version.V_6_3_0)) {
for (DocValueFormat format : formats) {
out.writeNamedWriteable(format);
}
}
out.writeIntArray(reverseMuls);
if (out.getVersion().onOrAfter(Version.V_7_16_0)) {
out.writeArray((o, order) -> order.writeTo(o), missingOrders);
}
out.writeList(buckets);
if (out.getVersion().onOrAfter(Version.V_6_3_0)) {
out.writeBoolean(afterKey != null);
if (afterKey != null) {
afterKey.writeTo(out);
}
}
if (out.getVersion().onOrAfter(Version.V_7_6_0)) {
out.writeBoolean(earlyTerminated);
}
}
@Override
public XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
return CompositeAggregation.toXContentFragment(this, builder, params);
}
@Override
public String getWriteableName() {
return CompositeAggregationBuilder.NAME;
}
@Override
public InternalComposite create(List newBuckets) {
/**
* This is used by pipeline aggregations to filter/remove buckets so we
* keep the afterKey
of the original aggregation in order
* to be able to retrieve the next page even if all buckets have been filtered.
*/
return new InternalComposite(
name,
size,
sourceNames,
formats,
newBuckets,
afterKey,
reverseMuls,
missingOrders,
earlyTerminated,
getMetadata()
);
}
@Override
public InternalBucket createBucket(InternalAggregations aggregations, InternalBucket prototype) {
return new InternalBucket(
prototype.sourceNames,
prototype.formats,
prototype.key,
prototype.reverseMuls,
prototype.missingOrders,
prototype.docCount,
aggregations
);
}
public int getSize() {
return size;
}
@Override
public List getBuckets() {
return buckets;
}
/**
* The formats used when writing the keys. Package private for testing.
*/
List getFormats() {
return formats;
}
@Override
public Map afterKey() {
if (afterKey != null) {
return new ArrayMap(sourceNames, formats, afterKey.values());
}
return null;
}
// Visible for tests
boolean isTerminatedEarly() {
return earlyTerminated;
}
// Visible for tests
int[] getReverseMuls() {
return reverseMuls;
}
@Override
public InternalAggregation reduce(List aggregations, ReduceContext reduceContext) {
PriorityQueue pq = new PriorityQueue(aggregations.size()) {
@Override
protected boolean lessThan(BucketIterator a, BucketIterator b) {
return a.compareTo(b) < 0;
}
};
boolean earlyTerminated = false;
for (InternalAggregation agg : aggregations) {
InternalComposite sortedAgg = (InternalComposite) agg;
earlyTerminated |= sortedAgg.earlyTerminated;
BucketIterator it = new BucketIterator(sortedAgg.buckets);
if (it.next() != null) {
pq.add(it);
}
}
InternalBucket lastBucket = null;
List buckets = new ArrayList<>();
List result = new ArrayList<>();
while (pq.size() > 0) {
BucketIterator bucketIt = pq.top();
if (lastBucket != null && bucketIt.current.compareKey(lastBucket) != 0) {
InternalBucket reduceBucket = reduceBucket(buckets, reduceContext);
buckets.clear();
result.add(reduceBucket);
if (result.size() >= size) {
break;
}
}
lastBucket = bucketIt.current;
buckets.add(bucketIt.current);
if (bucketIt.next() != null) {
pq.updateTop();
} else {
pq.pop();
}
}
if (buckets.size() > 0) {
InternalBucket reduceBucket = reduceBucket(buckets, reduceContext);
result.add(reduceBucket);
}
List reducedFormats = formats;
CompositeKey lastKey = null;
if (result.size() > 0) {
lastBucket = result.get(result.size() - 1);
/* Attach the formats from the last bucket to the reduced composite
* so that we can properly format the after key. */
reducedFormats = lastBucket.formats;
lastKey = lastBucket.getRawKey();
}
reduceContext.consumeBucketsAndMaybeBreak(result.size());
return new InternalComposite(
name,
size,
sourceNames,
reducedFormats,
result,
lastKey,
reverseMuls,
missingOrders,
earlyTerminated,
metadata
);
}
@Override
protected InternalBucket reduceBucket(List buckets, ReduceContext context) {
assert buckets.size() > 0;
List aggregations = new ArrayList<>(buckets.size());
long docCount = 0;
for (InternalBucket bucket : buckets) {
docCount += bucket.docCount;
aggregations.add(bucket.aggregations);
}
InternalAggregations aggs = InternalAggregations.reduce(aggregations, context);
/* Use the formats from the bucket because they'll be right to format
* the key. The formats on the InternalComposite doing the reducing are
* just whatever formats make sense for *its* index. This can be real
* trouble when the index doing the reducing is unmapped. */
List reducedFormats = buckets.get(0).formats;
return new InternalBucket(sourceNames, reducedFormats, buckets.get(0).key, reverseMuls, missingOrders, docCount, aggs);
}
@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;
InternalComposite that = (InternalComposite) obj;
return Objects.equals(size, that.size)
&& Objects.equals(buckets, that.buckets)
&& Objects.equals(afterKey, that.afterKey)
&& Arrays.equals(reverseMuls, that.reverseMuls)
&& Arrays.equals(missingOrders, that.missingOrders);
}
@Override
public int hashCode() {
return Objects.hash(super.hashCode(), size, buckets, afterKey, Arrays.hashCode(reverseMuls), Arrays.hashCode(missingOrders));
}
private static class BucketIterator implements Comparable {
final Iterator it;
InternalBucket current;
private BucketIterator(List buckets) {
this.it = buckets.iterator();
}
@Override
public int compareTo(BucketIterator other) {
return current.compareKey(other.current);
}
InternalBucket next() {
return current = it.hasNext() ? it.next() : null;
}
}
public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket
implements
CompositeAggregation.Bucket,
KeyComparable {
private final CompositeKey key;
private final long docCount;
private final InternalAggregations aggregations;
private final transient int[] reverseMuls;
private final transient MissingOrder[] missingOrders;
private final transient List sourceNames;
private final transient List formats;
InternalBucket(
List sourceNames,
List formats,
CompositeKey key,
int[] reverseMuls,
MissingOrder[] missingOrders,
long docCount,
InternalAggregations aggregations
) {
this.key = key;
this.docCount = docCount;
this.aggregations = aggregations;
this.reverseMuls = reverseMuls;
this.missingOrders = missingOrders;
this.sourceNames = sourceNames;
this.formats = formats;
}
InternalBucket(
StreamInput in,
List sourceNames,
List formats,
int[] reverseMuls,
MissingOrder[] missingOrders
) throws IOException {
this.key = new CompositeKey(in);
this.docCount = in.readVLong();
this.aggregations = InternalAggregations.readFrom(in);
this.reverseMuls = reverseMuls;
this.missingOrders = missingOrders;
this.sourceNames = sourceNames;
this.formats = formats;
}
@Override
public void writeTo(StreamOutput out) throws IOException {
key.writeTo(out);
out.writeVLong(docCount);
aggregations.writeTo(out);
}
@Override
public int hashCode() {
return Objects.hash(getClass(), docCount, key, aggregations);
}
@Override
public boolean equals(Object obj) {
if (obj == null || getClass() != obj.getClass()) {
return false;
}
InternalBucket that = (InternalBucket) obj;
return Objects.equals(docCount, that.docCount)
&& Objects.equals(key, that.key)
&& Objects.equals(aggregations, that.aggregations);
}
@Override
public Map getKey() {
// returns the formatted key in a map
return new ArrayMap(sourceNames, formats, key.values());
}
// get the raw key (without formatting to preserve the natural order).
// visible for testing
CompositeKey getRawKey() {
return key;
}
@Override
public String getKeyAsString() {
StringBuilder builder = new StringBuilder();
builder.append('{');
for (int i = 0; i < key.size(); i++) {
if (i > 0) {
builder.append(", ");
}
builder.append(sourceNames.get(i));
builder.append('=');
builder.append(formatObject(key.get(i), formats.get(i)));
}
builder.append('}');
return builder.toString();
}
@Override
public long getDocCount() {
return docCount;
}
@Override
public Aggregations getAggregations() {
return aggregations;
}
/**
* The formats used when writing the keys. Package private for testing.
*/
List getFormats() {
return formats;
}
@Override
public int compareKey(InternalBucket other) {
for (int i = 0; i < key.size(); i++) {
if (key.get(i) == null) {
if (other.key.get(i) == null) {
continue;
}
return -1 * missingOrders[i].compareAnyValueToMissing(reverseMuls[i]);
} else if (other.key.get(i) == null) {
return missingOrders[i].compareAnyValueToMissing(reverseMuls[i]);
}
assert key.get(i).getClass() == other.key.get(i).getClass();
@SuppressWarnings("unchecked")
int cmp = key.get(i).compareTo(other.key.get(i)) * reverseMuls[i];
if (cmp != 0) {
return cmp;
}
}
return 0;
}
@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
/**
* See {@link CompositeAggregation#bucketToXContent}
*/
throw new UnsupportedOperationException("not implemented");
}
}
/**
* Format obj
using the provided {@link DocValueFormat}.
* If the format is equals to {@link DocValueFormat#RAW}, the object is returned as is
* for numbers and a string for {@link BytesRef}s.
*
* This method will then attempt to parse the formatted value using the specified format,
* and throw an IllegalArgumentException if parsing fails. This in turn prevents us from
* returning an after_key which we can't subsequently parse into the original value.
*/
static Object formatObject(Object obj, DocValueFormat format) {
if (obj == null) {
return null;
}
Object formatted = obj;
Object parsed;
if (obj.getClass() == BytesRef.class) {
BytesRef value = (BytesRef) obj;
if (format == DocValueFormat.RAW) {
formatted = value.utf8ToString();
} else {
formatted = format.format(value);
}
parsed = format.parseBytesRef(formatted.toString());
if (parsed.equals(obj) == false) {
throw new IllegalArgumentException(
"Format ["
+ format
+ "] created output it couldn't parse for value ["
+ obj
+ "] "
+ "of type ["
+ obj.getClass()
+ "]. parsed value: ["
+ parsed
+ "("
+ parsed.getClass()
+ ")]"
);
}
} else if (obj.getClass() == Long.class) {
long value = (long) obj;
if (format == DocValueFormat.RAW) {
formatted = value;
} else {
formatted = format.format(value);
}
parsed = format.parseLong(
formatted.toString(),
false,
() -> { throw new UnsupportedOperationException("Using now() is not supported in after keys"); }
);
if (parsed.equals(((Number) obj).longValue()) == false) {
throw new IllegalArgumentException(
"Format ["
+ format
+ "] created output it couldn't parse for value ["
+ obj
+ "] "
+ "of type ["
+ obj.getClass()
+ "]. parsed value: ["
+ parsed
+ "("
+ parsed.getClass()
+ ")]"
);
}
} else if (obj.getClass() == Double.class) {
double value = (double) obj;
if (format == DocValueFormat.RAW) {
formatted = value;
} else {
formatted = format.format(value);
}
parsed = format.parseDouble(
formatted.toString(),
false,
() -> { throw new UnsupportedOperationException("Using now() is not supported in after keys"); }
);
if (parsed.equals(((Number) obj).doubleValue()) == false) {
throw new IllegalArgumentException(
"Format ["
+ format
+ "] created output it couldn't parse for value ["
+ obj
+ "] "
+ "of type ["
+ obj.getClass()
+ "]. parsed value: ["
+ parsed
+ "("
+ parsed.getClass()
+ ")]"
);
}
}
return formatted;
}
static class ArrayMap extends AbstractMap implements Comparable {
final List keys;
final Comparable>[] values;
final List formats;
ArrayMap(List keys, List formats, Comparable>[] values) {
assert keys.size() == values.length && keys.size() == formats.size();
this.keys = keys;
this.formats = formats;
this.values = values;
}
@Override
public int size() {
return values.length;
}
@Override
public Object get(Object key) {
for (int i = 0; i < keys.size(); i++) {
if (key.equals(keys.get(i))) {
return formatObject(values[i], formats.get(i));
}
}
return null;
}
@Override
public Set> entrySet() {
return new AbstractSet>() {
@Override
public Iterator> iterator() {
return new Iterator>() {
int pos = 0;
@Override
public boolean hasNext() {
return pos < values.length;
}
@Override
public Entry next() {
SimpleEntry entry = new SimpleEntry<>(
keys.get(pos),
formatObject(values[pos], formats.get(pos))
);
++pos;
return entry;
}
};
}
@Override
public int size() {
return keys.size();
}
};
}
@Override
@SuppressWarnings({ "rawtypes", "unchecked" })
public int compareTo(ArrayMap that) {
if (that == this) {
return 0;
}
int idx = 0;
int max = Math.min(this.keys.size(), that.keys.size());
while (idx < max) {
int compare = compareNullables(keys.get(idx), that.keys.get(idx));
if (compare == 0) {
compare = compareNullables((Comparable) values[idx], (Comparable) that.values[idx]);
}
if (compare != 0) {
return compare;
}
idx++;
}
if (idx < keys.size()) {
return 1;
}
if (idx < that.keys.size()) {
return -1;
}
return 0;
}
}
private static > int compareNullables(T a, T b) {
if (a == b) {
return 0;
}
if (a == null) {
return -1;
}
if (b == null) {
return 1;
}
return a.compareTo(b);
}
}