[Draft] Flink 2.0 support

Author: Abacn

gradle.properties

@@ -39,6 +39,6 @@ docker_image_default_repo_root=apache
 docker_image_default_repo_prefix=beam_
 
 # supported flink versions
-flink_versions=1.17,1.18,1.19,1.20
+flink_versions=1.17,1.18,1.19,1.20,2.0
 # supported python versions
 python_versions=3.10,3.11,3.12,3.13

runners/flink/2.0/build.gradle

@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+project.ext {
+  flink_major = '2.0'
+  flink_version = '2.0.1'
+  excluded_files = [
+    'main': [
+        // Used by DataSet API only
+        "org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapter.java",
+        "org/apache/beam/runners/flink/FlinkBatchPipelineTranslator.java",
+        "org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java",
+        "org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java",
+        "org/apache/beam/runners/flink/translation/functions/FlinkNonMergingReduceFunction.java",
+        // Moved to org.apache.flink.runtime.state.StateBackendFactory
+        "org/apache/beam/runners/flink/FlinkStateBackendFactory.java",
+    ],
+    'test': [
+        // Used by DataSet API only
+        "org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapterTest.java",
+        "org/apache/beam/runners/flink/batch/NonMergingGroupByKeyTest.java",
+        "org/apache/beam/runners/flink/batch/ReshuffleTest.java",
+    ]
+  ]
+}
+
+// Load the main build script which contains all build logic.
+apply from: "../flink_runner.gradle"

runners/flink/2.0/job-server-container/build.gradle

@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+def basePath = '../../job-server-container'
+
+project.ext {
+  resource_path = basePath
+}
+
+// Load the main build script which contains all build logic.
+apply from: "$basePath/flink_job_server_container.gradle"

runners/flink/2.0/job-server/build.gradle

@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+def basePath = '../../job-server'
+
+project.ext {
+  // Look for the source code in the parent module
+  main_source_dirs = ["$basePath/src/main/java"]
+  test_source_dirs = ["$basePath/src/test/java"]
+  main_resources_dirs = ["$basePath/src/main/resources"]
+  test_resources_dirs = ["$basePath/src/test/resources"]
+  archives_base_name = 'beam-runners-flink-2.0-job-server'
+}
+
+// Load the main build script which contains all build logic.
+apply from: "$basePath/flink_job_server.gradle"

runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java

@@ -0,0 +1,116 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink;
+
+import java.util.Map;
+import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
+import org.apache.beam.runners.flink.translation.utils.CountingPipelineVisitor;
+import org.apache.beam.runners.flink.translation.utils.LookupPipelineVisitor;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+
+/**
+ * Helper for {@link FlinkBatchPipelineTranslator} and translators in {@link
+ * FlinkBatchTransformTranslators}.
+ */
+@SuppressWarnings({
+  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+class FlinkBatchTranslationContext {
+  private final PipelineOptions options;
+
+  private AppliedPTransform<?, ?, ?> currentTransform;
+
+  private final CountingPipelineVisitor countingPipelineVisitor = new CountingPipelineVisitor();
+  private final LookupPipelineVisitor lookupPipelineVisitor = new LookupPipelineVisitor();
+
+  // ------------------------------------------------------------------------
+
+  FlinkBatchTranslationContext(PipelineOptions options) {
+    this.options = options;
+  }
+
+  void init(Pipeline pipeline) {
+    pipeline.traverseTopologically(countingPipelineVisitor);
+    pipeline.traverseTopologically(lookupPipelineVisitor);
+  }
+
+  public PipelineOptions getPipelineOptions() {
+    return options;
+  }
+
+  /**
+   * Sets the AppliedPTransform which carries input/output.
+   *
+   * @param currentTransform Current transformation.
+   */
+  void setCurrentTransform(AppliedPTransform<?, ?, ?> currentTransform) {
+    this.currentTransform = currentTransform;
+  }
+
+  AppliedPTransform<?, ?, ?> getCurrentTransform() {
+    return currentTransform;
+  }
+
+  Map<TupleTag<?>, Coder<?>> getOutputCoders(PTransform<?, ?> transform) {
+    return lookupPipelineVisitor.getOutputCoders(transform);
+  }
+
+  <T> TypeInformation<WindowedValue<T>> getTypeInfo(PCollection<T> collection) {
+    return getTypeInfo(collection.getCoder(), collection.getWindowingStrategy());
+  }
+
+  <T> TypeInformation<WindowedValue<T>> getTypeInfo(
+      Coder<T> coder, WindowingStrategy<?, ?> windowingStrategy) {
+    WindowedValues.FullWindowedValueCoder<T> windowedValueCoder =
+        WindowedValues.getFullCoder(coder, windowingStrategy.getWindowFn().windowCoder());
+
+    return new CoderTypeInformation<>(windowedValueCoder, options);
+  }
+
+  Map<TupleTag<?>, PCollection<?>> getInputs(PTransform<?, ?> transform) {
+    return lookupPipelineVisitor.getInputs(transform);
+  }
+
+  <T extends PValue> T getInput(PTransform<T, ?> transform) {
+    return lookupPipelineVisitor.getInput(transform);
+  }
+
+  Map<TupleTag<?>, PCollection<?>> getOutputs(PTransform<?, ?> transform) {
+    return lookupPipelineVisitor.getOutputs(transform);
+  }
+
+  <T extends PValue> T getOutput(PTransform<?, T> transform) {
+    return lookupPipelineVisitor.getOutput(transform);
+  }
+
+  /** {@link CountingPipelineVisitor#getNumConsumers(PValue)}. */
+  int getNumConsumers(PValue value) {
+    return countingPipelineVisitor.getNumConsumers(value);
+  }
+}