If the partition count or kafka IO size is large, then skip committin… (#37510)

* If the partition count or kafka IO size is large, then skip committing offsets that are not changed. Reduce kafka commit load

* Address PR review feedback for idle partition optimization

- Refactor commitCheckpointMark to use Java streams (per @johnjcasey)
  Changed from explicit for-loop to streams-based filtering for better
  code consistency with existing patterns

- Add debug logging for idle partitions (per @tomstepp)
  Log the count of idle partitions skipped during each commit to aid
  in monitoring and debugging the optimization

- Implement time-based periodic commits (per @tomstepp)
  Track last commit time per partition and ensure commits happen at
  least every 10 minutes even for idle partitions. This supports time
  lag monitoring use cases where customers track time since last commit.

- Add unit test for idle partition behavior (per @tomstepp)
  New test KafkaUnboundedReaderIdlePartitionTest verifies that:
  * Idle partitions are not committed repeatedly
  * Active partitions trigger commits correctly
  * Uses mock consumer to track commit calls

All changes maintain backward compatibility and follow Apache Beam
coding standards (spotless formatting applied).

* Fix test to follow Beam patterns for MockConsumer initialization

Rewrote KafkaUnboundedReaderIdlePartitionTest to follow the exact
pattern used in KafkaIOTest.java:
- Proper MockConsumer initialization with partition metadata
- Correct setup of beginning/end offsets
- Consumer records with proper offsets and timestamps
- schedulePollTask for record enqueueing based on position
- Override commitSync to track commit calls
- Use reader.start() before reader.advance()

This ensures the test properly initializes the Kafka consumer and
doesn't fail with IllegalStateException during source.split().

---------

Co-authored-by: Kishore Pola <kpola@paloaltonetworks.com>

Author: kishorepola

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

@@ -23,6 +23,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -79,6 +80,11 @@
  */
 class KafkaUnboundedReader<K, V> extends UnboundedReader<KafkaRecord<K, V>> {
 
+  // Track last successfully committed offsets to suppress no-op commits for idle partitions.
+  private final Map<TopicPartition, Long> lastCommittedOffsets = new HashMap<>();
+  // Track last commit time per partition to ensure periodic commits for time lag monitoring.
+  private final Map<TopicPartition, Instant> lastCommitTimes = new HashMap<>();
+
   ///////////////////// Reader API ////////////////////////////////////////////////////////////
   @SuppressWarnings("FutureReturnValueIgnored")
   @Override
@@ -375,6 +381,8 @@ public boolean offsetBasedDeduplicationSupported() {
   private static final Duration RECORDS_DEQUEUE_POLL_TIMEOUT_MAX = Duration.millis(20);
   private static final Duration RECORDS_ENQUEUE_POLL_TIMEOUT = Duration.millis(100);
   private static final Duration MIN_COMMIT_FAIL_LOG_INTERVAL = Duration.standardMinutes(10);
+  // Maximum time between commits for idle partitions (for time lag monitoring).
+  private static final Duration MAX_IDLE_COMMIT_INTERVAL = Duration.standardMinutes(10);
 
   // Use a separate thread to read Kafka messages. Kafka Consumer does all its work including
   // network I/O inside poll(). Polling only inside #advance(), especially with a small timeout
@@ -611,37 +619,88 @@ private void consumerPollLoop() {
 
   private void commitCheckpointMark() {
     KafkaCheckpointMark checkpointMark = finalizedCheckpointMark.getAndSet(null);
-    if (checkpointMark != null) {
-      LOG.debug("{}: Committing finalized checkpoint {}", this, checkpointMark);
-      Consumer<byte[], byte[]> consumer = Preconditions.checkStateNotNull(this.consumer);
-      Instant now = Instant.now();
+    if (checkpointMark == null) {
+      return;
+    }
 
-      try {
-        consumer.commitSync(
-            checkpointMark.getPartitions().stream()
-                .filter(p -> p.getNextOffset() != UNINITIALIZED_OFFSET)
-                .collect(
-                    Collectors.toMap(
-                        p -> new TopicPartition(p.getTopic(), p.getPartition()),
-                        p -> new OffsetAndMetadata(p.getNextOffset()))));
+    LOG.debug("{}: Committing finalized checkpoint {}", this, checkpointMark);
+    Consumer<byte[], byte[]> consumer = Preconditions.checkStateNotNull(this.consumer);
+    Instant now = Instant.now();
+
+    try {
+      // Commit only partitions whose offsets have advanced since the last successful commit
+      // for this reader, or partitions that haven't been committed within MAX_IDLE_COMMIT_INTERVAL.
+      // This suppresses no-op commits for idle partitions while ensuring periodic commits
+      // for time lag monitoring.
+      Map<TopicPartition, OffsetAndMetadata> toCommit =
+          checkpointMark.getPartitions().stream()
+              .filter(p -> p.getNextOffset() != UNINITIALIZED_OFFSET)
+              .filter(
+                  p -> {
+                    TopicPartition tp = new TopicPartition(p.getTopic(), p.getPartition());
+                    Long prev = lastCommittedOffsets.get(tp);
+                    long next = p.getNextOffset();
+                    Instant lastCommitTime = lastCommitTimes.get(tp);
+
+                    // Commit if offset has advanced
+                    if (prev == null || next > prev) {
+                      return true;
+                    }
+
+                    // Also commit if partition hasn't been committed within max idle interval
+                    if (lastCommitTime == null
+                        || now.isAfter(lastCommitTime.plus(MAX_IDLE_COMMIT_INTERVAL))) {
+                      return true;
+                    }
+
+                    return false;
+                  })
+              .collect(
+                  Collectors.toMap(
+                      p -> new TopicPartition(p.getTopic(), p.getPartition()),
+                      p -> new OffsetAndMetadata(p.getNextOffset())));
+
+      int totalPartitions = checkpointMark.getPartitions().size();
+      int idlePartitions = totalPartitions - toCommit.size();
+      if (idlePartitions > 0) {
+        LOG.debug(
+            "{}: Skipping commit for {} idle partitions ({} of {} partitions active)",
+            this,
+            idlePartitions,
+            toCommit.size(),
+            totalPartitions);
+      }
+
+      if (toCommit.isEmpty()) {
+        // Nothing advanced since last successful commit; avoid noisy commitSync().
+        return;
+      }
+
+      consumer.commitSync(toCommit);
+
+      // Only update after a successful commit.
+      for (Map.Entry<TopicPartition, OffsetAndMetadata> e : toCommit.entrySet()) {
+        lastCommittedOffsets.put(e.getKey(), e.getValue().offset());
+        lastCommitTimes.put(e.getKey(), now);
+      }
+
+      nextAllowedCommitFailLogTime = now.plus(MIN_COMMIT_FAIL_LOG_INTERVAL);
+    } catch (Exception e) {
+      // Log but ignore the exception. Committing consumer offsets to Kafka is not critical for
+      // KafkaIO because it relies on the offsets stored in KafkaCheckpointMark.
+      if (now.isAfter(nextAllowedCommitFailLogTime)) {
+        LOG.warn(
+            String.format(
+                "%s: Did not successfully commit finalized checkpoint for > %s. Current checkpoint: %s",
+                this, MIN_COMMIT_FAIL_LOG_INTERVAL, checkpointMark),
+            e);
         nextAllowedCommitFailLogTime = now.plus(MIN_COMMIT_FAIL_LOG_INTERVAL);
-      } catch (Exception e) {
-        // Log but ignore the exception. Committing consumer offsets to Kafka is not critical for
-        // KafkaIO because it relies on the offsets stored in KafkaCheckpointMark.
-        if (now.isAfter(nextAllowedCommitFailLogTime)) {
-          LOG.warn(
-              String.format(
-                  "%s: Did not successfully commit finalized checkpoint for > %s. Current checkpoint: %s",
-                  this, MIN_COMMIT_FAIL_LOG_INTERVAL, checkpointMark),
-              e);
-          nextAllowedCommitFailLogTime = now.plus(MIN_COMMIT_FAIL_LOG_INTERVAL);
-        } else {
-          LOG.info(
-              String.format(
-                  "%s: Could not commit finalized checkpoint. Commit will be retried with subsequent reads. Current checkpoint: %s",
-                  this, checkpointMark),
-              e);
-        }
+      } else {
+        LOG.info(
+            String.format(
+                "%s: Could not commit finalized checkpoint. Commit will be retried with subsequent reads. Current checkpoint: %s",
+                this, checkpointMark),
+            e);
       }
     }
   }