Examples with RuntimeUDFContext org.apache.flink.api.common.functions.util.RuntimeUDFContext used on opensource projects

Example 1 with RuntimeUDFContext

use of org.apache.flink.api.common.functions.util.RuntimeUDFContext in project flink by apache.

the class ReduceOperatorTest method testReduceCollectionWithRuntimeContext.

@Test
public void testReduceCollectionWithRuntimeContext() {
    try {
        final String taskName = "Test Task";
        final AtomicBoolean opened = new AtomicBoolean();
        final AtomicBoolean closed = new AtomicBoolean();
        final ReduceFunction<Tuple2<String, Integer>> reducer = new RichReduceFunction<Tuple2<String, Integer>>() {

            @Override
            public Tuple2<String, Integer> reduce(Tuple2<String, Integer> value1, Tuple2<String, Integer> value2) throws Exception {
                return new Tuple2<String, Integer>(value1.f0, value1.f1 + value2.f1);
            }

            @Override
            public void open(Configuration parameters) throws Exception {
                opened.set(true);
                RuntimeContext ctx = getRuntimeContext();
                assertEquals(0, ctx.getIndexOfThisSubtask());
                assertEquals(1, ctx.getNumberOfParallelSubtasks());
                assertEquals(taskName, ctx.getTaskName());
            }

            @Override
            public void close() throws Exception {
                closed.set(true);
            }
        };
        ReduceOperatorBase<Tuple2<String, Integer>, ReduceFunction<Tuple2<String, Integer>>> op = new ReduceOperatorBase<Tuple2<String, Integer>, ReduceFunction<Tuple2<String, Integer>>>(reducer, new UnaryOperatorInformation<Tuple2<String, Integer>, Tuple2<String, Integer>>(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, " + "Integer>")), new int[] { 0 }, "TestReducer");
        List<Tuple2<String, Integer>> input = new ArrayList<Tuple2<String, Integer>>(asList(new Tuple2<String, Integer>("foo", 1), new Tuple2<String, Integer>("foo", 3), new Tuple2<String, Integer>("bar", 2), new Tuple2<String, Integer>("bar", 4)));
        final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
        ExecutionConfig executionConfig = new ExecutionConfig();
        executionConfig.disableObjectReuse();
        List<Tuple2<String, Integer>> resultMutableSafe = op.executeOnCollections(input, new RuntimeUDFContext(taskInfo, null, executionConfig, new HashMap<String, Future<Path>>(), new HashMap<String, Accumulator<?, ?>>(), new UnregisteredMetricsGroup()), executionConfig);
        executionConfig.enableObjectReuse();
        List<Tuple2<String, Integer>> resultRegular = op.executeOnCollections(input, new RuntimeUDFContext(taskInfo, null, executionConfig, new HashMap<String, Future<Path>>(), new HashMap<String, Accumulator<?, ?>>(), new UnregisteredMetricsGroup()), executionConfig);
        Set<Tuple2<String, Integer>> resultSetMutableSafe = new HashSet<Tuple2<String, Integer>>(resultMutableSafe);
        Set<Tuple2<String, Integer>> resultSetRegular = new HashSet<Tuple2<String, Integer>>(resultRegular);
        Set<Tuple2<String, Integer>> expectedResult = new HashSet<Tuple2<String, Integer>>(asList(new Tuple2<String, Integer>("foo", 4), new Tuple2<String, Integer>("bar", 6)));
        assertEquals(expectedResult, resultSetMutableSafe);
        assertEquals(expectedResult, resultSetRegular);
        assertTrue(opened.get());
        assertTrue(closed.get());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 2 with RuntimeUDFContext

use of org.apache.flink.api.common.functions.util.RuntimeUDFContext in project flink by apache.

the class CollectionExecutor method executeUnaryOperator.

private <IN, OUT> List<OUT> executeUnaryOperator(SingleInputOperator<?, ?, ?> operator, int superStep) throws Exception {
    Operator<?> inputOp = operator.getInput();
    if (inputOp == null) {
        throw new InvalidProgramException("The unary operation " + operator.getName() + " has no input.");
    }
    @SuppressWarnings("unchecked") List<IN> inputData = (List<IN>) execute(inputOp, superStep);
    @SuppressWarnings("unchecked") SingleInputOperator<IN, OUT, ?> typedOp = (SingleInputOperator<IN, OUT, ?>) operator;
    // build the runtime context and compute broadcast variables, if necessary
    TaskInfo taskInfo = new TaskInfo(typedOp.getName(), 1, 0, 1, 0);
    RuntimeUDFContext ctx;
    MetricGroup metrics = new UnregisteredMetricsGroup();
    if (RichFunction.class.isAssignableFrom(typedOp.getUserCodeWrapper().getUserCodeClass())) {
        ctx = superStep == 0 ? new RuntimeUDFContext(taskInfo, classLoader, executionConfig, cachedFiles, accumulators, metrics) : new IterationRuntimeUDFContext(taskInfo, classLoader, executionConfig, cachedFiles, accumulators, metrics);
        for (Map.Entry<String, Operator<?>> bcInputs : operator.getBroadcastInputs().entrySet()) {
            List<?> bcData = execute(bcInputs.getValue());
            ctx.setBroadcastVariable(bcInputs.getKey(), bcData);
        }
    } else {
        ctx = null;
    }
    return typedOp.executeOnCollections(inputData, ctx, executionConfig);
}

Example 3 with RuntimeUDFContext

use of org.apache.flink.api.common.functions.util.RuntimeUDFContext in project flink by apache.

the class InnerJoinOperatorBaseTest method testJoinRich.

@Test
public void testJoinRich() {
    final AtomicBoolean opened = new AtomicBoolean(false);
    final AtomicBoolean closed = new AtomicBoolean(false);
    final String taskName = "Test rich join function";
    final RichFlatJoinFunction<String, String, Integer> joiner = new RichFlatJoinFunction<String, String, Integer>() {

        @Override
        public void open(Configuration parameters) throws Exception {
            opened.compareAndSet(false, true);
            assertEquals(0, getRuntimeContext().getIndexOfThisSubtask());
            assertEquals(1, getRuntimeContext().getNumberOfParallelSubtasks());
        }

        @Override
        public void close() throws Exception {
            closed.compareAndSet(false, true);
        }

        @Override
        public void join(String first, String second, Collector<Integer> out) throws Exception {
            out.collect(first.length());
            out.collect(second.length());
        }
    };
    InnerJoinOperatorBase<String, String, Integer, RichFlatJoinFunction<String, String, Integer>> base = new InnerJoinOperatorBase<String, String, Integer, RichFlatJoinFunction<String, String, Integer>>(joiner, new BinaryOperatorInformation<String, String, Integer>(BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO), new int[0], new int[0], taskName);
    final List<String> inputData1 = new ArrayList<String>(Arrays.asList("foo", "bar", "foobar"));
    final List<String> inputData2 = new ArrayList<String>(Arrays.asList("foobar", "foo"));
    final List<Integer> expected = new ArrayList<Integer>(Arrays.asList(3, 3, 6, 6));
    try {
        final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
        final HashMap<String, Accumulator<?, ?>> accumulatorMap = new HashMap<String, Accumulator<?, ?>>();
        final HashMap<String, Future<Path>> cpTasks = new HashMap<>();
        ExecutionConfig executionConfig = new ExecutionConfig();
        executionConfig.disableObjectReuse();
        List<Integer> resultSafe = base.executeOnCollections(inputData1, inputData2, new RuntimeUDFContext(taskInfo, null, executionConfig, cpTasks, accumulatorMap, new UnregisteredMetricsGroup()), executionConfig);
        executionConfig.enableObjectReuse();
        List<Integer> resultRegular = base.executeOnCollections(inputData1, inputData2, new RuntimeUDFContext(taskInfo, null, executionConfig, cpTasks, accumulatorMap, new UnregisteredMetricsGroup()), executionConfig);
        assertEquals(expected, resultSafe);
        assertEquals(expected, resultRegular);
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
    assertTrue(opened.get());
    assertTrue(closed.get());
}

Example 4 with RuntimeUDFContext

use of org.apache.flink.api.common.functions.util.RuntimeUDFContext in project flink by apache.

the class OuterJoinOperatorBaseTest method setup.

@SuppressWarnings({ "rawtypes", "unchecked" })
@Before
public void setup() {
    joiner = new MockRichFlatJoinFunction();
    baseOperator = new OuterJoinOperatorBase(joiner, new BinaryOperatorInformation(BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO), new int[0], new int[0], "TestJoiner", null);
    executionConfig = new ExecutionConfig();
    String taskName = "Test rich outer join function";
    TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
    HashMap<String, Accumulator<?, ?>> accumulatorMap = new HashMap<>();
    HashMap<String, Future<Path>> cpTasks = new HashMap<>();
    runtimeContext = new RuntimeUDFContext(taskInfo, null, executionConfig, cpTasks, accumulatorMap, new UnregisteredMetricsGroup());
}

Example 5 with RuntimeUDFContext

use of org.apache.flink.api.common.functions.util.RuntimeUDFContext in project flink by apache.

the class PartitionMapOperatorTest method testMapPartitionWithRuntimeContext.

@Test
public void testMapPartitionWithRuntimeContext() {
    try {
        final String taskName = "Test Task";
        final AtomicBoolean opened = new AtomicBoolean();
        final AtomicBoolean closed = new AtomicBoolean();
        final MapPartitionFunction<String, Integer> parser = new RichMapPartitionFunction<String, Integer>() {

            @Override
            public void open(Configuration parameters) throws Exception {
                opened.set(true);
                RuntimeContext ctx = getRuntimeContext();
                assertEquals(0, ctx.getIndexOfThisSubtask());
                assertEquals(1, ctx.getNumberOfParallelSubtasks());
                assertEquals(taskName, ctx.getTaskName());
            }

            @Override
            public void mapPartition(Iterable<String> values, Collector<Integer> out) {
                for (String s : values) {
                    out.collect(Integer.parseInt(s));
                }
            }

            @Override
            public void close() throws Exception {
                closed.set(true);
            }
        };
        MapPartitionOperatorBase<String, Integer, MapPartitionFunction<String, Integer>> op = new MapPartitionOperatorBase<String, Integer, MapPartitionFunction<String, Integer>>(parser, new UnaryOperatorInformation<String, Integer>(BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO), taskName);
        List<String> input = new ArrayList<String>(asList("1", "2", "3", "4", "5", "6"));
        final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
        ExecutionConfig executionConfig = new ExecutionConfig();
        executionConfig.disableObjectReuse();
        List<Integer> resultMutableSafe = op.executeOnCollections(input, new RuntimeUDFContext(taskInfo, null, executionConfig, new HashMap<String, Future<Path>>(), new HashMap<String, Accumulator<?, ?>>(), new UnregisteredMetricsGroup()), executionConfig);
        executionConfig.enableObjectReuse();
        List<Integer> resultRegular = op.executeOnCollections(input, new RuntimeUDFContext(taskInfo, null, executionConfig, new HashMap<String, Future<Path>>(), new HashMap<String, Accumulator<?, ?>>(), new UnregisteredMetricsGroup()), executionConfig);
        assertEquals(asList(1, 2, 3, 4, 5, 6), resultMutableSafe);
        assertEquals(asList(1, 2, 3, 4, 5, 6), resultRegular);
        assertTrue(opened.get());
        assertTrue(closed.get());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}