Examples with ScrollState org.springframework.data.elasticsearch.client.util.ScrollState used on opensource projects

Example 1 with ScrollState

use of org.springframework.data.elasticsearch.client.util.ScrollState in project spring-data-elasticsearch by spring-projects.

the class StreamQueries method streamResults.

/**
 * Stream query results using {@link SearchScrollHits}.
 *
 * @param maxCount the maximum number of entities to return, a value of 0 means that all available entities are
 *          returned
 * @param searchHits the initial hits
 * @param continueScrollFunction function to continue scrolling applies to the current scrollId.
 * @param clearScrollConsumer consumer to clear the scroll context by accepting the scrollIds to clear.
 * @param <T> the entity type
 * @return the {@link SearchHitsIterator}.
 */
static <T> SearchHitsIterator<T> streamResults(int maxCount, SearchScrollHits<T> searchHits, Function<String, SearchScrollHits<T>> continueScrollFunction, Consumer<List<String>> clearScrollConsumer) {
    Assert.notNull(searchHits, "searchHits must not be null.");
    Assert.notNull(searchHits.getScrollId(), "scrollId of searchHits must not be null.");
    Assert.notNull(continueScrollFunction, "continueScrollFunction must not be null.");
    Assert.notNull(clearScrollConsumer, "clearScrollConsumer must not be null.");
    AggregationsContainer<?> aggregations = searchHits.getAggregations();
    float maxScore = searchHits.getMaxScore();
    long totalHits = searchHits.getTotalHits();
    TotalHitsRelation totalHitsRelation = searchHits.getTotalHitsRelation();
    return new SearchHitsIterator<T>() {

        private volatile AtomicInteger currentCount = new AtomicInteger();

        private volatile Iterator<SearchHit<T>> currentScrollHits = searchHits.iterator();

        private volatile boolean continueScroll = currentScrollHits.hasNext();

        private volatile ScrollState scrollState = new ScrollState(searchHits.getScrollId());

        private volatile boolean isClosed = false;

        @Override
        public void close() {
            if (!isClosed) {
                clearScrollConsumer.accept(scrollState.getScrollIds());
                isClosed = true;
            }
        }

        @Override
        @Nullable
        public AggregationsContainer<?> getAggregations() {
            return aggregations;
        }

        @Override
        public float getMaxScore() {
            return maxScore;
        }

        @Override
        public long getTotalHits() {
            return totalHits;
        }

        @Override
        public TotalHitsRelation getTotalHitsRelation() {
            return totalHitsRelation;
        }

        @Override
        public boolean hasNext() {
            boolean hasNext = false;
            if (!isClosed && continueScroll && (maxCount <= 0 || currentCount.get() < maxCount)) {
                if (!currentScrollHits.hasNext()) {
                    SearchScrollHits<T> nextPage = continueScrollFunction.apply(scrollState.getScrollId());
                    currentScrollHits = nextPage.iterator();
                    scrollState.updateScrollId(nextPage.getScrollId());
                    continueScroll = currentScrollHits.hasNext();
                }
                hasNext = currentScrollHits.hasNext();
            }
            if (!hasNext) {
                close();
            }
            return hasNext;
        }

        @Override
        public SearchHit<T> next() {
            if (hasNext()) {
                currentCount.incrementAndGet();
                return currentScrollHits.next();
            }
            throw new NoSuchElementException();
        }

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