@Override
public DynamicSplitResult requestDynamicSplit(DynamicSplitRequest splitRequest) {
checkNotNull(splitRequest);
ApproximateSplitRequest splitProgress = splitRequestToApproximateSplitRequest(splitRequest);
com.google.api.services.dataflow.model.Position cloudPosition = splitProgress.getPosition();
if (cloudPosition == null) {
LOG.warn("Concat only supports split at a Position. Requested: {}", splitRequest);
return null;
}
ConcatPosition concatPosition = cloudPosition.getConcatPosition();
if (concatPosition == null) {
LOG.warn(
"ConcatReader only supports split at a ConcatPosition. Requested: {}", cloudPosition);
return null;
}
if (rangeTracker.trySplitAtPosition(concatPosition.getIndex())) {
com.google.api.services.dataflow.model.Position positionToSplit =
new com.google.api.services.dataflow.model.Position();
positionToSplit.setConcatPosition(
new ConcatPosition().setIndex((int) rangeTracker.getStopPosition().longValue()));
return new DynamicSplitResultWithPosition(cloudPositionToReaderPosition(positionToSplit));
} else {
LOG.debug("Could not perform the dynamic split request " + splitRequest);
return null;
}
}
@Override
protected boolean hasNextImpl() throws IOException {
for (; ; ) {
if (currentIterator != null && currentIterator.hasNext()) {
break;
}
if (currentIterator != null) {
currentIterator.close();
}
currentIteratorIndex++;
if (currentIteratorIndex == sources.size()) {
// All sources were read.
return false;
}
Source currentSource = sources.get(currentIteratorIndex);
try {
@SuppressWarnings("unchecked")
NativeReader<T> currentReader =
(NativeReader<T>)
registry.create(
currentSource, options, executionContext, addCounterMutator, operationName);
currentIterator = (LegacyReaderIterator) currentReader.iterator();
isAtFirstRecordInCurrentSource = true;
} catch (Exception e) {
throw new IOException("Failed to create a reader for source: " + currentSource, e);
}
}
return rangeTracker.tryReturnRecordAt(isAtFirstRecordInCurrentSource, currentIteratorIndex);
}
@Override
public DynamicSplitResult requestDynamicSplit(DynamicSplitRequest splitRequest) {
checkNotNull(splitRequest);
com.google.api.services.dataflow.model.Position splitPosition =
splitRequestToApproximateSplitRequest(splitRequest).getPosition();
if (splitPosition == null) {
LOG.warn("InMemoryReader only supports split at a Position. Requested: {}", splitRequest);
return null;
}
Long splitIndex = splitPosition.getRecordIndex();
if (splitIndex == null) {
LOG.warn(
"InMemoryReader only supports split at a record index. Requested: {}", splitPosition);
return null;
}
if (!tracker.trySplitAtPosition(splitIndex)) {
return null;
}
return new DynamicSplitResultWithPosition(cloudPositionToReaderPosition(splitPosition));
}
@Override
protected boolean hasNextImpl() {
return tracker.tryReturnRecordAt(true, nextIndex);
}
@Override
public double getRemainingParallelism() {
return tracker.getStopPosition() - nextIndex;
}