[FLINK-9470] Allow querying the key in KeyedProcessFunction

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/KeyedProcessFunction.java

@@ -110,6 +110,11 @@ public abstract class KeyedProcessFunction<K, I, O> extends AbstractRichFunction
 		 * @param value The record to emit.
 		 */
 		public abstract <X> void output(OutputTag<X> outputTag, X value);
+
+		/**
+		 * Get key of the element being processed.
+		 */
+		public abstract K getCurrentKey();
 	}
 
 	/**
@@ -124,6 +129,7 @@ public abstract class KeyedProcessFunction<K, I, O> extends AbstractRichFunction
 		/**
 		 * Get key of the firing timer.
 		 */
+		@Override
 		public abstract K getCurrentKey();
 	}
 

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/co/KeyedBroadcastProcessFunction.java

@@ -158,6 +158,12 @@ public abstract class KeyedBroadcastProcessFunction<KS, IN1, IN2, OUT> extends B
 		 * A {@link TimerService} for querying time and registering timers.
 		 */
 		public abstract TimerService timerService();
+
+
+		/**
+		 * Get key of the element being processed.
+		 */
+		public abstract KS getCurrentKey();
 	}
 
 	/**
@@ -174,6 +180,7 @@ public abstract class KeyedBroadcastProcessFunction<KS, IN1, IN2, OUT> extends B
 		/**
 		 * Get the key of the firing timer.
 		 */
+		@Override
 		public abstract KS getCurrentKey();
 	}
 }

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/KeyedProcessOperator.java

@@ -131,6 +131,12 @@ public class KeyedProcessOperator<K, IN, OUT>
 
 			output.collect(outputTag, new StreamRecord<>(value, element.getTimestamp()));
 		}
+
+		@Override
+		@SuppressWarnings("unchecked")
+		public K getCurrentKey() {
+			return (K) KeyedProcessOperator.this.getCurrentKey();
+		}
 	}
 
 	private class OnTimerContextImpl extends KeyedProcessFunction<K, IN, OUT>.OnTimerContext {

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/co/CoBroadcastWithKeyedOperator.java

@@ -288,6 +288,13 @@ public class CoBroadcastWithKeyedOperator<KS, IN1, IN2, OUT>
 			}
 			return state;
 		}
+
+		@Override
+		@SuppressWarnings("unchecked")
+		public KS getCurrentKey() {
+			return (KS) CoBroadcastWithKeyedOperator.this.getCurrentKey();
+		}
+
 	}
 
 	private class OnTimerContextImpl extends KeyedBroadcastProcessFunction<KS, IN1, IN2, OUT>.OnTimerContext {

flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/KeyedProcessOperatorTest.java

@@ -23,6 +23,7 @@ import org.apache.flink.api.common.state.ValueStateDescriptor;
 import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
 import org.apache.flink.api.common.typeutils.base.IntSerializer;
 import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
 import org.apache.flink.streaming.api.TimeDomain;
 import org.apache.flink.streaming.api.TimerService;
@@ -43,6 +44,7 @@ import org.junit.rules.ExpectedException;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Tests {@link KeyedProcessOperator}.
@@ -52,6 +54,48 @@ public class KeyedProcessOperatorTest extends TestLogger {
 	@Rule
 	public ExpectedException expectedException = ExpectedException.none();
 
+	@Test
+	public void testKeyQuerying() throws Exception {
+
+		class KeyQueryingProcessFunction extends KeyedProcessFunction<Integer, Tuple2<Integer, String>, String> {
+
+			@Override
+			public void processElement(
+				Tuple2<Integer, String> value,
+				Context ctx,
+				Collector<String> out) throws Exception {
+
+				assertTrue("Did not get expected key.", ctx.getCurrentKey().equals(value.f0));
+
+				// we check that we receive this output, to ensure that the assert was actually checked
+				out.collect(value.f1);
+			}
+		}
+
+		KeyedProcessOperator<Integer, Tuple2<Integer, String>, String> operator =
+			new KeyedProcessOperator<>(new KeyQueryingProcessFunction());
+
+		try (
+			OneInputStreamOperatorTestHarness<Tuple2<Integer, String>, String> testHarness =
+				new KeyedOneInputStreamOperatorTestHarness<>(operator, (in) -> in.f0 , BasicTypeInfo.INT_TYPE_INFO)) {
+
+			testHarness.setup();
+			testHarness.open();
+
+			testHarness.processElement(new StreamRecord<>(Tuple2.of(5, "5"), 12L));
+			testHarness.processElement(new StreamRecord<>(Tuple2.of(42, "42"), 13L));
+
+			ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+			expectedOutput.add(new StreamRecord<>("5", 12L));
+			expectedOutput.add(new StreamRecord<>("42", 13L));
+
+			TestHarnessUtil.assertOutputEquals(
+				"Output was not correct.",
+				expectedOutput,
+				testHarness.getOutput());
+		}
+	}
+
 	@Test
 	public void testTimestampAndWatermarkQuerying() throws Exception {
 

flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/co/CoBroadcastWithKeyedOperatorTest.java

@@ -25,6 +25,7 @@ import org.apache.flink.api.common.state.ValueStateDescriptor;
 import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
 import org.apache.flink.runtime.state.KeyedStateFunction;
 import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction;
@@ -70,6 +71,56 @@ public class CoBroadcastWithKeyedOperatorTest {
 					BasicTypeInfo.INT_TYPE_INFO
 			);
 
+	@Test
+	public void testKeyQuerying() throws Exception {
+
+		class KeyQueryingProcessFunction extends KeyedBroadcastProcessFunction<Integer, Tuple2<Integer, String>, String, String> {
+
+			@Override
+			public void processElement(
+				Tuple2<Integer, String> value,
+				ReadOnlyContext ctx,
+				Collector<String> out) throws Exception {
+				assertTrue("Did not get expected key.", ctx.getCurrentKey().equals(value.f0));
+
+				// we check that we receive this output, to ensure that the assert was actually checked
+				out.collect(value.f1);
+
+			}
+
+			@Override
+			public void processBroadcastElement(
+				String value,
+				Context ctx,
+				Collector<String> out) throws Exception {
+
+			}
+		}
+
+		CoBroadcastWithKeyedOperator<Integer, Tuple2<Integer, String>, String, String> operator =
+			new CoBroadcastWithKeyedOperator<>(new KeyQueryingProcessFunction(), Collections.emptyList());
+
+		try (
+			TwoInputStreamOperatorTestHarness<Tuple2<Integer, String>, String, String> testHarness =
+				new KeyedTwoInputStreamOperatorTestHarness<>(operator, (in) -> in.f0 , null, BasicTypeInfo.INT_TYPE_INFO)) {
+
+			testHarness.setup();
+			testHarness.open();
+
+			testHarness.processElement1(new StreamRecord<>(Tuple2.of(5, "5"), 12L));
+			testHarness.processElement1(new StreamRecord<>(Tuple2.of(42, "42"), 13L));
+
+			ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+			expectedOutput.add(new StreamRecord<>("5", 12L));
+			expectedOutput.add(new StreamRecord<>("42", 13L));
+
+			TestHarnessUtil.assertOutputEquals(
+				"Output was not correct.",
+				expectedOutput,
+				testHarness.getOutput());
+		}
+	}
+
 	/** Test the iteration over the keyed state on the broadcast side. */
 	@Test
 	public void testAccessToKeyedStateIt() throws Exception {