InternalDistinctDateHistogramFacet.java

package com.pearson.entech.elasticsearch.search.facet.approx.datehistogram;

import java.io.IOException;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

import org.elasticsearch.ElasticSearchException;
import org.elasticsearch.common.CacheRecycler;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.trove.ExtTLongObjectHashMap;
import org.elasticsearch.common.trove.iterator.TLongObjectIterator;
import org.elasticsearch.common.trove.procedure.TLongObjectProcedure;
import org.elasticsearch.common.trove.procedure.TObjectProcedure;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentBuilderString;
import org.elasticsearch.search.facet.Facet;
import org.elasticsearch.search.facet.InternalFacet;

import com.clearspring.analytics.stream.cardinality.CardinalityMergeException;

/**
 */
public abstract class InternalDistinctDateHistogramFacet extends InternalFacet implements DistinctDateHistogramFacet {

    public static final String TYPE = "distinct_date_histogram";
    protected ComparatorType comparatorType;

    ExtTLongObjectHashMap<DistinctCountPayload> counts;
    DistinctCountPayload overallCount;
    boolean cachedCounts;
    protected String name;

    public InternalDistinctDateHistogramFacet() {}

    public InternalDistinctDateHistogramFacet(final String facetName) {
        super(facetName);
        this.name = facetName;
    }

    @Override
    public String getType() {
        return TYPE;
    }

    public static void registerStreams() {
        LongInternalDistinctDateHistogramFacet.registerStreams();
        StringInternalDistinctDateHistogramFacet.registerStreams();
    }

    /**
     * A histogram entry representing a single entry within the result of a histogram facet.
     *
     * It holds a set of distinct values and the time.
     */
    public static class InternalEntry implements Entry {
        private final long time;
        private final long distinctCount;
        private final long totalCount;

        public InternalEntry(final long time, final long distinctCount, final long totalCount) {
            this.time = time;
            this.distinctCount = distinctCount;
            this.totalCount = totalCount;
        }

        @Override
        public long getTime() {
            return time;
        }

        @Override
        public long getDistinctCount() {
            return this.distinctCount;
        }

        @Override
        public long getTotalCount() {
            return this.totalCount;
        }
    }

    public List<Entry> entries() {

        // Materialize entries
        final Entry[] entries = new InternalEntry[counts.size()];
        counts.forEachEntry(new TLongObjectProcedure<DistinctCountPayload>() {
            int idx = 0;

            @Override
            public boolean execute(final long histoKey, final DistinctCountPayload payload) {
                entries[idx] = new InternalEntry(histoKey,
                        payload.getCardinality().cardinality(), payload.getCount());
                idx++;
                return true;
            }
        });

        // Sort and return them
        Arrays.sort(entries, comparatorType.comparator());
        return Arrays.asList(entries);
    }

    @Override
    public List<Entry> getEntries() {
        return entries();
    }

    @Override
    public Iterator<Entry> iterator() {
        return entries().iterator();
    }

    void releaseCache() {
        if(cachedCounts) {
            CacheRecycler.pushLongObjectMap(counts);
            cachedCounts = false;
            counts = null;
        }
    }

    static final class Fields {
        static final XContentBuilderString _TYPE = new XContentBuilderString("_type");
        static final XContentBuilderString ENTRIES = new XContentBuilderString("entries");
        static final XContentBuilderString TIME = new XContentBuilderString("time");
        static final XContentBuilderString COUNT = new XContentBuilderString("count");
        static final XContentBuilderString DISTINCT_COUNT = new XContentBuilderString("distinct_count");
        static final XContentBuilderString TOTAL_COUNT = new XContentBuilderString("total_count");
        static final XContentBuilderString TOTAL_DISTINCT_COUNT = new XContentBuilderString("total_distinct_count");
    }

    @Override
    public Facet reduce(final List<Facet> facets) {
        final ExtTLongObjectHashMap<DistinctCountPayload> counts = CacheRecycler.popLongObjectMap();

        System.out.println("Merging " + facets.size() + " facets");
        for(final Facet facet : facets) {
            final InternalDistinctDateHistogramFacet histoFacet = (InternalDistinctDateHistogramFacet) facet;
            System.out.println("Merging facet " + histoFacet);
            for(final TLongObjectIterator<DistinctCountPayload> it = histoFacet.counts.iterator(); it.hasNext();) {
                it.advance();
                final long facetKey = it.key();
                it.value().mergeInto(counts, facetKey);
            }
            histoFacet.releaseCache();
        }

        final InternalDistinctDateHistogramFacet ret = newFacet();
        ret.comparatorType = comparatorType;
        ret.counts = counts;
        ret.cachedCounts = true;
        System.out.println("Returning merged facet " + ret);
        return ret;
    }

    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()).append(": ");
        counts.forEachEntry(new TLongObjectProcedure<DistinctCountPayload>() {
            @Override
            public boolean execute(final long histoKey, final DistinctCountPayload count) {
                sb.append(String.format("%d={%s}, ", histoKey, count));
                return true;
            }
        });
        return counts.size() > 0 ? sb.substring(0, sb.lastIndexOf(",")) : sb.toString();
    }

    protected abstract InternalDistinctDateHistogramFacet newFacet();

    /**
     * Builds the final JSON result.
     *
     * For each time entry we provide the number of distinct values in the time range.
     */
    @Override
    public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException {
        builder.startObject(this.name);
        builder.field(Fields._TYPE, TYPE);
        builder.field(Fields.TOTAL_COUNT, getTotalCount());
        builder.field(Fields.TOTAL_DISTINCT_COUNT, getDistinctCount());
        builder.startArray(Fields.ENTRIES);
        for(final Entry entry : entries()) {
            builder.startObject();
            builder.field(Fields.TIME, entry.getTime());
            builder.field(Fields.COUNT, entry.getTotalCount());
            builder.field(Fields.DISTINCT_COUNT, entry.getDistinctCount());
            builder.endObject();
        }
        builder.endArray();
        builder.endObject();
        return builder;
    }

    @Override
    public long getDistinctCount() {
        return overallCount().getCardinality().cardinality();
    }

    @Override
    public long getTotalCount() {
        return overallCount().getCount();
    }

    private synchronized DistinctCountPayload overallCount() {
        if(overallCount == null) {
            overallCount = new DistinctCountPayload(Integer.MAX_VALUE);
            counts.forEachValue(new TObjectProcedure<DistinctCountPayload>() {
                @Override
                public boolean execute(final DistinctCountPayload slice) {
                    try {
                        overallCount.merge(slice);
                    } catch(final CardinalityMergeException e) {
                        throw new IllegalStateException(e);
                    }
                    return true;
                }
            });
        }
        return overallCount;
    }

    /**
     * The reader for the internal transport protocol.
     */
    @Override
    public void readFrom(final StreamInput in) throws IOException {
        name = in.readString();
        comparatorType = ComparatorType.fromId(in.readByte());
        final int size = in.readVInt();
        counts = CacheRecycler.popLongObjectMap();
        cachedCounts = true;
        for(int i = 0; i < size; i++) {
            final long key = in.readLong();
            final long count = in.readVLong();
            final int cardLength = in.readVInt();
            final byte[] cardinality = new byte[cardLength];
            in.read(cardinality, 0, cardLength);
            try {
                counts.put(key, new DistinctCountPayload(count, cardinality));
            } catch(final ClassNotFoundException e) {
                throw new ElasticSearchException("Unable to deserialize facet cardinality object", e);
            }
        }
    }

    /**
     * The writer for the internal transport protocol.
     */
    @Override
    public void writeTo(final StreamOutput out) throws IOException {
        out.writeString(getName());
        out.writeByte(comparatorType.id());
        out.writeVInt(counts.size());
        for(final TLongObjectIterator<DistinctCountPayload> it = counts.iterator(); it.hasNext();) {
            it.advance();
            out.writeLong(it.key());
            out.writeVLong(it.value().getCount());
            final byte[] cardinality = it.value().cardinalityBytes();
            out.writeVInt(cardinality.length);
            out.write(cardinality);
        }
        releaseCache();
    }

}