org.vertexium.query.DefaultGraphQueryIterable 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.*;
import org.vertexium.property.StreamingPropertyValue;
import org.vertexium.scoring.ScoringStrategy;
import org.vertexium.util.CloseableIterator;
import org.vertexium.util.CloseableUtils;

import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;

import static org.vertexium.util.IterableUtils.count;
import static org.vertexium.util.IterableUtils.toList;
import static org.vertexium.util.Preconditions.checkNotNull;

public class DefaultGraphQueryIterable implements
    Iterable,
    QueryResultsIterable,
    IterableWithScores {
    private final QueryParameters parameters;
    private final Iterable iterable;
    private final boolean evaluateQueryString;
    private final boolean evaluateHasContainers;

    public DefaultGraphQueryIterable(
        QueryParameters parameters,
        Iterable iterable,
        boolean evaluateQueryString,
        boolean evaluateHasContainers,
        boolean evaluateSortContainers
    ) {
        checkNotNull(iterable, "iterable cannot be null");
        this.parameters = parameters;
        this.evaluateQueryString = evaluateQueryString;
        this.evaluateHasContainers = evaluateHasContainers;
        if (evaluateSortContainers && this.parameters.getSortContainers().size() > 0) {
            this.iterable = sortUsingSortContainers(iterable, parameters.getSortContainers());
        } else if (evaluateHasContainers && this.parameters.getScoringStrategy() != null) {
            this.iterable = sortUsingScoringStrategy(iterable, parameters.getScoringStrategy());
        } else {
            this.iterable = iterable;
        }
    }

    private Iterable sortUsingScoringStrategy(Iterable iterable, ScoringStrategy scoringStrategy) {
        List list = toList(iterable);
        list.sort(new ScoringStrategyComparator<>(scoringStrategy));
        return list;
    }

    private List sortUsingSortContainers(Iterable iterable, List sortContainers) {
        List list = toList(iterable);
        list.sort(new SortContainersComparator<>(sortContainers));
        return list;
    }

    @Override
    public Iterator iterator() {
        return iterator(false);
    }

    protected Iterator iterator(final boolean iterateAll) {
        final Iterator it = iterable.iterator();

        return new CloseableIterator() {
            public T next;
            public T current;
            public long count;

            @Override
            public boolean hasNext() {
                loadNext();
                if (next == null) {
                    close();
                }
                return next != null;
            }

            @Override
            public T next() {
                loadNext();
                if (next == null) {
                    throw new NoSuchElementException();
                }
                this.current = this.next;
                this.next = null;
                return this.current;
            }

            @Override
            public void remove() {
                throw new UnsupportedOperationException();
            }

            @Override
            public void close() {
                CloseableUtils.closeQuietly(it);
                DefaultGraphQueryIterable.this.close();
            }

            private void loadNext() {
                if (this.next != null) {
                    return;
                }

                if (!iterateAll && parameters.getLimit() != null && (this.count >= parameters.getSkip() + parameters.getLimit())) {
                    return;
                }

                while (it.hasNext()) {
                    T elem = it.next();
                    VertexiumObject vertexiumElem = elem instanceof VertexiumObject ? (VertexiumObject) elem : null;

                    boolean match = true;
                    if (evaluateHasContainers && vertexiumElem != null) {
                        for (QueryBase.HasContainer has : parameters.getHasContainers()) {
                            if (!has.isMatch(vertexiumElem)) {
                                match = false;
                                break;
                            }
                        }
                        if (vertexiumElem instanceof Edge && parameters.getEdgeLabels().size() > 0) {
                            Edge edge = (Edge) vertexiumElem;
                            if (!parameters.getEdgeLabels().contains(edge.getLabel())) {
                                match = false;
                            }
                        }
                        if (parameters.getIds() != null) {
                            if (vertexiumElem instanceof Element) {
                                if (!parameters.getIds().contains(((Element) vertexiumElem).getId())) {
                                    match = false;
                                }
                            } else if (vertexiumElem instanceof ExtendedDataRow) {
                                if (!parameters.getIds().contains(((ExtendedDataRow) vertexiumElem).getId().getElementId())) {
                                    match = false;
                                }
                            } else {
                                throw new VertexiumException("Unhandled element type: " + vertexiumElem.getClass().getName());
                            }
                        }

                        if (parameters.getMinScore() != null) {
                            if (parameters.getScoringStrategy() == null) {
                                match = false;
                            } else {
                                Double elementScore = parameters.getScoringStrategy().getScore(vertexiumElem);
                                if (elementScore == null) {
                                    match = false;
                                } else {
                                    match = elementScore >= parameters.getMinScore();
                                }
                            }
                        }
                    }
                    if (!match) {
                        continue;
                    }
                    if (evaluateQueryString
                        && vertexiumElem != null
                        && parameters instanceof QueryStringQueryParameters
                        && ((QueryStringQueryParameters) parameters).getQueryString() != null
                        && !evaluateQueryString(vertexiumElem, ((QueryStringQueryParameters) parameters).getQueryString())
                    ) {
                        continue;
                    }

                    this.count++;
                    if (!iterateAll && (this.count <= parameters.getSkip())) {
                        continue;
                    }

                    this.next = elem;
                    break;
                }
            }
        };
    }

    protected boolean evaluateQueryString(VertexiumObject vertexiumObject, String queryString) {
        if (vertexiumObject instanceof Element) {
            return evaluateQueryString((Element) vertexiumObject, queryString);
        } else if (vertexiumObject instanceof ExtendedDataRow) {
            return evaluateQueryString((ExtendedDataRow) vertexiumObject, queryString);
        } else {
            throw new VertexiumException("Unhandled VertexiumObject type: " + vertexiumObject.getClass().getName());
        }
    }

    private boolean evaluateQueryString(Element element, String queryString) {
        for (Property property : element.getProperties()) {
            if (evaluateQueryStringOnValue(property.getValue(), queryString)) {
                return true;
            }
        }
        return false;
    }

    private boolean evaluateQueryString(ExtendedDataRow extendedDataRow, String queryString) {
        for (Property property : extendedDataRow.getProperties()) {
            if (evaluateQueryStringOnValue(property.getValue(), queryString)) {
                return true;
            }
        }
        return false;
    }

    private boolean evaluateQueryStringOnValue(Object value, String queryString) {
        if (value == null) {
            return false;
        }
        if (queryString.equals("*")) {
            return true;
        }
        if (value instanceof StreamingPropertyValue) {
            value = ((StreamingPropertyValue) value).readToString();
        }
        String valueString = value.toString().toLowerCase();
        return valueString.contains(queryString.toLowerCase());
    }

    @Override
    public long getTotalHits() {
        // a limit could be set on a query which could prevent all items being returned
        return count(this.iterator(true));
    }

    @Override
    public void close() {
        CloseableUtils.closeQuietly(iterable);
    }

    @Override
    public  TResult getAggregationResult(String name, Class resultType) {
        throw new VertexiumException("Could not find aggregation with name: " + name);
    }

    @Override
    public Double getScore(Object id) {
        if (parameters.getScoringStrategy() != null) {
            VertexiumObject vertexiumObject = findVertexiumObjectById(id);
            if (vertexiumObject != null) {
                return parameters.getScoringStrategy().getScore(vertexiumObject);
            }
        }
        return 0.0;
    }

    private VertexiumObject findVertexiumObjectById(Object id) {
        Iterator it = iterator(true);
        while (it.hasNext()) {
            T obj = it.next();
            if (obj instanceof VertexiumObject) {
                VertexiumObject vertexiumObject = (VertexiumObject) obj;
                if (vertexiumObject.getId().equals(id)) {
                    return vertexiumObject;
                }
            }
        }
        return null;
    }
}