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;
        }
        return false;
    }

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

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

        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 (TResult) new TermsResult(buckets);
    }

    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) {
        Map> elementsByProperty = new HashMap<>();
        while (it.hasNext()) {
            T elem = it.next();
            Iterable