org.apache.flink.streaming.api.operators.KeyedProcessOperator Java Examples

/**
 * Returns an initialized test harness for {@link KeyedProcessFunction}.
 *
 * @param function instance of a {@link KeyedCoProcessFunction} under test
 * @param <K> key type
 * @param <IN> type of input stream elements
 * @param <OUT> type of output stream elements
 * @return {@link KeyedOneInputStreamOperatorTestHarness} wrapped around {@code function}
 */
public static <K, IN, OUT>
KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> forKeyedProcessFunction(
	final KeyedProcessFunction<K, IN, OUT> function,
	final KeySelector<IN, K> keySelector,
	final TypeInformation<K> keyType) throws Exception {

	KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness =
		new KeyedOneInputStreamOperatorTestHarness<>(
			new KeyedProcessOperator<>(
				Preconditions.checkNotNull(function)),
			keySelector,
			keyType,
			1,
			1,
			0);
	testHarness.open();
	return testHarness;
}
 

@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
public void testTimerStateRestorable() throws Exception {
	Path path = new Path(folder.newFolder().toURI());

	OperatorSubtaskState state;
	KeyedStateBootstrapOperator<Long, Long> bootstrapOperator = new KeyedStateBootstrapOperator<>(0L, path, new TimerBootstrapFunction());
	try (KeyedOneInputStreamOperatorTestHarness<Long, Long, TaggedOperatorSubtaskState> harness = getHarness(bootstrapOperator)) {
		processElements(harness, 1L, 2L, 3L);
		state = getState(bootstrapOperator, harness);
	}

	KeyedProcessOperator<Long, Long, Tuple3<Long, Long, TimeDomain>> procOperator = new KeyedProcessOperator<>(new SimpleProcessFunction());
	try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Tuple3<Long, Long, TimeDomain>> harness = getHarness(procOperator, state)) {
		harness.processWatermark(EVENT_TIMER);
		harness.setProcessingTime(PROC_TIMER);

		assertHarnessOutput(harness,
			Tuple3.of(1L, EVENT_TIMER, TimeDomain.EVENT_TIME),
			Tuple3.of(2L, EVENT_TIMER, TimeDomain.EVENT_TIME),
			Tuple3.of(3L, EVENT_TIMER, TimeDomain.EVENT_TIME),
			Tuple3.of(1L, PROC_TIMER, TimeDomain.PROCESSING_TIME),
			Tuple3.of(2L, PROC_TIMER, TimeDomain.PROCESSING_TIME),
			Tuple3.of(3L, PROC_TIMER, TimeDomain.PROCESSING_TIME));
	}
}
 

@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
public void testStateCleanup() throws Exception {
	ProcTimeRangeBoundedPrecedingFunction<RowData> function = new ProcTimeRangeBoundedPrecedingFunction<>(
		aggsHandleFunction, accTypes, inputFieldTypes, 2000);
	KeyedProcessOperator<RowData, RowData, RowData> operator = new KeyedProcessOperator<>(function);

	OneInputStreamOperatorTestHarness<RowData, RowData> testHarness = createTestHarness(operator);

	testHarness.open();

	AbstractKeyedStateBackend stateBackend = (AbstractKeyedStateBackend) operator.getKeyedStateBackend();

	assertEquals("Initial state is not empty", 0, stateBackend.numKeyValueStateEntries());

	// put some records
	testHarness.setProcessingTime(100);
	testHarness.processElement(insertRecord("key", 1L));
	testHarness.processElement(insertRecord("key", 1L));
	testHarness.setProcessingTime(500);
	testHarness.processElement(insertRecord("key", 1L));

	testHarness.setProcessingTime(1000);
	// at this moment we expect the function to have some records in state

	testHarness.setProcessingTime(4000);
	// at this moment the function should have cleaned up states

	assertEquals("State has not been cleaned up", 0, stateBackend.numKeyValueStateEntries());
}
 

private KeyedOneInputStreamOperatorTestHarness<String, String, Integer> getHarnessTest(
	KeySelector<String, String> keySelector,
	int maxParallelism,
	int taskParallelism,
	int subtaskIdx) throws Exception {
	return new KeyedOneInputStreamOperatorTestHarness<>(
		new KeyedProcessOperator<>(new TestKeyedFunction()),
		keySelector,
		BasicTypeInfo.STRING_TYPE_INFO,
		maxParallelism,
		taskParallelism,
		subtaskIdx);
}
 

OneInputStreamOperatorTestHarness<RowData, RowData> createTestHarness(
		AbstractTopNFunction rankFunction)
		throws Exception {
	KeyedProcessOperator<RowData, RowData, RowData> operator = new KeyedProcessOperator<>(rankFunction);
	rankFunction.setKeyContext(operator);
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, keySelector.getProducedType());
}
 

private KeyedOneInputStreamOperatorTestHarness<String, String, Integer> getHarnessTest(
	KeySelector<String, String> keySelector,
	int maxParallelism,
	int taskParallelism,
	int subtaskIdx) throws Exception {
	return new KeyedOneInputStreamOperatorTestHarness<>(
		new KeyedProcessOperator<>(new TestKeyedFunction()),
		keySelector,
		BasicTypeInfo.STRING_TYPE_INFO,
		maxParallelism,
		taskParallelism,
		subtaskIdx);
}
 

@Test
public void testStateCleanup() throws Exception {
	RowTimeRangeBoundedPrecedingFunction<RowData> function = new RowTimeRangeBoundedPrecedingFunction<>(
		aggsHandleFunction, accTypes, inputFieldTypes, 2000, 2);
	KeyedProcessOperator<RowData, RowData, RowData> operator = new KeyedProcessOperator<>(function);

	OneInputStreamOperatorTestHarness<RowData, RowData> testHarness = createTestHarness(operator);

	testHarness.open();

	AbstractKeyedStateBackend stateBackend = (AbstractKeyedStateBackend) operator.getKeyedStateBackend();

	assertEquals("Initial state is not empty", 0, stateBackend.numKeyValueStateEntries());

	// put some records
	testHarness.processElement(insertRecord("key", 1L, 100L));
	testHarness.processElement(insertRecord("key", 1L, 100L));
	testHarness.processElement(insertRecord("key", 1L, 500L));

	testHarness.processWatermark(new Watermark(1000L));
	// at this moment we expect the function to have some records in state

	testHarness.processWatermark(new Watermark(4000L));
	// at this moment the function should have cleaned up states

	assertEquals("State has not been cleaned up", 0, stateBackend.numKeyValueStateEntries());
}
 

private KeyedOneInputStreamOperatorTestHarness<String, String, Integer> getHarnessTest(
	KeySelector<String, String> keySelector,
	int maxParallelism,
	int taskParallelism,
	int subtaskIdx) throws Exception {
	return new KeyedOneInputStreamOperatorTestHarness<>(
		new KeyedProcessOperator<>(new TestKeyedFunction()),
		keySelector,
		BasicTypeInfo.STRING_TYPE_INFO,
		maxParallelism,
		taskParallelism,
		subtaskIdx);
}
 

OneInputStreamOperatorTestHarness<BaseRow, BaseRow> createTestHarness(
		AbstractTopNFunction rankFunction)
		throws Exception {
	KeyedProcessOperator<BaseRow, BaseRow, BaseRow> operator = new KeyedProcessOperator<>(rankFunction);
	rankFunction.setKeyContext(operator);
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, keySelector.getProducedType());
}
 

private OneInputStreamOperatorTestHarness<RowData, RowData> createTestHarness(
		KeyedProcessOperator<RowData, RowData, RowData> operator) throws Exception {
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, keyType);
}
 

private OneInputStreamOperatorTestHarness<RowData, RowData> createTestHarness(
		KeyedProcessOperator<RowData, RowData, RowData> operator) throws Exception {
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, keyType);
}
 

private OneInputStreamOperatorTestHarness<RowData, RowData> createTestHarness(
		DeduplicateKeepLastRowFunction func)
		throws Exception {
	KeyedProcessOperator<RowData, RowData, RowData> operator = new KeyedProcessOperator<>(func);
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, rowKeySelector, rowKeySelector.getProducedType());
}
 

private OneInputStreamOperatorTestHarness<RowData, RowData> createTestHarness(
		DeduplicateKeepFirstRowFunction func)
		throws Exception {
	KeyedProcessOperator<RowData, RowData, RowData> operator = new KeyedProcessOperator<>(func);
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, rowKeySelector, rowKeySelector.getProducedType());
}
 

private OneInputStreamOperatorTestHarness<BaseRow, BaseRow> createTestHarness(
		DeduplicateKeepFirstRowFunction func)
		throws Exception {
	KeyedProcessOperator<BaseRow, BaseRow, BaseRow> operator = new KeyedProcessOperator<>(func);
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, rowKeySelector, rowKeySelector.getProducedType());
}
 

private OneInputStreamOperatorTestHarness<BaseRow, BaseRow> createTestHarness(
		DeduplicateKeepLastRowFunction func)
		throws Exception {
	KeyedProcessOperator<BaseRow, BaseRow, BaseRow> operator = new KeyedProcessOperator<>(func);
	return new KeyedOneInputStreamOperatorTestHarness<>(operator, rowKeySelector, rowKeySelector.getProducedType());
}
 

/**
 * Applies the given {@link KeyedProcessFunction} on the input stream, thereby creating a transformed output stream.
 *
 * <p>The function will be called for every element in the input streams and can produce zero
 * or more output elements. Contrary to the {@link DataStream#flatMap(FlatMapFunction)}
 * function, this function can also query the time and set timers. When reacting to the firing
 * of set timers the function can directly emit elements and/or register yet more timers.
 *
 * @param keyedProcessFunction The {@link KeyedProcessFunction} that is called for each element in the stream.
 *
 * @param outputType {@link TypeInformation} for the result type of the function.
 *
 * @param <R> The type of elements emitted by the {@code KeyedProcessFunction}.
 *
 * @return The transformed {@link DataStream}.
 */
@Internal
public <R> SingleOutputStreamOperator<R> process(
		KeyedProcessFunction<KEY, T, R> keyedProcessFunction,
		TypeInformation<R> outputType) {

	KeyedProcessOperator<KEY, T, R> operator = new KeyedProcessOperator<>(clean(keyedProcessFunction));
	return transform("KeyedProcess", outputType, operator);
}
 

/**
 * Applies the given {@link KeyedProcessFunction} on the input stream, thereby creating a transformed output stream.
 *
 * <p>The function will be called for every element in the input streams and can produce zero
 * or more output elements. Contrary to the {@link DataStream#flatMap(FlatMapFunction)}
 * function, this function can also query the time and set timers. When reacting to the firing
 * of set timers the function can directly emit elements and/or register yet more timers.
 *
 * @param keyedProcessFunction The {@link KeyedProcessFunction} that is called for each element in the stream.
 *
 * @param outputType {@link TypeInformation} for the result type of the function.
 *
 * @param <R> The type of elements emitted by the {@code KeyedProcessFunction}.
 *
 * @return The transformed {@link DataStream}.
 */
@Internal
public <R> SingleOutputStreamOperator<R> process(
		KeyedProcessFunction<KEY, T, R> keyedProcessFunction,
		TypeInformation<R> outputType) {

	KeyedProcessOperator<KEY, T, R> operator = new KeyedProcessOperator<>(clean(keyedProcessFunction));
	return transform("KeyedProcess", outputType, operator);
}
 

/**
 * Applies the given {@link KeyedProcessFunction} on the input stream, thereby creating a transformed output stream.
 *
 * <p>The function will be called for every element in the input streams and can produce zero
 * or more output elements. Contrary to the {@link DataStream#flatMap(FlatMapFunction)}
 * function, this function can also query the time and set timers. When reacting to the firing
 * of set timers the function can directly emit elements and/or register yet more timers.
 *
 * @param keyedProcessFunction The {@link KeyedProcessFunction} that is called for each element in the stream.
 *
 * @param outputType {@link TypeInformation} for the result type of the function.
 *
 * @param <R> The type of elements emitted by the {@code KeyedProcessFunction}.
 *
 * @return The transformed {@link DataStream}.
 */
@Internal
public <R> SingleOutputStreamOperator<R> process(
		KeyedProcessFunction<KEY, T, R> keyedProcessFunction,
		TypeInformation<R> outputType) {

	KeyedProcessOperator<KEY, T, R> operator = new KeyedProcessOperator<>(clean(keyedProcessFunction));
	return transform("KeyedProcess", outputType, operator);
}