Github user zentol commented on a diff in the pull request:
https://github.com/apache/flink/pull/3501#discussion_r105407025
--- Diff:
flink-streaming-java/src/test/java/org/apache/flink/streaming/api/DataStreamTest.java
---
@@ -906,6 +919,243 @@ public void testChannelSelectors() {
}
/////////////////////////////////////////////////////////////
+ // KeyBy testing
+ /////////////////////////////////////////////////////////////
+
+ @Rule
+ public ExpectedException expectedException = ExpectedException.none();
+
+ @Test
+ public void testPrimitiveArrayKeyRejection() {
+
+ KeySelector<Tuple2<Integer[], String>, int[]> keySelector =
+ new KeySelector<Tuple2<Integer[], String>,
int[]>() {
+
+ @Override
+ public int[] getKey(Tuple2<Integer[], String> value)
throws Exception {
+ int[] ks = new int[value.f0.length];
+ for (int i = 0; i < ks.length; i++) {
+ ks[i] = value.f0[i];
+ }
+ return ks;
+ }
+ };
+
+ testKeyRejection(keySelector,
PrimitiveArrayTypeInfo.INT_PRIMITIVE_ARRAY_TYPE_INFO);
+ }
+
+ @Test
+ public void testBasicArrayKeyRejection() {
+
+ KeySelector<Tuple2<Integer[], String>, Integer[]> keySelector =
+ new KeySelector<Tuple2<Integer[], String>,
Integer[]>() {
+
+ @Override
+ public Integer[] getKey(Tuple2<Integer[], String>
value) throws Exception {
+ return value.f0;
+ }
+ };
+
+ testKeyRejection(keySelector,
BasicArrayTypeInfo.INT_ARRAY_TYPE_INFO);
+ }
+
+ @Test
+ public void testObjectArrayKeyRejection() {
+
+ KeySelector<Tuple2<Integer[], String>, Object[]> keySelector =
+ new KeySelector<Tuple2<Integer[], String>,
Object[]>() {
+
+ @Override
+ public Object[]
getKey(Tuple2<Integer[], String> value) throws Exception {
+ Object[] ks = new
TestClass[value.f0.length];
+ for (int i = 0; i < ks.length;
i++) {
+ ks[i] = new
TestClass(value.f0[i]);
+ }
+ return ks;
+ }
+ };
+
+ ObjectArrayTypeInfo<Object[], Object> keyTypeInfo =
ObjectArrayTypeInfo.getInfoFor(
+ Object[].class, new
GenericTypeInfo<>(Object.class));
+
+ testKeyRejection(keySelector, keyTypeInfo);
+ }
+
+ private <K> void testKeyRejection(KeySelector<Tuple2<Integer[],
String>, K> keySelector, TypeInformation<K> expectedKeyType) {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+
+ DataStream<Tuple2<Integer[], String>> input = env.fromElements(
+ new Tuple2<>(new Integer[] {1, 2}, "barfoo")
+ );
+
+ Assert.assertEquals(expectedKeyType,
TypeExtractor.getKeySelectorTypes(keySelector, input.getType()));
+
+ // adjust the rule
+ expectedException.expect(InvalidProgramException.class);
+ expectedException.expectMessage(new StringStartsWith("Type " +
expectedKeyType + " cannot be used as key."));
+
+ input.keyBy(keySelector);
+ }
+
+ //////////////// Composite Key Tests : POJOs
////////////////
+
+ @Test
+ public void testPOJOWithNestedArrayNoHashCodeKeyRejection() {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+
+ DataStream<POJOWithHashCode> input = env.fromElements(
+ new POJOWithHashCode(new int[] {1, 2}));
+
+ TypeInformation<?> expectedTypeInfo = new
TupleTypeInfo<Tuple1<int[]>>(
+
PrimitiveArrayTypeInfo.INT_PRIMITIVE_ARRAY_TYPE_INFO);
+
+ // adjust the rule
+ expectedException.expect(InvalidProgramException.class);
+ expectedException.expectMessage(new StringStartsWith("Type " +
expectedTypeInfo + " cannot be used as key."));
+
+ input.keyBy("id");
+ }
+
+ @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});
+ }
+ });
+ }
+
+ @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);
+ }
+
+ //////////////// Composite Key Tests : Tuples
////////////////
+
+ @Test
+ public void testTupleNestedArrayKeyRejection() {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+
+ DataStream<Tuple2<Integer[], String>> input = env.fromElements(
+ new Tuple2<>(new Integer[] {1, 2},
"test-test"));
+
+ TypeInformation<?> expectedTypeInfo = new
TupleTypeInfo<Tuple2<Integer[], String>>(
+ BasicArrayTypeInfo.INT_ARRAY_TYPE_INFO,
BasicTypeInfo.STRING_TYPE_INFO);
+
+ // adjust the rule
+ expectedException.expect(InvalidProgramException.class);
+ expectedException.expectMessage(new StringStartsWith("Type " +
expectedTypeInfo + " cannot be used as key."));
+
+
+ input.keyBy(new KeySelector<Tuple2<Integer[],String>,
Tuple2<Integer[],String>>() {
+ @Override
+ public Tuple2<Integer[], String>
getKey(Tuple2<Integer[], String> value) throws Exception {
+ return value;
+ }
+ });
+ }
+
+ @Test
+ public void testPrimitiveKeyAcceptance() throws Exception {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ env.setMaxParallelism(1);
+
+ DataStream<Integer> input = env.fromElements(new
Integer(10000));
+
+ KeyedStream<Integer, Object> keyedStream = input.keyBy(new
KeySelector<Integer, Object>() {
+ @Override
+ public Object getKey(Integer value) throws Exception {
+ return value;
+ }
+ });
+
+ keyedStream.addSink(new SinkFunction<Integer>() {
+ @Override
+ public void invoke(Integer value) throws Exception {
+ Assert.assertEquals(10000L, (long) value);
+ }
+ });
+
+ env.execute();
--- End diff --
We don't need to execute a job, correct? just the keyBy passing should be
enough.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
