[FLINK-4804] [py] Fix first() failing when applied to groupings

apache · Oct 21, 2016 · 5c83e78 · 5c83e78
1 parent 41d5167
commit 5c83e78
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/flink-libraries/flink-python/src/main/java/org/apache/flink/python/api/PythonPlanBinder.java b/flink-libraries/flink-python/src/main/java/org/apache/flink/python/api/PythonPlanBinder.java
@@ -459,9 +459,22 @@ private void createDistinctOperation(PythonOperationInfo info) throws IOExceptio
 				.map(new KeyDiscarder()).setParallelism(getParallelism(info)).name("DistinctPostStep"));
 	}
 
+	@SuppressWarnings("unchecked")
 	private void createFirstOperation(PythonOperationInfo info) throws IOException {
-		DataSet op = (DataSet) sets.get(info.parentID);
+		Object op = sets.get(info.parentID);
-		sets.put(info.setID, op.first(info.count).setParallelism(getParallelism(info)).name("First"));
+		if (op instanceof DataSet) {
+			sets.put(info.setID, ((DataSet) op).first(info.count).setParallelism(getParallelism(info)).name("First"));
+			return;
+		}
+		if (op instanceof UnsortedGrouping) {
+			sets.put(info.setID, ((UnsortedGrouping) op).first(info.count).setParallelism(getParallelism(info)).name("First")
+				.map(new KeyDiscarder()).setParallelism(getParallelism(info)).name("FirstPostStep"));
+			return;
+		}
+		if (op instanceof SortedGrouping) {
+			sets.put(info.setID, ((SortedGrouping) op).first(info.count).setParallelism(getParallelism(info)).name("First")
+				.map(new KeyDiscarder()).setParallelism(getParallelism(info)).name("FirstPostStep"));
+		}
 	}
 
 	private void createGroupOperation(PythonOperationInfo info) throws IOException {

diff --git a/flink-libraries/flink-python/src/test/python/org/apache/flink/python/api/test_main2.py b/flink-libraries/flink-python/src/test/python/org/apache/flink/python/api/test_main2.py
@@ -126,6 +126,10 @@ def map(self, value):
     d1 \
         .first(1) \
         .map_partition(Verify([1], "First")).output()
+    d4 \
+        .group_by(0) \
+        .first(1) \
+        .map_partition(Verify([(1, 0.5, "hello", True), (2, 0.4, "world", False)], "Grouped First")).output()
     d1 \
         .rebalance()
     d6 \