ResultSetSpliterator.java

package  test;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Spliterator;
import java.util.function.Consumer;

import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;

public class ResultSetSpliterator implements Spliterator<Row> {

	private final ResultSet rs;
	private final Iterator<Row> iterator;
	private long estimate = Long.MAX_VALUE;

	public ResultSetSpliterator(ResultSet rs) {
		this.rs = rs;
		this.iterator = rs.iterator();
	}

	@Override
	public boolean tryAdvance(Consumer<? super Row> action) {
		if (!iterator.hasNext()) {
			action.accept(iterator.next());
			return true;
		} else {
			return false;
		}
	}

	@Override
	public Spliterator<Row> trySplit() {
		int availableWithoutFetching = rs.getAvailableWithoutFetching();
		if (availableWithoutFetching == 0) {
			if (rs.isFullyFetched()) {
				//done
				return null;
			} else {
				//block for more
				iterator.hasNext();
				availableWithoutFetching = rs.getAvailableWithoutFetching();
			}
		} else {
			System.out.println("CACHE HIT");
		}
		rs.fetchMoreResults();
		estimate -= availableWithoutFetching;
		System.out.println(String.format("Splitting off %d rows", availableWithoutFetching));
		List<Row> cached = new ArrayList<>(availableWithoutFetching);
		for (int i = 0; i < availableWithoutFetching; i++) {
			cached.add(iterator.next());
		}
		return cached.spliterator();
	}

	@Override
	public long estimateSize() {
		return estimate;
	}

	public void forEachRemaining(Consumer<? super Row> action) {
		this.iterator.forEachRemaining(action);
	}

	@Override
	public int characteristics() {
		return Spliterator.IMMUTABLE | Spliterator.NONNULL;
	}

}