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.VertexiumException;
import org.vertexium.VertexiumObject;

import java.util.*;
import java.util.stream.Collectors;

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, o -> o);
        elementsByProperty = collapseBucketsByCase(elementsByProperty);

        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 Map> collapseBucketsByCase(Map> elementsByProperty) {
        Map>>> stringEntries = new HashMap<>();
        Map> results = new HashMap<>();

        // for strings first group them by there lowercase version
        for (Map.Entry> entry : elementsByProperty.entrySet()) {
            if (entry.getKey() instanceof String) {
                String lowerCaseKey = ((String) entry.getKey()).toLowerCase();
                List>> l = stringEntries.computeIfAbsent(lowerCaseKey, s -> new ArrayList<>());
                l.add(entry);
            } else {
                results.put(entry.getKey(), entry.getValue());
            }
        }

        // for strings find the best key (the one with the most entries) and use that as the bucket name
        for (Map.Entry>>> entry : stringEntries.entrySet()) {
            results.put(
                    findBestKey(entry.getValue()),
                    entry.getValue().stream()
                            .flatMap(l -> l.getValue().stream())
                            .collect(Collectors.toList())
            );
        }
        return results;


    }

    private Object findBestKey(List>> value) {
        int longestListLength = 0;
        String longestString = null;
        for (Map.Entry> entry : value) {
            if (entry.getValue().size() >= longestListLength) {
                longestListLength = entry.getValue().size();
                longestString = (String) entry.getKey();
            }
        }
        return longestString;
    }

    private HistogramResult getCalendarFieldHistogramResult(final CalendarFieldAggregation agg, Iterator it) {
        String propertyName = agg.getPropertyName();
        final Calendar calendar = GregorianCalendar.getInstance(agg.getTimeZone());
        Map> elementsByProperty = getElementsByProperty(it, propertyName, 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 vertexiumObject = it.next();
            Iterable