Examples with KeySelector org.apache.flink.api.java.functions.KeySelector used on opensource projects

Example 46 with KeySelector

use of org.apache.flink.api.java.functions.KeySelector in project flink by apache.

the class JoinITCase method testDefaultJoinOnTwoCustomTypeInputsWithInnerClassKeyExtractorsClosureCleaner.

@Test
public void testDefaultJoinOnTwoCustomTypeInputsWithInnerClassKeyExtractorsClosureCleaner() throws Exception {
    /*
         * (Default) Join on two custom type inputs with key extractors, implemented as inner classes to test closure
         * cleaning
         */
    final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
    DataSet<CustomType> ds1 = CollectionDataSets.getCustomTypeDataSet(env);
    DataSet<CustomType> ds2 = CollectionDataSets.getSmallCustomTypeDataSet(env);
    DataSet<Tuple2<CustomType, CustomType>> joinDs = ds1.join(ds2).where(new KeySelector<CustomType, Integer>() {

        @Override
        public Integer getKey(CustomType value) {
            return value.myInt;
        }
    }).equalTo(new KeySelector<CustomType, Integer>() {

        @Override
        public Integer getKey(CustomType value) throws Exception {
            return value.myInt;
        }
    });
    List<Tuple2<CustomType, CustomType>> result = joinDs.collect();
    String expected = "1,0,Hi,1,0,Hi\n" + "2,1,Hello,2,1,Hello\n" + "2,1,Hello,2,2,Hello world\n" + "2,2,Hello world,2,1,Hello\n" + "2,2,Hello world,2,2,Hello world\n";
    compareResultAsTuples(result, expected);
}

Example 47 with KeySelector

use of org.apache.flink.api.java.functions.KeySelector in project flink by apache.

the class DataStreamTest method testPOJOnoHashCodeKeyRejection.

@Test
public void testPOJOnoHashCodeKeyRejection() {
    KeySelector<POJOWithoutHashCode, POJOWithoutHashCode> keySelector = new KeySelector<POJOWithoutHashCode, POJOWithoutHashCode>() {

        @Override
        public POJOWithoutHashCode getKey(POJOWithoutHashCode value) throws Exception {
            return value;
        }
    };
    StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
    DataStream<POJOWithoutHashCode> input = env.fromElements(new POJOWithoutHashCode(new int[] { 1, 2 }));
    // adjust the rule
    expectedException.expect(InvalidProgramException.class);
    input.keyBy(keySelector);
}

Example 48 with KeySelector

use of org.apache.flink.api.java.functions.KeySelector in project flink by apache.

the class DataStreamTest method testPOJOWithNestedArrayAndHashCodeWorkAround.

@Test
public void testPOJOWithNestedArrayAndHashCodeWorkAround() {
    StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
    DataStream<POJOWithHashCode> input = env.fromElements(new POJOWithHashCode(new int[] { 1, 2 }));
    input.keyBy(new KeySelector<POJOWithHashCode, POJOWithHashCode>() {

        @Override
        public POJOWithHashCode getKey(POJOWithHashCode value) throws Exception {
            return value;
        }
    }).addSink(new SinkFunction<POJOWithHashCode>() {

        @Override
        public void invoke(POJOWithHashCode value) throws Exception {
            Assert.assertEquals(value.getId(), new int[] { 1, 2 });
        }
    });
}

Example 49 with KeySelector

use of org.apache.flink.api.java.functions.KeySelector in project flink by apache.

the class DataStreamTest method sinkKeyTest.

@Test
public void sinkKeyTest() {
    StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
    DataStreamSink<Long> sink = env.generateSequence(1, 100).print();
    assertEquals(0, getStreamGraph(env).getStreamNode(sink.getTransformation().getId()).getStatePartitioners().length);
    assertTrue(getStreamGraph(env).getStreamNode(sink.getTransformation().getId()).getInEdges().get(0).getPartitioner() instanceof ForwardPartitioner);
    KeySelector<Long, Long> key1 = new KeySelector<Long, Long>() {

        private static final long serialVersionUID = 1L;

        @Override
        public Long getKey(Long value) throws Exception {
            return (long) 0;
        }
    };
    DataStreamSink<Long> sink2 = env.generateSequence(1, 100).keyBy(key1).print();
    assertEquals(1, getStreamGraph(env).getStreamNode(sink2.getTransformation().getId()).getStatePartitioners().length);
    assertNotNull(getStreamGraph(env).getStreamNode(sink2.getTransformation().getId()).getStateKeySerializer());
    assertNotNull(getStreamGraph(env).getStreamNode(sink2.getTransformation().getId()).getStateKeySerializer());
    assertEquals(key1, getStreamGraph(env).getStreamNode(sink2.getTransformation().getId()).getStatePartitioners()[0]);
    assertTrue(getStreamGraph(env).getStreamNode(sink2.getTransformation().getId()).getInEdges().get(0).getPartitioner() instanceof KeyGroupStreamPartitioner);
    KeySelector<Long, Long> key2 = new KeySelector<Long, Long>() {

        private static final long serialVersionUID = 1L;

        @Override
        public Long getKey(Long value) throws Exception {
            return (long) 0;
        }
    };
    DataStreamSink<Long> sink3 = env.generateSequence(1, 100).keyBy(key2).print();
    assertEquals(1, getStreamGraph(env).getStreamNode(sink3.getTransformation().getId()).getStatePartitioners().length);
    assertEquals(key2, getStreamGraph(env).getStreamNode(sink3.getTransformation().getId()).getStatePartitioners()[0]);
    assertTrue(getStreamGraph(env).getStreamNode(sink3.getTransformation().getId()).getInEdges().get(0).getPartitioner() instanceof KeyGroupStreamPartitioner);
}

Example 50 with KeySelector

use of org.apache.flink.api.java.functions.KeySelector in project flink by apache.

the class NotifyCheckpointAbortedITCase method testNotifyCheckpointAborted.

/**
 * Verify operators would be notified as checkpoint aborted.
 *
 * <p>The job would run with at least two checkpoints. The 1st checkpoint would fail due to add
 * checkpoint to store, and the 2nd checkpoint would decline by async checkpoint phase of
 * 'DeclineSink'.
 *
 * <p>The job graph looks like: NormalSource --> keyBy --> NormalMap --> DeclineSink
 */
@Test(timeout = TEST_TIMEOUT)
public void testNotifyCheckpointAborted() throws Exception {
    final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
    env.enableCheckpointing(200, CheckpointingMode.EXACTLY_ONCE);
    env.getCheckpointConfig().enableUnalignedCheckpoints(unalignedCheckpointEnabled);
    env.getCheckpointConfig().setTolerableCheckpointFailureNumber(1);
    env.disableOperatorChaining();
    env.setParallelism(1);
    final StateBackend failingStateBackend = new DeclineSinkFailingStateBackend(checkpointPath);
    env.setStateBackend(failingStateBackend);
    env.addSource(new NormalSource()).name("NormalSource").keyBy((KeySelector<Tuple2<Integer, Integer>, Integer>) value -> value.f0).transform("NormalMap", TypeInformation.of(Integer.class), new NormalMap()).transform(DECLINE_SINK_NAME, TypeInformation.of(Object.class), new DeclineSink());
    final ClusterClient<?> clusterClient = cluster.getClusterClient();
    JobGraph jobGraph = env.getStreamGraph().getJobGraph();
    JobID jobID = jobGraph.getJobID();
    clusterClient.submitJob(jobGraph).get();
    TestingCompletedCheckpointStore.addCheckpointLatch.await();
    log.info("The checkpoint to abort is ready to add to checkpoint store.");
    TestingCompletedCheckpointStore.abortCheckpointLatch.trigger();
    log.info("Verifying whether all operators have been notified of checkpoint-1 aborted.");
    verifyAllOperatorsNotifyAborted();
    log.info("Verified that all operators have been notified of checkpoint-1 aborted.");
    resetAllOperatorsNotifyAbortedLatches();
    verifyAllOperatorsNotifyAbortedTimes(1);
    NormalSource.waitLatch.trigger();
    log.info("Verifying whether all operators have been notified of checkpoint-2 aborted.");
    verifyAllOperatorsNotifyAborted();
    log.info("Verified that all operators have been notified of checkpoint-2 aborted.");
    verifyAllOperatorsNotifyAbortedTimes(2);
    clusterClient.cancel(jobID).get();
    log.info("Test is verified successfully as expected.");
}