org.vertexium.elasticsearch.utils.InfiniteScrollIterable 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.elasticsearch.action.search.SearchResponse;
import org.vertexium.query.AggregationResult;
import org.vertexium.query.QueryResultsIterable;
import org.vertexium.util.CloseableIterator;
import org.vertexium.util.CloseableUtils;

import java.util.ArrayList;
import java.util.Iterator;

public abstract class InfiniteScrollIterable implements QueryResultsIterable {
    private SearchResponse firstResponse;
    private QueryResultsIterable firstIterable;
    private boolean initCalled;
    private boolean firstCall;

    protected abstract SearchResponse getInitialSearchResponse();

    protected abstract SearchResponse getNextSearchResponse(String scrollId);

    protected abstract QueryResultsIterable searchResponseToIterable(SearchResponse searchResponse);

    protected abstract void closeScroll(String scrollId);

    @Override
    public void close() {
    }

    private void init() {
        if (initCalled) {
            return;
        }
        firstResponse = getInitialSearchResponse();
        if (firstResponse == null) {
            firstIterable = null;
        } else {
            firstIterable = searchResponseToIterable(firstResponse);
        }
        firstCall = true;
        initCalled = true;
    }

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

    @Override
    public long getTotalHits() {
        init();
        if (firstIterable == null) {
            return 0;
        }
        return firstIterable.getTotalHits();
    }

    @Override
    public Iterator iterator() {
        init();
        if (firstResponse == null) {
            return new ArrayList().iterator();
        }

        SearchResponse response;
        Iterator it;
        if (firstCall) {
            response = firstResponse;
            it = firstIterable.iterator();
            firstCall = false;
        } else {
            response = getInitialSearchResponse();
            it = searchResponseToIterable(response).iterator();
        }
        String scrollId = response.getScrollId();
        return new InfiniteIterator(scrollId, it);
    }

    private class InfiniteIterator implements CloseableIterator {
        private final String scrollId;
        private Iterator it;
        private T next;
        private T current;
        private long currentResultNumber = 0;

        public InfiniteIterator(String scrollId, Iterator it) {
            this.scrollId = scrollId;
            this.it = it;
        }

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

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

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

            if (it.hasNext()) {
                this.next = it.next();
                currentResultNumber++;
            } else {
                CloseableUtils.closeQuietly(it);
                it = null;

                if (getTotalHits() > currentResultNumber) {
                    QueryResultsIterable iterable = searchResponseToIterable(getNextSearchResponse(scrollId));
                    it = iterable.iterator();
                    if (!it.hasNext()) {
                        it = null;
                    } else {
                        this.next = it.next();
                        currentResultNumber++;
                    }
                }
            }
        }

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

        @Override
        public void close() {
            closeScroll(scrollId);
            CloseableUtils.closeQuietly(it);
            InfiniteScrollIterable.this.close();
        }
    }
}