Move Kafka read with redistribute override to Kafka IO package. (#36887)

* Add kafka read override to Dataflow java runner.

* Fix spot bugs (spacing)

* Add unit test of redistribute override

* Update test dependencies via gradle

* Add logic and test case for explicitly disabled.

* Add explicitly enabled test case

* Use boolean asserts over assertThat, assert each read is visited only once, refine suppressed lint warnings to just instanceof on matches method.

* Move kafka read with redistribute override to Kafka IO package.

* Lint fixes

* Remove Kafka test dependencies for Dataflow worker

* Ignore abandoned nodes in the test since we just need to replace the transforms.

Author: tomstepp

runners/google-cloud-dataflow-java/build.gradle

@@ -129,8 +129,6 @@ dependencies {
   testImplementation library.java.google_cloud_dataflow_java_proto_library_all
   testImplementation library.java.jackson_dataformat_yaml
   testImplementation library.java.mockito_inline
-  testImplementation project(":sdks:java:io:kafka")
-  testImplementation library.java.kafka_clients
   validatesRunner project(path: ":sdks:java:core", configuration: "shadowTest")
   validatesRunner project(path: project.path, configuration: "testRuntimeMigration")
   validatesRunner library.java.hamcrest

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java

@@ -659,10 +659,7 @@ private List<PTransformOverride> getOverrides(boolean streaming) {
 
       try {
         overridesBuilder.add(KafkaIO.Read.KAFKA_READ_OVERRIDE);
-        overridesBuilder.add(
-            PTransformOverride.of(
-                KafkaReadWithRedistributeOverride.matcher(),
-                new KafkaReadWithRedistributeOverride.Factory()));
+        overridesBuilder.add(KafkaIO.Read.KAFKA_REDISTRIBUTE_OVERRIDE);
       } catch (NoClassDefFoundError e) {
         // Do nothing. io-kafka is an optional dependency of runners-google-cloud-dataflow-java
         // and only needed when KafkaIO is used in the pipeline.

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

@@ -1817,6 +1817,13 @@ private boolean runnerPrefersLegacyRead(PipelineOptions options) {
       return true;
     }
 
+    /** A {@link PTransformOverride} for runners to override redistributed Kafka Read transforms. */
+    @Internal
+    public static final PTransformOverride KAFKA_REDISTRIBUTE_OVERRIDE =
+        PTransformOverride.of(
+            KafkaReadWithRedistributeOverride.matcher(),
+            new KafkaReadWithRedistributeOverride.Factory<>());
+
     /**
      * A {@link PTransformOverride} for runners to swap {@link ReadFromKafkaViaSDF} to legacy Kafka
      * read if runners doesn't have a good support on executing unbounded Splittable DoFn.

…w/KafkaReadWithRedistributeOverride.java …a/KafkaReadWithRedistributeOverride.javarunners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/KafkaReadWithRedistributeOverride.java renamed to sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadWithRedistributeOverride.java runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/KafkaReadWithRedistributeOverride.java renamed to sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadWithRedistributeOverride.java

@@ -15,11 +15,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.runners.dataflow;
+package org.apache.beam.sdk.io.kafka;
 
 import java.util.Map;
-import org.apache.beam.sdk.io.kafka.KafkaIO;
-import org.apache.beam.sdk.io.kafka.KafkaRecord;
 import org.apache.beam.sdk.runners.AppliedPTransform;
 import org.apache.beam.sdk.runners.PTransformMatcher;
 import org.apache.beam.sdk.runners.PTransformOverrideFactory;
@@ -48,8 +46,8 @@ public boolean matches(AppliedPTransform<?, ?, ?> application) {
   }
 
   /**
-   * {@link PTransformOverrideFactory} for {@link KafkaIO.Read} that enables {@code
-   * withOffsetDeduplication} when {@code withRedistribute} is enabled.
+   * {@link PTransformOverrideFactory} for {@link org.apache.beam.sdk.io.kafka.KafkaIO.Read} that
+   * enables {@code withOffsetDeduplication} when {@code withRedistribute} is enabled.
    */
   static class Factory<K, V>
       implements PTransformOverrideFactory<

…fkaReadWithRedistributeOverrideTest.java …fkaReadWithRedistributeOverrideTest.javarunners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/KafkaReadWithRedistributeOverrideTest.java renamed to sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadWithRedistributeOverrideTest.java runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/KafkaReadWithRedistributeOverrideTest.java renamed to sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadWithRedistributeOverrideTest.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.runners.dataflow;
+package org.apache.beam.sdk.io.kafka;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.nullValue;
@@ -25,7 +25,6 @@
 import java.io.Serializable;
 import java.util.Collections;
 import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.io.kafka.KafkaIO;
 import org.apache.beam.sdk.runners.PTransformOverride;
 import org.apache.beam.sdk.runners.TransformHierarchy.Node;
 import org.apache.beam.sdk.testing.TestPipeline;
@@ -129,5 +128,6 @@ public void leaveCompositeTransform(Node node) {
           }
         };
     p.traverseTopologically(visitor);
+    p.enableAbandonedNodeEnforcement(false);
   }
 }