Examples with PythonKeyedCoProcessOperator org.apache.flink.streaming.api.operators.python.PythonKeyedCoProcessOperator used on opensource projects

Example 1 with PythonKeyedCoProcessOperator

use of org.apache.flink.streaming.api.operators.python.PythonKeyedCoProcessOperator in project flink by apache.

the class PythonOperatorChainingOptimizerTest method testChainingTwoInputOperators.

@Test
public void testChainingTwoInputOperators() {
    PythonKeyedCoProcessOperator<?> keyedCoProcessOperator1 = createCoKeyedProcessOperator("f1", new RowTypeInfo(Types.INT(), Types.STRING()), new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
    PythonProcessOperator<?, ?> processOperator1 = createProcessOperator("f2", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
    PythonProcessOperator<?, ?> processOperator2 = createProcessOperator("f3", new RowTypeInfo(Types.INT(), Types.INT()), Types.LONG());
    PythonKeyedProcessOperator<?> keyedProcessOperator2 = createKeyedProcessOperator("f4", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
    PythonProcessOperator<?, ?> processOperator3 = createProcessOperator("f5", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
    Transformation<?> sourceTransformation1 = mock(SourceTransformation.class);
    Transformation<?> sourceTransformation2 = mock(SourceTransformation.class);
    TwoInputTransformation<?, ?, ?> keyedCoProcessTransformation = new TwoInputTransformation(sourceTransformation1, sourceTransformation2, "keyedCoProcess", keyedCoProcessOperator1, keyedCoProcessOperator1.getProducedType(), 2);
    Transformation<?> processTransformation1 = new OneInputTransformation(keyedCoProcessTransformation, "process", processOperator1, processOperator1.getProducedType(), 2);
    Transformation<?> processTransformation2 = new OneInputTransformation(processTransformation1, "process", processOperator2, processOperator2.getProducedType(), 2);
    OneInputTransformation<?, ?> keyedProcessTransformation = new OneInputTransformation(processTransformation2, "keyedProcess", keyedProcessOperator2, keyedProcessOperator2.getProducedType(), 2);
    Transformation<?> processTransformation3 = new OneInputTransformation(keyedProcessTransformation, "process", processOperator3, processOperator3.getProducedType(), 2);
    List<Transformation<?>> transformations = new ArrayList<>();
    transformations.add(sourceTransformation1);
    transformations.add(sourceTransformation2);
    transformations.add(keyedCoProcessTransformation);
    transformations.add(processTransformation1);
    transformations.add(processTransformation2);
    transformations.add(keyedProcessTransformation);
    transformations.add(processTransformation3);
    List<Transformation<?>> optimized = PythonOperatorChainingOptimizer.optimize(transformations);
    assertEquals(4, optimized.size());
    TwoInputTransformation<?, ?, ?> chainedTransformation1 = (TwoInputTransformation<?, ?, ?>) optimized.get(2);
    assertEquals(sourceTransformation1.getOutputType(), chainedTransformation1.getInputType1());
    assertEquals(sourceTransformation2.getOutputType(), chainedTransformation1.getInputType2());
    assertEquals(processOperator2.getProducedType(), chainedTransformation1.getOutputType());
    OneInputTransformation<?, ?> chainedTransformation2 = (OneInputTransformation<?, ?>) optimized.get(3);
    assertEquals(processOperator2.getProducedType(), chainedTransformation2.getInputType());
    assertEquals(processOperator3.getProducedType(), chainedTransformation2.getOutputType());
    TwoInputStreamOperator<?, ?, ?> chainedOperator1 = chainedTransformation1.getOperator();
    assertTrue(chainedOperator1 instanceof PythonKeyedCoProcessOperator);
    validateChainedPythonFunctions(((PythonKeyedCoProcessOperator<?>) chainedOperator1).getPythonFunctionInfo(), "f3", "f2", "f1");
    OneInputStreamOperator<?, ?> chainedOperator2 = chainedTransformation2.getOperator();
    assertTrue(chainedOperator2 instanceof PythonKeyedProcessOperator);
    validateChainedPythonFunctions(((PythonKeyedProcessOperator<?>) chainedOperator2).getPythonFunctionInfo(), "f5", "f4");
}