org.apache.flink.state.api.output.TaggedOperatorSubtaskState Java Examples

@Test
public void testOperatorSpecificMaxParallelismRespected() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(4);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.setMaxParallelism(1)
		.transform(new ExampleStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(4);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals("The parallelism of a data set should be constrained my the savepoint max parallelism", 1, getParallelism(result));
}
 

@Test
public void testMaxParallelismRespected() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(10);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(4);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals(
		"The parallelism of a data set should be constrained my the savepoint max parallelism",
		4,
		getParallelism(result));
}
 

@Test
public void testDefaultParallelismRespectedWhenLessThanMaxParallelism() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(4);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(10);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals(
		"The parallelism of a data set should not change when less than the max parallelism of the savepoint",
		ExecutionConfig.PARALLELISM_DEFAULT,
		getParallelism(result));
}
 

@Test
public void testBroadcastStateTransformationParallelism() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(10);

	DataSet<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleBroadcastStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(4);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals("Broadcast transformations should always be run at parallelism 1",
		1,
		getParallelism(result));
}
 

@VisibleForTesting
StreamConfig getConfig(OperatorID operatorID, StateBackend stateBackend, StreamOperator<TaggedOperatorSubtaskState> operator) {
	final StreamConfig config;
	if (keyType == null) {
		config = new BoundedStreamConfig();
	} else {
		TypeSerializer<?> keySerializer = keyType.createSerializer(dataSet.getExecutionEnvironment().getConfig());
		config = new BoundedStreamConfig(keySerializer, originalKeySelector);
	}

	config.setStreamOperator(operator);
	config.setOperatorName(operatorID.toHexString());
	config.setOperatorID(operatorID);
	config.setStateBackend(stateBackend);
	return config;
}
 

@Test
public void testBroadcastStateTransformationParallelism() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(10);

	DataSet<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleBroadcastStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(4);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals("Broadcast transformations should always be run at parallelism 1",
		1,
		getParallelism(result));
}
 

@Test
public void testDefaultParallelismRespectedWhenLessThanMaxParallelism() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(4);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(10);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals(
		"The parallelism of a data set should not change when less than the max parallelism of the savepoint",
		ExecutionConfig.PARALLELISM_DEFAULT,
		getParallelism(result));
}
 

@Test
public void testMaxParallelismRespected() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(10);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(4);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals(
		"The parallelism of a data set should be constrained my the savepoint max parallelism",
		4,
		getParallelism(result));
}
 

@Test
public void testOperatorSpecificMaxParallelismRespected() {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(4);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.setMaxParallelism(1)
		.transform(new ExampleStateBootstrapFunction());

	int maxParallelism = transformation.getMaxParallelism(4);
	DataSet<TaggedOperatorSubtaskState> result = transformation.writeOperatorSubtaskStates(
		OperatorIDGenerator.fromUid("uid"),
		new MemoryStateBackend(),
		new Path(),
		maxParallelism
	);

	Assert.assertEquals("The parallelism of a data set should be constrained my the savepoint max parallelism", 1, getParallelism(result));
}
 

@VisibleForTesting
StreamConfig getConfig(OperatorID operatorID, StateBackend stateBackend, StreamOperator<TaggedOperatorSubtaskState> operator) {
	// Eagerly perform a deep copy of the configuration, otherwise it will result in undefined behavior
	// when deploying with multiple bootstrap transformations.
	Configuration deepCopy = new Configuration(dataSet.getExecutionEnvironment().getConfiguration());
	final StreamConfig config = new StreamConfig(deepCopy);
	config.setChainStart();
	config.setCheckpointingEnabled(true);
	config.setCheckpointMode(CheckpointingMode.EXACTLY_ONCE);

	if (keyType != null) {
		TypeSerializer<?> keySerializer = keyType.createSerializer(dataSet.getExecutionEnvironment().getConfig());

		config.setStateKeySerializer(keySerializer);
		config.setStatePartitioner(0, originalKeySelector);
	}

	config.setStreamOperator(operator);
	config.setOperatorName(operatorID.toHexString());
	config.setOperatorID(operatorID);
	config.setStateBackend(stateBackend);
	return config;
}
 

@Override
public void endInput() throws Exception {
  // bootstrap dataset is now completely processed;
  // take a snapshot of the function states
  final TaggedOperatorSubtaskState state =
      SnapshotUtils.snapshot(
          this,
          getRuntimeContext().getIndexOfThisSubtask(),
          snapshotTimestamp,
          getContainingTask().getCheckpointStorage(),
          snapshotPath);

  output.collect(new StreamRecord<>(state));
}
 

private static <T> int getParallelism(MapPartitionOperator<T, TaggedOperatorSubtaskState> subtaskStates) {
	int parallelism = subtaskStates.getParallelism();
	if (parallelism == ExecutionConfig.PARALLELISM_DEFAULT) {
		parallelism = subtaskStates.getExecutionEnvironment().getParallelism();
	}

	return parallelism;
}
 

@VisibleForTesting
MapPartitionOperator<T, TaggedOperatorSubtaskState> writeOperatorSubtaskStates(
	OperatorID operatorID,
	StateBackend stateBackend,
	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, operator);

	BoundedOneInputStreamTaskRunner<T> operatorRunner = new BoundedOneInputStreamTaskRunner<>(
		config,
		localMaxParallelism);

	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;
}
 

@Override
public void endInput() throws Exception {
	TaggedOperatorSubtaskState state = SnapshotUtils.snapshot(
		this,
		getRuntimeContext().getIndexOfThisSubtask(),
		timestamp,
		getContainingTask().getConfiguration().isExactlyOnceCheckpointMode(),
		getContainingTask().getConfiguration().isUnalignedCheckpointsEnabled(),
		getContainingTask().getCheckpointStorage(),
		savepointPath);

	output.collect(new StreamRecord<>(state));
}
 

@Override
public void endInput() throws Exception {
	TaggedOperatorSubtaskState state = SnapshotUtils.snapshot(
		this,
		getRuntimeContext().getIndexOfThisSubtask(),
		timestamp,
		getContainingTask().getConfiguration().isExactlyOnceCheckpointMode(),
		getContainingTask().getConfiguration().isUnalignedCheckpointsEnabled(),
		getContainingTask().getCheckpointStorage(),
		savepointPath);

	output.collect(new StreamRecord<>(state));
}
 

@Override
public void endInput() throws Exception {
	TaggedOperatorSubtaskState state = SnapshotUtils.snapshot(
		this,
		getRuntimeContext().getIndexOfThisSubtask(),
		timestamp,
		getContainingTask().getConfiguration().isExactlyOnceCheckpointMode(),
		getContainingTask().getConfiguration().isUnalignedCheckpointsEnabled(),
		getContainingTask().getCheckpointStorage(),
		savepointPath);

	output.collect(new StreamRecord<>(state));
}
 

@Override
public void endInput() throws Exception {
  // bootstrap dataset is now completely processed;
  // take a snapshot of the function states
  final TaggedOperatorSubtaskState state =
      SnapshotUtils.snapshot(
          this,
          getRuntimeContext().getIndexOfThisSubtask(),
          snapshotTimestamp,
          getContainingTask().getCheckpointStorage(),
          snapshotPath);

  output.collect(new StreamRecord<>(state));
}
 

private static <T> int getParallelism(MapPartitionOperator<T, TaggedOperatorSubtaskState> subtaskStates) {
	int parallelism = subtaskStates.getParallelism();
	if (parallelism == ExecutionConfig.PARALLELISM_DEFAULT) {
		parallelism = subtaskStates.getExecutionEnvironment().getParallelism();
	}

	return parallelism;
}
 

@VisibleForTesting
MapPartitionOperator<T, TaggedOperatorSubtaskState> writeOperatorSubtaskStates(
	OperatorID operatorID,
	StateBackend stateBackend,
	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, operator);

	BoundedOneInputStreamTaskRunner<T> operatorRunner = new BoundedOneInputStreamTaskRunner<>(
		config,
		localMaxParallelism
	);

	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;
}
 

@Override
public void endInput() throws Exception {
	TaggedOperatorSubtaskState state = SnapshotUtils.snapshot(
		this,
		getRuntimeContext().getIndexOfThisSubtask(),
		timestamp,
		getContainingTask().getCheckpointStorage(),
		savepointPath);

	output.collect(new StreamRecord<>(state));
}
 

@Override
public void endInput() throws Exception {
	TaggedOperatorSubtaskState state = SnapshotUtils.snapshot(
		this,
		getRuntimeContext().getIndexOfThisSubtask(),
		timestamp,
		getContainingTask().getCheckpointStorage(),
		savepointPath);

	output.collect(new StreamRecord<>(state));
}
 

@Override
public void endInput() throws Exception {
	TaggedOperatorSubtaskState state = SnapshotUtils.snapshot(
		this,
		getRuntimeContext().getIndexOfThisSubtask(),
		timestamp,
		getContainingTask().getCheckpointStorage(),
		savepointPath);

	output.collect(new StreamRecord<>(state));
}
 

/**
 * Creates a {@link StreamOperator} to be used for generating and snapshotting state.
 *
 * @param savepointTimestamp the timestamp to associate with the generated savepoint.
 * @param savepointPath the path to write the savepoint to.
 *
 * @return a stream operator for writing the savepoint.
 */
StreamOperator<TaggedOperatorSubtaskState> createOperator(long savepointTimestamp, Path savepointPath);
 

/**
 * Creates a {@link StreamOperator} to be used for generating and snapshotting state.
 *
 * @param savepointTimestamp the timestamp to associate with the generated savepoint.
 * @param savepointPath the path to write the savepoint to.
 *
 * @return a stream operator for writing the savepoint.
 */
StreamOperator<TaggedOperatorSubtaskState> createOperator(long savepointTimestamp, Path savepointPath);