Fix Kafka with Redistribute

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

@@ -890,7 +890,6 @@ static <K, V> void setupExternalBuilder(
           builder.setRedistributeNumKeys(0);
           builder.setAllowDuplicates(false);
         }
-        System.out.println("xxx builder service" + builder.toString());
       }
 
       private static <T> Coder<T> resolveCoder(Class<Deserializer<T>> deserializer) {
@@ -1697,11 +1696,12 @@ public PCollection<KafkaRecord<K, V>> expand(PBegin input) {
         }
 
         if (kafkaRead.isRedistributed()) {
-          // fail here instead.
-          checkArgument(
-              kafkaRead.isCommitOffsetsInFinalizeEnabled(),
-              "commitOffsetsInFinalize() can't be enabled with isRedistributed");
+          if (kafkaRead.isCommitOffsetsInFinalizeEnabled()) {
+            LOG.warn(
+                "Offsets committed due to usage of commitOffsetsInFinalize() may not capture all work processed due to use of withRedistribute()");
+          }
           PCollection<KafkaRecord<K, V>> output = input.getPipeline().apply(transform);
+
           if (kafkaRead.getRedistributeNumKeys() == 0) {
             return output.apply(
                 "Insert Redistribute",
@@ -2654,6 +2654,12 @@ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescriptor>
         if (getRedistributeNumKeys() == 0) {
           LOG.warn("This will create a key per record, which is sub-optimal for most use cases.");
         }
+        if (isCommitOffsetEnabled() || configuredKafkaCommit()) {
+          LOG.warn(
+              "Either auto_commit is set, or commitOffsetEnabled is enabled (or both), but since "
+                  + "withRestribute() is enabled, the runner may have additional work processed that "
+                  + "is ahead of the current checkpoint");
+        }
       }
 
       if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
@@ -2687,8 +2693,7 @@ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescriptor>
                             .getSchemaCoder(KafkaSourceDescriptor.class),
                         recordCoder));
 
-        boolean applyCommitOffsets =
-            isCommitOffsetEnabled() && !configuredKafkaCommit() && !isRedistribute();
+        boolean applyCommitOffsets = isCommitOffsetEnabled() && !configuredKafkaCommit();
         if (!applyCommitOffsets) {
           return outputWithDescriptor
               .apply(MapElements.into(new TypeDescriptor<KafkaRecord<K, V>>() {}).via(KV::getValue))
@@ -2710,6 +2715,15 @@ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescriptor>
           if (Comparators.lexicographical(Comparator.<String>naturalOrder())
                   .compare(requestedVersion, targetVersion)
               < 0) {
+            // Redistribute is not allowed with commits prior to 2.59.0, since there is a Reshuffle
+            // prior to the redistribute. The reshuffle will occur before commits are offsetted and
+            // before outputting KafkaRecords. Adding a redistrube then afterwards doesn't provide
+            // additional performance benefit.
+            checkArgument(
+                !isRedistribute(),
+                "Can not enable isRedistribute() while committing offsets prior to "
+                    + String.join(".", targetVersion));
+
             return expand259Commits(
                 outputWithDescriptor, recordCoder, input.getPipeline().getSchemaRegistry());
           }

sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java

@@ -88,6 +88,7 @@
 import org.apache.beam.sdk.metrics.SinkMetrics;
 import org.apache.beam.sdk.metrics.SourceMetrics;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.StreamingOptions;
 import org.apache.beam.sdk.testing.ExpectedLogs;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
@@ -628,33 +629,40 @@ public void testRiskyConfigurationWarnsProperly() {
   }
 
   @Test
-  public void testCommitOffsetsInFinalizeAndRedistributeErrors() {
-    thrown.expect(Exception.class);
-    thrown.expectMessage("commitOffsetsInFinalize() can't be enabled with isRedistributed");
-
+  public void testCommitOffsetsInFinalizeAndRedistributeWarnings() {
     int numElements = 1000;
 
     PCollection<Long> input =
         p.apply(
                 mkKafkaReadTransform(numElements, numElements, new ValueAsTimestampFn(), true, 0)
+                    .commitOffsetsInFinalize()
                     .withConsumerConfigUpdates(
-                        ImmutableMap.of(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true))
+                        ImmutableMap.of(ConsumerConfig.GROUP_ID_CONFIG, "group_id"))
                     .withoutMetadata())
             .apply(Values.create());
 
     addCountingAsserts(input, numElements);
     p.run();
+
+    kafkaIOExpectedLogs.verifyWarn(
+        "Offsets committed due to usage of commitOffsetsInFinalize() may not capture all work processed due to use of withRedistribute()");
   }
 
   @Test
   public void testNumKeysIgnoredWithRedistributeNotEnabled() {
+    thrown.expect(Exception.class);
+    thrown.expectMessage(
+        "withRedistributeNumKeys is ignored if withRedistribute() is not enabled on the transform");
+
     int numElements = 1000;
 
     PCollection<Long> input =
         p.apply(
                 mkKafkaReadTransform(numElements, numElements, new ValueAsTimestampFn(), false, 0)
+                    .withRedistributeNumKeys(100)
+                    .commitOffsetsInFinalize()
                     .withConsumerConfigUpdates(
-                        ImmutableMap.of(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true))
+                        ImmutableMap.of(ConsumerConfig.GROUP_ID_CONFIG, "group_id"))
                     .withoutMetadata())
             .apply(Values.create());
 
@@ -663,6 +671,32 @@ public void testNumKeysIgnoredWithRedistributeNotEnabled() {
     p.run();
   }
 
+  @Test
+  public void testDisableRedistributeKafkaOffsetLegacy() {
+    thrown.expect(Exception.class);
+    thrown.expectMessage(
+        "Can not enable isRedistribute() while committing offsets prior to 2.60.0");
+    p.getOptions().as(StreamingOptions.class).setUpdateCompatibilityVersion("2.59.0");
+
+    p.apply(
+            Create.of(
+                KafkaSourceDescriptor.of(
+                    new TopicPartition("topic", 1),
+                    null,
+                    null,
+                    null,
+                    null,
+                    ImmutableList.of("8.8.8.8:9092"))))
+        .apply(
+            KafkaIO.<Long, Long>readSourceDescriptors()
+                .withKeyDeserializer(LongDeserializer.class)
+                .withValueDeserializer(LongDeserializer.class)
+                .withRedistribute()
+                .withProcessingTime()
+                .commitOffsets());
+    p.run();
+  }
+
   @Test
   public void testUnreachableKafkaBrokers() {
     // Expect an exception when the Kafka brokers are not reachable on the workers.