org.vertexium.elasticsearch.utils.PagingIterable Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of vertexium-elasticsearch-singledocument Show documentation
There is a newer version: 3.0.4
Show newest version
package org.vertexium.elasticsearch.utils;

import org.vertexium.VertexiumException;
import org.vertexium.elasticsearch.ElasticSearchGraphQueryIterable;
import org.vertexium.query.*;

import java.io.IOException;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;

public abstract class PagingIterable implements
        Iterable,
        IterableWithTotalHits,
        IterableWithScores,
        IterableWithHistogramResults,
        IterableWithTermsResults,
        IterableWithGeohashResults,
        IterableWithStatisticsResults,
        QueryResultsIterable {
    private final long skip;
    private final Long limit;
    private boolean isFirstCallToIterator;
    private final ElasticSearchGraphQueryIterable firstIterable;
    private final int pageSize;

    public PagingIterable(long skip, Long limit, int pageSize) {
        this.skip = skip;
        this.limit = limit;
        this.pageSize = pageSize;

        // This is a bit of a hack. Because the underlying iterable is the iterable with geohash results, histogram results, etc.
        //   we need to grab the first iterable to get the results out.
        int firstIterableLimit = Math.min(pageSize, limit == null ? Integer.MAX_VALUE : limit.intValue());
        this.firstIterable = getPageIterable((int) this.skip, firstIterableLimit, true);
        this.isFirstCallToIterator = true;
    }

    @Override
    public GeohashResult getGeohashResults(String name) {
        return this.firstIterable.getGeohashResults(name);
    }

    @Override
    public HistogramResult getHistogramResults(String name) {
        return this.firstIterable.getHistogramResults(name);
    }

    @Override
    public StatisticsResult getStatisticsResults(String name) {
        return this.firstIterable.getStatisticsResults(name);
    }

    @Override
    public TermsResult getTermsResults(String name) {
        return this.firstIterable.getTermsResults(name);
    }

    @Override
    public Map getScores() {
        return this.firstIterable.getScores();
    }

    @Override
    public void close() throws IOException {

    }

    @Override
    public  TResult getAggregationResult(String name, Class resultType) {
        return this.firstIterable.getAggregationResult(name, resultType);
    }

    @Override
    public long getTotalHits() {
        return this.firstIterable.getTotalHits();
    }

    protected abstract ElasticSearchGraphQueryIterable getPageIterable(int skip, int limit, boolean includeAggregations);

    @Override
    public Iterator iterator() {
        MyIterator it = new MyIterator<>(isFirstCallToIterator ? firstIterable : null, skip, limit, pageSize, PagingIterable.this::getPageIterable);
        isFirstCallToIterator = false;
        return it;
    }

    private interface GetPageIterableFunction {
        Iterable getPageIterable(int skip, int limit, boolean includeAggregations);
    }

    private static class MyIterator implements Iterator {
        private Iterable firstIterable;
        private final int pageSize;
        private final GetPageIterableFunction getPageIterableFunction;
        private int nextSkip;
        private int limit;
        private int currentIteratorCount;
        private Iterator currentIterator;

        public MyIterator(Iterable firstIterable, long skip, Long limit, int pageSize, GetPageIterableFunction getPageIterableFunction) {
            this.firstIterable = firstIterable;
            this.pageSize = pageSize;
            this.getPageIterableFunction = getPageIterableFunction;
            this.nextSkip = (int) skip;
            this.limit = (int) (limit == null ? Integer.MAX_VALUE : limit);
            this.currentIterator = getNextIterator();
        }

        @Override
        public boolean hasNext() {
            while (true) {
                if (currentIterator == null) {
                    if (currentIteratorCount == 0) {
                        return false;
                    }
                    currentIterator = getNextIterator();
                    if (currentIterator == null) {
                        return false;
                    }
                }
                if (currentIterator.hasNext()) {
                    return true;
                }
                currentIterator = null;
            }
        }

        @Override
        public T next() {
            if (hasNext()) {
                limit--;
                currentIteratorCount++;
                return currentIterator.next();
            }
            throw new NoSuchElementException();
        }

        private Iterator getNextIterator() {
            if (limit <= 0) {
                return null;
            }
            int limit = Math.min(pageSize, this.limit);
            currentIteratorCount = 0;
            if (firstIterable == null) {
                firstIterable = getPageIterableFunction.getPageIterable(nextSkip, limit, false);
            }
            Iterator it = firstIterable.iterator();
            firstIterable = null;
            nextSkip += limit;
            return it;
        }

        @Override
        public void remove() {
            throw new VertexiumException("remove not implemented");
        }
    }
}