Search in sources :

Example 1 with RichFlatJoinFunction

use of org.apache.flink.api.common.functions.RichFlatJoinFunction 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());
}
Also used : Accumulator(org.apache.flink.api.common.accumulators.Accumulator) RichFlatJoinFunction(org.apache.flink.api.common.functions.RichFlatJoinFunction) UnregisteredMetricsGroup(org.apache.flink.metrics.groups.UnregisteredMetricsGroup) Configuration(org.apache.flink.configuration.Configuration) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) TaskInfo(org.apache.flink.api.common.TaskInfo) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) Collector(org.apache.flink.util.Collector) RuntimeUDFContext(org.apache.flink.api.common.functions.util.RuntimeUDFContext) Future(java.util.concurrent.Future) Test(org.junit.Test)

Aggregations

ArrayList (java.util.ArrayList)1 HashMap (java.util.HashMap)1 Future (java.util.concurrent.Future)1 AtomicBoolean (java.util.concurrent.atomic.AtomicBoolean)1 ExecutionConfig (org.apache.flink.api.common.ExecutionConfig)1 TaskInfo (org.apache.flink.api.common.TaskInfo)1 Accumulator (org.apache.flink.api.common.accumulators.Accumulator)1 RichFlatJoinFunction (org.apache.flink.api.common.functions.RichFlatJoinFunction)1 RuntimeUDFContext (org.apache.flink.api.common.functions.util.RuntimeUDFContext)1 Configuration (org.apache.flink.configuration.Configuration)1 UnregisteredMetricsGroup (org.apache.flink.metrics.groups.UnregisteredMetricsGroup)1 Collector (org.apache.flink.util.Collector)1 Test (org.junit.Test)1