All Downloads are FREE. Search and download functionalities are using the official Maven repository.

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);
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy