org.apache.flink.state.api.functions.KeyedStateReaderFunction Java Examples

/**
 * Creates an input format for reading partitioned state from an operator in a savepoint.
 *
 * @param operatorState The state to be queried.
 * @param stateBackend  The state backed used to snapshot the operator.
 * @param keyType       The type information describing the key type.
 * @param userFunction  The {@link KeyedStateReaderFunction} called for each key in the operator.
 */
public KeyedStateInputFormat(
	OperatorState operatorState,
	StateBackend stateBackend,
	TypeInformation<K> keyType,
	KeyedStateReaderFunction<K, OUT> userFunction) {
	Preconditions.checkNotNull(operatorState, "The operator state cannot be null");
	Preconditions.checkNotNull(stateBackend, "The state backend cannot be null");
	Preconditions.checkNotNull(keyType, "The key type information cannot be null");
	Preconditions.checkNotNull(userFunction, "The userfunction cannot be null");

	this.operatorState = operatorState;
	this.stateBackend = stateBackend;
	this.keyType = keyType;
	this.userFunction = userFunction;
}
 

@Nonnull
private List<Integer> readInputSplit(KeyGroupRangeInputSplit split, KeyedStateReaderFunction<Integer, Integer> userFunction) throws IOException {
	KeyedStateInputFormat<Integer, VoidNamespace, Integer> format = new KeyedStateInputFormat<>(
		new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4),
		new MemoryStateBackend(),
		new Configuration(),
		new KeyedStateReaderOperator<>(userFunction, Types.INT));

	List<Integer> data = new ArrayList<>();

	format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));

	format.openInputFormat();
	format.open(split);

	while (!format.reachedEnd()) {
		data.add(format.nextRecord(0));
	}

	format.close();
	format.closeInputFormat();

	data.sort(Comparator.comparingInt(id -> id));
	return data;
}
 

@Test
public void testReadTime() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new KeyedProcessOperator<>(new StatefulFunctionWithTime()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(new TimerReaderFunction(), Types.INT));
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new TimerReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 1, 2, 2, 3, 3), data);
}
 

@Test(expected = IOException.class)
public void testInvalidProcessReaderFunctionFails() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(new ReaderFunction(), Types.INT));
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new InvalidReaderFunction();

	readInputSplit(split, userFunction);

	Assert.fail("KeyedStateReaderFunction did not fail on invalid RuntimeContext use");
}
 

@Test
public void testReadMultipleOutputPerKey() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(new ReaderFunction(), Types.INT));
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new DoubleReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 1, 2, 2, 3, 3), data);
}
 

@Test
public void testReadState() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(new ReaderFunction(), Types.INT));
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new ReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 2, 3), data);
}
 

/**
 * Read keyed state from an operator in a {@code Savepoint}.
 * @param uid The uid of the operator.
 * @param function The {@link KeyedStateReaderFunction} that is called for each key in state.
 * @param keyTypeInfo The type information of the key in state.
 * @param outTypeInfo The type information of the output of the transform reader function.
 * @param <K> The type of the key in state.
 * @param <OUT> The output type of the transform function.
 * @return A {@code DataSet} of objects read from keyed state.
 * @throws IOException If the savepoint does not contain operator state with the given uid.
 */
public <K, OUT> DataSet<OUT> readKeyedState(
	String uid,
	KeyedStateReaderFunction<K, OUT> function,
	TypeInformation<K> keyTypeInfo,
	TypeInformation<OUT> outTypeInfo) throws IOException {

	OperatorState operatorState = metadata.getOperatorState(uid);
	KeyedStateInputFormat<K, VoidNamespace, OUT> inputFormat = new KeyedStateInputFormat<>(
		operatorState,
		stateBackend,
		env.getConfiguration(),
		new KeyedStateReaderOperator<>(function, keyTypeInfo));

	return env.createInput(inputFormat, outTypeInfo);
}
 

@Nonnull
private List<Integer> readInputSplit(KeyGroupRangeInputSplit split, KeyedStateReaderFunction<Integer, Integer> userFunction) throws IOException {
	KeyedStateInputFormat<Integer, Integer> format = new KeyedStateInputFormat<>(
		new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4),
		new MemoryStateBackend(),
		Types.INT,
		userFunction);

	List<Integer> data = new ArrayList<>();

	format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));

	format.openInputFormat();
	format.open(split);

	while (!format.reachedEnd()) {
		data.add(format.nextRecord(0));
	}

	format.close();
	format.closeInputFormat();

	data.sort(Comparator.comparingInt(id -> id));
	return data;
}
 

@Test
public void testReadTime() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new KeyedProcessOperator<>(new StatefulFunctionWithTime()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new TimeReaderFunction());
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new TimeReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 1, 2, 2, 3, 3), data);
}
 

@Test(expected = IOException.class)
public void testInvalidProcessReaderFunctionFails() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new ReaderFunction());
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new InvalidReaderFunction();

	readInputSplit(split, userFunction);

	Assert.fail("KeyedStateReaderFunction did not fail on invalid RuntimeContext use");
}
 

@Test
public void testReadMultipleOutputPerKey() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new ReaderFunction());
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new DoubleReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 1, 2, 2, 3, 3), data);
}
 

@Test
public void testReadState() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new ReaderFunction());
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new ReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 2, 3), data);
}
 

/**
 * Read keyed state from an operator in a {@code Savepoint}.
 * @param uid The uid of the operator.
 * @param function The {@link KeyedStateReaderFunction} that is called for each key in state.
 * @param keyTypeInfo The type information of the key in state.
 * @param outTypeInfo The type information of the output of the transform reader function.
 * @param <K> The type of the key in state.
 * @param <OUT> The output type of the transform function.
 * @return A {@code DataSet} of objects read from keyed state.
 */
public <K, OUT> DataSet<OUT> readKeyedState(
	String uid,
	KeyedStateReaderFunction<K, OUT> function,
	TypeInformation<K> keyTypeInfo,
	TypeInformation<OUT> outTypeInfo) throws IOException {

	OperatorState operatorState = metadata.getOperatorState(uid);
	KeyedStateInputFormat<K, OUT> inputFormat = new KeyedStateInputFormat<>(
		operatorState,
		stateBackend,
		keyTypeInfo,
		function);

	return env.createInput(inputFormat, outTypeInfo);
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	Set<Long> eventTimers = ctx.registeredEventTimeTimers();
	Assert.assertEquals("Each key should have exactly one event timer for key " + key, 1, eventTimers.size());

	out.collect(eventTimers.iterator().next().intValue());

	Set<Long> procTimers = ctx.registeredProcessingTimeTimers();
	Assert.assertEquals("Each key should have exactly one processing timer for key " + key, 1, procTimers.size());

	out.collect(procTimers.iterator().next().intValue());
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	Set<Long> eventTimers = ctx.registeredEventTimeTimers();
	Assert.assertEquals("Each key should have exactly one event timer for key " + key, 1, eventTimers.size());

	out.collect(eventTimers.iterator().next().intValue());

	Set<Long> procTimers = ctx.registeredProcessingTimeTimers();
	Assert.assertEquals("Each key should have exactly one processing timer for key " + key, 1, procTimers.size());

	out.collect(procTimers.iterator().next().intValue());
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	out.collect(state.value());
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	out.collect(state.value());
	out.collect(state.value());
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	ValueState<Integer> state = getRuntimeContext().getState(stateDescriptor);
	out.collect(state.value());
}
 

public KeyedStateReaderOperator(KeyedStateReaderFunction<KEY, OUT> function, TypeInformation<KEY> keyType) {
	super(function, keyType, VoidNamespaceSerializer.INSTANCE);
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	out.collect(state.value());
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	out.collect(state.value());
	out.collect(state.value());
}
 

@Override
public void readKey(Integer key, KeyedStateReaderFunction.Context ctx, Collector<Integer> out) throws Exception {
	ValueState<Integer> state = getRuntimeContext().getState(stateDescriptor);
	out.collect(state.value());
}