Provide kafka unbounded reader to checkpoint mark when offset based deduplication is supported.

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

@@ -21,7 +21,6 @@
 
 import java.io.Serializable;
 import java.util.List;
-import java.util.Optional;
 import org.apache.avro.reflect.AvroIgnore;
 import org.apache.beam.sdk.coders.DefaultCoder;
 import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
@@ -35,21 +34,24 @@
  */
 @DefaultCoder(AvroCoder.class)
 public class KafkaCheckpointMark implements UnboundedSource.CheckpointMark {
+  private static final long OFFSET_DEDUP_PARTITIONS_PER_SPLIT = 1;
 
   private List<PartitionMark> partitions;
 
-  @AvroIgnore
-  private Optional<KafkaUnboundedReader<?, ?>> reader; // Present when offsets need to be committed.
+  @AvroIgnore private KafkaUnboundedReader<?, ?> reader;
+
+  private boolean commitOffsetsInFinalize;
 
   @SuppressWarnings("initialization") // Avro will set the fields by breaking abstraction
   private KafkaCheckpointMark() {} // for Avro
 
-  private static final long OFFSET_DEDUP_PARTITIONS_PER_SPLIT = 1;
-
   public KafkaCheckpointMark(
-      List<PartitionMark> partitions, Optional<KafkaUnboundedReader<?, ?>> reader) {
+      List<PartitionMark> partitions,
+      KafkaUnboundedReader<?, ?> reader,
+      boolean commitOffsetsInFinalize) {
     this.partitions = partitions;
     this.reader = reader;
+    this.commitOffsetsInFinalize = commitOffsetsInFinalize;
   }
 
   public List<PartitionMark> getPartitions() {
@@ -58,7 +60,10 @@ public List<PartitionMark> getPartitions() {
 
   @Override
   public void finalizeCheckpoint() {
-    reader.ifPresent(r -> r.finalizeCheckpointMarkAsync(this));
+    if (!commitOffsetsInFinalize) {
+      return;
+    }
+    reader.finalizeCheckpointMarkAsync(this);
     // Is it ok to commit asynchronously, or should we wait till this (or newer) is committed?
     // Often multiple marks would be finalized at once, since we only need to finalize the latest,
     // it is better to wait a little while. Currently maximum delay is same as KAFKA_POLL_TIMEOUT
@@ -72,11 +77,7 @@ public String toString() {
 
   @Override
   public byte[] getOffsetLimit() {
-    if (!reader.isPresent()) {
-      throw new RuntimeException(
-          "KafkaCheckpointMark reader is not present while calling getOffsetLimit().");
-    }
-    if (!reader.get().offsetBasedDeduplicationSupported()) {
+    if (!reader.offsetBasedDeduplicationSupported()) {
       throw new RuntimeException(
           "Unexpected getOffsetLimit() called while KafkaUnboundedReader not configured for offset deduplication.");
     }

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

@@ -268,7 +268,8 @@ public CheckpointMark getCheckpointMark() {
                         p.nextOffset,
                         p.lastWatermark.getMillis()))
             .collect(Collectors.toList()),
-        source.getSpec().isCommitOffsetsInFinalizeEnabled() ? Optional.of(this) : Optional.empty());
+        this,
+        source.getSpec().isCommitOffsetsInFinalizeEnabled());
   }
 
   @Override