@Override public ListenableFuture<SplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, Lifespan lifespan, int maxSize) { ListenableFuture<SplitBatch> batch = splitSource.getNextBatch(partitionHandle, lifespan, maxSize); return Futures.transform(batch, splitBatch -> new SplitBatch( splitBatch.getSplits().stream() .filter(input -> ThreadLocalRandom.current().nextDouble() < sampleRatio) .collect(toImmutableList()), splitBatch.isLastBatch()), directExecutor()); }
private ListenableFuture<?> fetchSplits() { if (splits.size() >= min) { return immediateFuture(null); } ListenableFuture<SplitBatch> future = splitSource.getNextBatch(partitionHandle, lifespan, max - splits.size()); return Futures.transformAsync(future, splitBatch -> { splits.addAll(splitBatch.getSplits()); if (splitBatch.isLastBatch()) { noMoreSplits = true; return immediateFuture(null); } return fetchSplits(); }, directExecutor()); } }
public NextBatchResult assertNoMoreSplits(boolean expectedNoMoreSplits) { assertEquals(splitBatch.isLastBatch(), expectedNoMoreSplits); return this; } }
scheduleGroup.nextSplitBatchFuture = null; pendingSplits.addAll(nextSplits.getSplits()); if (nextSplits.isLastBatch()) { if (scheduleGroup.state == ScheduleGroupState.INITIALIZED && pendingSplits.isEmpty()) {
if (splitBatch.isLastBatch()) { break;