@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 assertSize(int expectedSize) { assertEquals(splitBatch.getSplits().size(), expectedSize); return this; }
private static List<Split> getNextBatch(SplitSource splitSource) { return getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000)).getSplits(); }
SplitBatch nextSplits = getFutureValue(scheduleGroup.nextSplitBatchFuture); scheduleGroup.nextSplitBatchFuture = null; pendingSplits.addAll(nextSplits.getSplits()); if (nextSplits.isLastBatch()) { if (scheduleGroup.state == ScheduleGroupState.INITIALIZED && pendingSplits.isEmpty()) {
while (!Thread.currentThread().isInterrupted()) { SplitBatch splitBatch = getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000)); List<Split> splits = splitBatch.getSplits();
private static List<Split> getNextBatch(SplitSource splitSource) { return getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000)).getSplits(); }
private static List<Split> getNextBatch(SplitSource splitSource) { return getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000)).getSplits(); }