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

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

There is a newer version: 4.10.0
Show newest version
package org.vertexium.query;

import org.vertexium.*;

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

import static org.vertexium.query.TermsResult.NOT_COMPUTED;

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;
        }
        if (agg instanceof CardinalityAggregation) {
            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);
        }
        if (agg instanceof CardinalityAggregation) {
            return (TResult) getCardinalityAggregationResult((CardinalityAggregation) agg, it);
        }
        throw new VertexiumException("Unhandled aggregation: " + agg.getClass().getName());
    }

    private CardinalityResult getCardinalityAggregationResult(CardinalityAggregation agg, Iterator it) {
        String fieldName = agg.getPropertyName();

        if (Element.ID_PROPERTY_NAME.equals(fieldName)
            || Edge.LABEL_PROPERTY_NAME.equals(fieldName)
            || Edge.OUT_VERTEX_ID_PROPERTY_NAME.equals(fieldName)
            || Edge.IN_VERTEX_ID_PROPERTY_NAME.equals(fieldName)
            || ExtendedDataRow.TABLE_NAME.equals(fieldName)
            || ExtendedDataRow.ROW_ID.equals(fieldName)
            || ExtendedDataRow.ELEMENT_ID.equals(fieldName)
            || ExtendedDataRow.ELEMENT_TYPE.equals(fieldName)) {
            Set values = new HashSet<>();
            while (it.hasNext()) {
                T vertexiumObject = it.next();
                Iterable propertyValues = vertexiumObject.getPropertyValues(fieldName);
                for (Object propertyValue : propertyValues) {
                    values.add(propertyValue);
                }
            }
            return new CardinalityResult(values.size());
        } else {
            throw new VertexiumException("Cannot use cardinality aggregation on properties with visibility: " + fieldName);
        }
    }

    private TermsResult getTermsAggregationResult(TermsAggregation agg, Iterator it) {
        String propertyName = agg.getPropertyName();
        String missingFlag = "PROPERTY_MISSING";
        Map> elementsByProperty = getElementsByProperty(it, propertyName, o -> o, missingFlag);
        elementsByProperty = collapseBucketsByCase(elementsByProperty);

        long hasNotCount = NOT_COMPUTED;
        if (agg.isIncludeHasNotCount()) {
            hasNotCount = elementsByProperty.containsKey(missingFlag) ? elementsByProperty.get(missingFlag).size() : 0;
        }
        elementsByProperty.remove(missingFlag);

        long other = 0;
        List buckets = new ArrayList<>();
        for (Map.Entry> entry : elementsByProperty.entrySet()) {
            Object key = entry.getKey();
            int count = entry.getValue().size();
            if (agg.getSize() == null || buckets.size() < agg.getSize()) {
                Map nestedResults = getNestedResults(agg.getNestedAggregations(), entry.getValue());
                buckets.add(new TermsBucket(key, count, nestedResults));
            } else {
                other += count;
            }
        }
        return new TermsResult(buckets, other, 0, hasNotCount);
    }

    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());
        }, null);

        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, TKey missingFlag) {
        Map> elementsByProperty = new HashMap<>();
        while (it.hasNext()) {
            T vertexiumObject = it.next();
            Iterable values = vertexiumObject.getPropertyValues(propertyName);
            boolean hasValues = false;
            for (Object value : values) {
                hasValues = true;
                TKey convertedValue = valueConverter.convert(value);
                elementsByProperty.computeIfAbsent(convertedValue, k -> new ArrayList<>())
                    .add(vertexiumObject);
            }

            if (!hasValues && missingFlag != null) {
                elementsByProperty.computeIfAbsent(missingFlag, k -> new ArrayList<>()).add(vertexiumObject);
            }
        }
        return elementsByProperty;
    }

    private interface ValueConverter {
        T convert(Object o);
    }
}