use of org.apache.flink.state.api.output.partitioner.KeyGroupRangePartitioner in project flink by apache.
the class BootstrapTransformation method writeOperatorSubtaskStates.
private MapPartitionOperator<T, TaggedOperatorSubtaskState> writeOperatorSubtaskStates(OperatorID operatorID, @Nullable StateBackend stateBackend, Configuration additionalConfig, Path savepointPath, int localMaxParallelism) {
DataSet<T> input = dataSet;
if (originalKeySelector != null) {
input = dataSet.partitionCustom(new KeyGroupRangePartitioner(localMaxParallelism), hashKeySelector);
}
StreamOperator<TaggedOperatorSubtaskState> operator = factory.createOperator(System.currentTimeMillis(), savepointPath);
operator = dataSet.clean(operator);
final StreamConfig config = getConfig(operatorID, stateBackend, additionalConfig, operator);
BoundedOneInputStreamTaskRunner<T> operatorRunner = new BoundedOneInputStreamTaskRunner<>(config, localMaxParallelism, timestamper);
MapPartitionOperator<T, TaggedOperatorSubtaskState> subtaskStates = input.mapPartition(operatorRunner).name(operatorID.toHexString());
if (operator instanceof BroadcastStateBootstrapOperator) {
subtaskStates = subtaskStates.setParallelism(1);
} else {
int currentParallelism = getParallelism(subtaskStates);
if (currentParallelism > localMaxParallelism) {
subtaskStates.setParallelism(localMaxParallelism);
}
}
return subtaskStates;
}
Aggregations