org.apache.flink.core.io.InputSplitAssigner Java Examples

private InputSplitAssigner getAssigner() {
	InputSplit[] inputSplits = new InputSplit[NUM_SPLITS];
	int[] taskInputSplitSize = new int[NUM_TASKS];
	int[] taskInputSplitStartIndex = new int[NUM_TASKS];
	for (int i = 0; i < NUM_SPLITS; i++) {
		inputSplits[i] = new MockInputSplit(i);
	}
	int currentSplitIndex = 0;
	for (int i = 0; i < NUM_TASKS; i++) {
		taskInputSplitSize[i] = i + 1;
		taskInputSplitStartIndex[i] = currentSplitIndex;
		currentSplitIndex += taskInputSplitSize[i];
	}
	return new AbstractParallelSourceBase.PreAssignedInputSplitAssigner(
			inputSplits, taskInputSplitSize, taskInputSplitStartIndex);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
	LocatableInputSplit[] locatableInputSplits = new LocatableInputSplit[inputSplits.length];
	for (int i = 0; i < inputSplits.length; i++) {
		locatableInputSplits[i] = (LocatableInputSplit) inputSplits[i];
	}
	if (source.isParallelReadDisabled()) {
		return new LocatableInputSplitAssigner(locatableInputSplits);
	}
	return super.getInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(final SourceInputSplit[] sourceInputSplits) {
	return new InputSplitAssigner() {
		private int index = 0;
		private final SourceInputSplit[] splits = sourceInputSplits;
		@Override
		public InputSplit getNextInputSplit(String host, int taskId) {
			if (index < splits.length) {
				return splits[index++];
			} else {
				return null;
			}
		}
	};
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
	if (!source.initInputSplitInMaster || source.disableParallelRead) {
		return new DefaultInputSplitAssigner(inputSplits);
	} else {
		return new PreAssignedInputSplitAssigner(inputSplits, taskInputSplitSize, taskInputSplitStartIndex);
	}
}
 

@Test
public void testReturnSplits() {
	InputSplitAssigner assigner = getAssigner();
	assertEquals(3, assigner.getNextInputSplit("", 2).getSplitNumber());
	assertEquals(4, assigner.getNextInputSplit("", 2).getSplitNumber());

	// Return split 3 back.
	assigner.returnInputSplit(Collections.singletonList(new MockInputSplit(3)), 2);

	assertEquals(3, assigner.getNextInputSplit("", 2).getSplitNumber());
	assertEquals(5, assigner.getNextInputSplit("", 2).getSplitNumber());
	assertNull(assigner.getNextInputSplit("", 2));
}
 

@Test
public void testInitializeAndAssign() {
	InputSplitAssigner assigner = getAssigner();
	assertEquals(0, assigner.getNextInputSplit("", 0).getSplitNumber());
	assertNull(assigner.getNextInputSplit("", 0));
	assertEquals(1, assigner.getNextInputSplit("", 1).getSplitNumber());
	assertEquals(2, assigner.getNextInputSplit("", 1).getSplitNumber());
	assertNull(assigner.getNextInputSplit("", 1));
	assertEquals(3, assigner.getNextInputSplit("", 2).getSplitNumber());
	assertEquals(4, assigner.getNextInputSplit("", 2).getSplitNumber());
	assertEquals(5, assigner.getNextInputSplit("", 2).getSplitNumber());
	assertNull(assigner.getNextInputSplit("", 2));
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(GenericInputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
	return null;
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(GenericInputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

private Execution resetForNewExecutionInternal(final long timestamp, final long originatingGlobalModVersion) {
	final Execution oldExecution = currentExecution;
	final ExecutionState oldState = oldExecution.getState();

	if (oldState.isTerminal()) {
		if (oldState == FINISHED) {
			// pipelined partitions are released in Execution#cancel(), covering both job failures and vertex resets
			// do not release pipelined partitions here to save RPC calls
			oldExecution.handlePartitionCleanup(false, true);
			getExecutionGraph().getPartitionReleaseStrategy().vertexUnfinished(executionVertexId);
		}

		priorExecutions.add(oldExecution.archive());

		final Execution newExecution = new Execution(
			getExecutionGraph().getFutureExecutor(),
			this,
			oldExecution.getAttemptNumber() + 1,
			originatingGlobalModVersion,
			timestamp,
			timeout);

		currentExecution = newExecution;

		synchronized (inputSplits) {
			InputSplitAssigner assigner = jobVertex.getSplitAssigner();
			if (assigner != null) {
				assigner.returnInputSplit(inputSplits, getParallelSubtaskIndex());
				inputSplits.clear();
			}
		}

		CoLocationGroup grp = jobVertex.getCoLocationGroup();
		if (grp != null) {
			locationConstraint = grp.getLocationConstraint(subTaskIndex);
		}

		// register this execution at the execution graph, to receive call backs
		getExecutionGraph().registerExecution(newExecution);

		// if the execution was 'FINISHED' before, tell the ExecutionGraph that
		// we take one step back on the road to reaching global FINISHED
		if (oldState == FINISHED) {
			getExecutionGraph().vertexUnFinished();
		}

		// reset the intermediate results
		for (IntermediateResultPartition resultPartition : resultPartitions.values()) {
			resultPartition.resetForNewExecution();
		}

		return newExecution;
	}
	else {
		throw new IllegalStateException("Cannot reset a vertex that is in non-terminal state " + oldState);
	}
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(KeyGroupRangeInputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(OperatorStateInputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(S[] inputSplits) {
	return new ReplicatingInputSplitAssigner(inputSplits);
}
 

public InputSplitAssigner getSplitAssigner() {
	return splitAssigner;
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
InputSplitAssigner getInputSplitAssigner(T[] inputSplits);
 

@Override
public InputSplitAssigner getInputSplitAssigner(HadoopInputSplit[] inputSplits) {
	return new LocatableInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(GenericInputSplit[] genericInputSplits) {
  return new DefaultInputSplitAssigner(genericInputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(final SourceInputSplit[] sourceInputSplits) {
  return new DefaultInputSplitAssigner(sourceInputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(KuduInputSplit[] inputSplits) {
    return new LocatableInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(PravegaInputSplit[] inputSplits) {
    return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(TweetFileInputSplit[] inputSplits) {
    return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(CsvFileInputSplit[] inputSplits) {
    return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(KuduInputSplit[] inputSplits) {
    return new LocatableInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(KeyedStateInputSplit[] inputSplits) {
	return new DefaultInputSplitAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(FileCopyTaskInputSplit[] inputSplits) {
	return new FileCopyTaskAssigner(inputSplits);
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
	return null;
}
 

@Override
public InputSplitAssigner getInputSplitAssigner(FileCopyTaskInputSplit[] inputSplits) {
	return new FileCopyTaskAssigner(inputSplits);
}