org.vertexium.query.DefaultGraphQueryIterableWithAggregations Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of vertexium-core Show documentation
There is a newer version: 4.10.0
package org.vertexium.query;

import org.vertexium.Element;
import org.vertexium.VertexiumException;

import java.util.*;

public class DefaultGraphQueryIterableWithAggregations extends DefaultGraphQueryIterable {
    private final Collection aggregations;

    public DefaultGraphQueryIterableWithAggregations(
            QueryParameters parameters,
            Iterable iterable,
            boolean evaluateQueryString,
            boolean evaluateHasContainers,
            boolean evaluateSortContainers,
            Collection aggregations
    ) {
        super(parameters, iterable, evaluateQueryString, evaluateHasContainers, evaluateSortContainers);
        this.aggregations = aggregations;
    }

    @Override
    public  TResult getAggregationResult(String name, Class resultType) {
        for (Aggregation agg : this.aggregations) {
            if (agg.getAggregationName().equals(name)) {
                return getAggregationResult(agg, this.iterator(true));
            }
        }
        return super.getAggregationResult(name, resultType);
    }

    public static boolean isAggregationSupported(Aggregation agg) {
        if (agg instanceof TermsAggregation) {
            return true;
        }
        if (agg instanceof CalendarFieldAggregation) {
            return true;
        }
        return false;
    }

    @SuppressWarnings("unchecked")
    public  TResult getAggregationResult(Aggregation agg, Iterator it) {
        if (agg instanceof TermsAggregation) {
            return (TResult) getTermsAggregationResult((TermsAggregation) agg, it);
        }
        if (agg instanceof CalendarFieldAggregation) {
            return (TResult) getCalendarFieldHistogramResult((CalendarFieldAggregation) agg, it);
        }
        throw new VertexiumException("Unhandled aggregation: " + agg.getClass().getName());
    }

    private TermsResult getTermsAggregationResult(TermsAggregation agg, Iterator it) {
        String propertyName = agg.getPropertyName();
        Map> elementsByProperty = getElementsByProperty(it, propertyName, new IdentityValueConverter());

        List buckets = new ArrayList<>();
        for (Map.Entry> entry : elementsByProperty.entrySet()) {
            Object key = entry.getKey();
            int count = entry.getValue().size();
            Map nestedResults = getNestedResults(agg.getNestedAggregations(), entry.getValue());
            buckets.add(new TermsBucket(key, count, nestedResults));
        }
        return new TermsResult(buckets);
    }

    private HistogramResult getCalendarFieldHistogramResult(final CalendarFieldAggregation agg, Iterator it) {
        String propertyName = agg.getPropertyName();
        final Calendar calendar = GregorianCalendar.getInstance(agg.getTimeZone());
        Map> elementsByProperty = getElementsByProperty(it, propertyName, new ValueConverter() {
            @Override
            public Integer convert(Object o) {
                Date d = (Date) o;
                calendar.setTime(d);
                //noinspection MagicConstant
                return calendar.get(agg.getCalendarField());
            }
        });

        Map buckets = new HashMap<>(24);
        for (Map.Entry> entry : elementsByProperty.entrySet()) {
            int key = entry.getKey();
            int count = entry.getValue().size();
            Map nestedResults = getNestedResults(agg.getNestedAggregations(), entry.getValue());
            buckets.put(key, new HistogramBucket(key, count, nestedResults));
        }
        return new HistogramResult(buckets.values());
    }

    private Map getNestedResults(Iterable nestedAggregations, List elements) {
        Map results = new HashMap<>();
        for (Aggregation nestedAggregation : nestedAggregations) {
            AggregationResult nestedResult = getAggregationResult(nestedAggregation, elements.iterator());
            results.put(nestedAggregation.getAggregationName(), nestedResult);
        }
        return results;
    }

    private  Map> getElementsByProperty(Iterator it, String propertyName, ValueConverter valueConverter) {
        Map> elementsByProperty = new HashMap<>();
        while (it.hasNext()) {
            T elem = it.next();
            Iterable