[runners-flink] apache#30621 use groupBy for Reshuffle in batch

je-ik · Mar 13, 2024 · 5697c67 · 5697c67
1 parent 289bbc4
commit 5697c67
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 10 deletions.
diff --git a/runners/flink/flink_runner.gradle b/runners/flink/flink_runner.gradle
@@ -192,13 +192,11 @@ dependencies {
 
   if (flink_version.compareTo("1.14") >= 0) {
     implementation "org.apache.flink:flink-runtime:$flink_version"
-    implementation "org.apache.flink:flink-optimizer:$flink_version"
     implementation "org.apache.flink:flink-metrics-core:$flink_version"
     testImplementation "org.apache.flink:flink-runtime:$flink_version:tests"
     testImplementation "org.apache.flink:flink-rpc-akka:$flink_version"
   } else {
     implementation "org.apache.flink:flink-runtime_2.12:$flink_version"
-    implementation "org.apache.flink:flink-optimizer_2.12:$flink_version"
     testImplementation "org.apache.flink:flink-runtime_2.12:$flink_version:tests"
   }
   testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")

diff --git a/...ers/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java b/...ers/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java
@@ -94,8 +94,6 @@
 import org.apache.flink.api.java.operators.MapOperator;
 import org.apache.flink.api.java.operators.SingleInputUdfOperator;
 import org.apache.flink.api.java.operators.UnsortedGrouping;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.optimizer.Optimizer;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
@@ -413,12 +411,15 @@ public void translateNode(
               outputType,
               FlinkIdentityFunction.of(),
               getCurrentTransformName(context));
-      final Configuration partitionOptions = new Configuration();
-      partitionOptions.setString(
-          Optimizer.HINT_SHIP_STRATEGY, Optimizer.HINT_SHIP_STRATEGY_REPARTITION);
-      context.setOutputDataSet(
-          context.getOutput(transform),
-          retypedDataSet.map(FlinkIdentityFunction.of()).withParameters(partitionOptions));
+      WindowedValue.WindowedValueCoder<KV<K, InputT>> kvWvCoder =
+          (WindowedValue.WindowedValueCoder<KV<K, InputT>>) outputType.getCoder();
+      KvCoder<K, InputT> kvCoder = (KvCoder<K, InputT>) kvWvCoder.getValueCoder();
+      DataSet<WindowedValue<KV<K, InputT>>> reshuffle =
+          retypedDataSet
+              .groupBy(new KvKeySelector<>(kvCoder.getKeyCoder()))
+              .<WindowedValue<KV<K, InputT>>>reduceGroup((i, c) -> i.forEach(c::collect))
+              .returns(outputType);
+      context.setOutputDataSet(context.getOutput(transform), reshuffle);
     }
   }