[core] Support watermark partition markdone mode (#5284)

Author: xiangyuf

docs/layouts/shortcodes/generated/flink_connector_configuration.html

@@ -86,6 +86,12 @@
             <td>Duration</td>
             <td>Set a time duration when a partition has no new data after this time duration, mark the done status to indicate that the data is ready.</td>
         </tr>
+        <tr>
+            <td><h5>partition.mark-done-action.mode</h5></td>
+            <td style="word-wrap: break-word;">process-time</td>
+            <td><p>Enum</p></td>
+            <td>How to trigger partition mark done action.<br /><br />Possible values:<ul><li>"process-time": Based on the time of the machine, mark the partition done once the processing time passes period time plus delay.</li><li>"watermark": Based on the watermark of the input, mark the partition done once the watermark passes period time plus delay.</li></ul></td>
+        </tr>
         <tr>
             <td><h5>partition.time-interval</h5></td>
             <td style="word-wrap: break-word;">(none)</td>

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/FlinkConnectorOptions.java

@@ -356,6 +356,12 @@ public class FlinkConnectorOptions {
                     .withDescription(
                             "Allow sink committer and writer operator to be chained together");
 
+    public static final ConfigOption<PartitionMarkDoneActionMode> PARTITION_MARK_DONE_MODE =
+            key("partition.mark-done-action.mode")
+                    .enumType(PartitionMarkDoneActionMode.class)
+                    .defaultValue(PartitionMarkDoneActionMode.PROCESS_TIME)
+                    .withDescription("How to trigger partition mark done action.");
+
     public static final ConfigOption<Duration> PARTITION_IDLE_TIME_TO_DONE =
             key("partition.idle-time-to-done")
                     .durationType()
@@ -544,4 +550,32 @@ public InlineElement getDescription() {
             return text(description);
         }
     }
+
+    /** The mode for partition mark done. */
+    public enum PartitionMarkDoneActionMode implements DescribedEnum {
+        PROCESS_TIME(
+                "process-time",
+                "Based on the time of the machine, mark the partition done once the processing time passes period time plus delay."),
+        WATERMARK(
+                "watermark",
+                "Based on the watermark of the input, mark the partition done once the watermark passes period time plus delay.");
+
+        private final String value;
+        private final String description;
+
+        PartitionMarkDoneActionMode(String value, String description) {
+            this.value = value;
+            this.description = description;
+        }
+
+        @Override
+        public String toString() {
+            return value;
+        }
+
+        @Override
+        public InlineElement getDescription() {
+            return text(description);
+        }
+    }
 }

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/partition/PartitionMarkDone.java

@@ -21,6 +21,7 @@
 import org.apache.paimon.CoreOptions;
 import org.apache.paimon.CoreOptions.MergeEngine;
 import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.flink.FlinkConnectorOptions.PartitionMarkDoneActionMode;
 import org.apache.paimon.manifest.ManifestCommittable;
 import org.apache.paimon.options.Options;
 import org.apache.paimon.partition.actions.PartitionMarkDoneAction;
@@ -32,24 +33,33 @@
 import org.apache.paimon.utils.PartitionPathUtils;
 
 import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.time.Duration;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
 import static org.apache.paimon.CoreOptions.PARTITION_MARK_DONE_WHEN_END_INPUT;
 import static org.apache.paimon.flink.FlinkConnectorOptions.PARTITION_IDLE_TIME_TO_DONE;
+import static org.apache.paimon.flink.FlinkConnectorOptions.PARTITION_MARK_DONE_MODE;
 
 /** Mark partition done. */
 public class PartitionMarkDone implements PartitionListener {
 
+    private static final Logger LOG = LoggerFactory.getLogger(PartitionMarkDone.class);
+
     private final InternalRowPartitionComputer partitionComputer;
     private final PartitionMarkDoneTrigger trigger;
     private final List<PartitionMarkDoneAction> actions;
     private final boolean waitCompaction;
+    private final PartitionMarkDoneActionMode partitionMarkDoneActionMode;
 
     public static Optional<PartitionMarkDone> create(
             ClassLoader cl,
@@ -86,7 +96,12 @@ public static Optional<PartitionMarkDone> create(
                                 || coreOptions.mergeEngine() == MergeEngine.FIRST_ROW);
 
         return Optional.of(
-                new PartitionMarkDone(partitionComputer, trigger, actions, waitCompaction));
+                new PartitionMarkDone(
+                        partitionComputer,
+                        trigger,
+                        actions,
+                        waitCompaction,
+                        options.get(PARTITION_MARK_DONE_MODE)));
     }
 
     private static boolean disablePartitionMarkDone(
@@ -108,15 +123,25 @@ public PartitionMarkDone(
             InternalRowPartitionComputer partitionComputer,
             PartitionMarkDoneTrigger trigger,
             List<PartitionMarkDoneAction> actions,
-            boolean waitCompaction) {
+            boolean waitCompaction,
+            PartitionMarkDoneActionMode partitionMarkDoneActionMode) {
         this.partitionComputer = partitionComputer;
         this.trigger = trigger;
         this.actions = actions;
         this.waitCompaction = waitCompaction;
+        this.partitionMarkDoneActionMode = partitionMarkDoneActionMode;
     }
 
     @Override
     public void notifyCommittable(List<ManifestCommittable> committables) {
+        if (partitionMarkDoneActionMode == PartitionMarkDoneActionMode.WATERMARK) {
+            markDoneByWatermark(committables);
+        } else {
+            markDoneByProcessTime(committables);
+        }
+    }
+
+    private void markDoneByProcessTime(List<ManifestCommittable> committables) {
         Set<BinaryRow> partitions = new HashSet<>();
         boolean endInput = false;
         for (ManifestCommittable committable : committables) {
@@ -141,6 +166,58 @@ public void notifyCommittable(List<ManifestCommittable> committables) {
         markDone(trigger.donePartitions(endInput), actions);
     }
 
+    private void markDoneByWatermark(List<ManifestCommittable> committables) {
+        // extract watermarks from committables and update partition watermarks
+        Tuple2<Map<BinaryRow, Long>, Boolean> extractedWatermarks =
+                extractPartitionWatermarks(committables);
+        Map<BinaryRow, Long> partitionWatermarks = extractedWatermarks.f0;
+        boolean endInput = extractedWatermarks.f1;
+        Optional<Long> latestWatermark = partitionWatermarks.values().stream().max(Long::compareTo);
+
+        if (!latestWatermark.isPresent()) {
+            LOG.warn("No watermark found in this batch of committables, skip partition mark done.");
+            return;
+        }
+
+        partitionWatermarks.forEach(
+                (row, value) -> {
+                    String partition =
+                            PartitionPathUtils.generatePartitionPath(
+                                    partitionComputer.generatePartValues(row));
+                    trigger.notifyPartition(partition, value);
+                });
+
+        markDone(trigger.donePartitions(endInput, latestWatermark.get(), true), actions);
+    }
+
+    private Tuple2<Map<BinaryRow, Long>, Boolean> extractPartitionWatermarks(
+            List<ManifestCommittable> committables) {
+        boolean endInput = false;
+        Map<BinaryRow, Long> partitionWatermarks = new HashMap<>();
+        for (ManifestCommittable committable : committables) {
+            Long watermark = committable.watermark();
+            if (watermark != null) {
+                for (CommitMessage commitMessage : committable.fileCommittables()) {
+                    CommitMessageImpl message = (CommitMessageImpl) commitMessage;
+                    if (waitCompaction
+                            || !message.indexIncrement().isEmpty()
+                            || !message.newFilesIncrement().isEmpty()) {
+                        partitionWatermarks.compute(
+                                message.partition(),
+                                (partition, old) ->
+                                        old == null ? watermark : Math.max(old, watermark));
+                    }
+                }
+            }
+
+            if (committable.identifier() == Long.MAX_VALUE) {
+                endInput = true;
+            }
+        }
+
+        return Tuple2.of(partitionWatermarks, endInput);
+    }
+
     public static void markDone(List<String> partitions, List<PartitionMarkDoneAction> actions) {
         for (String partition : partitions) {
             try {

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/partition/PartitionMarkDoneTrigger.java

@@ -111,11 +111,16 @@ void notifyPartition(String partition, long currentTimeMillis) {
     }
 
     public List<String> donePartitions(boolean endInput) {
-        return donePartitions(endInput, System.currentTimeMillis());
+        return donePartitions(endInput, System.currentTimeMillis(), false);
     }
 
-    @VisibleForTesting
     List<String> donePartitions(boolean endInput, long currentTimeMillis) {
+        return donePartitions(endInput, currentTimeMillis, false);
+    }
+
+    @VisibleForTesting
+    List<String> donePartitions(
+            boolean endInput, long currentTimeMillis, boolean watermarkEnabled) {
         if (endInput && markDoneWhenEndInput) {
             return new ArrayList<>(pendingPartitions.keySet());
         }
@@ -131,11 +136,21 @@ List<String> donePartitions(boolean endInput, long currentTimeMillis) {
             String partition = entry.getKey();
 
             long lastUpdateTime = entry.getValue();
-            long partitionStartTime =
-                    extractDateTime(partition)
-                            .atZone(ZoneId.systemDefault())
-                            .toInstant()
-                            .toEpochMilli();
+            long partitionStartTime;
+            if (watermarkEnabled) {
+                // watermark should be compared as UTC time
+                partitionStartTime =
+                        extractDateTime(partition)
+                                .atZone(ZoneId.of("UTC"))
+                                .toInstant()
+                                .toEpochMilli();
+            } else {
+                partitionStartTime =
+                        extractDateTime(partition)
+                                .atZone(ZoneId.systemDefault())
+                                .toInstant()
+                                .toEpochMilli();
+            }
             long partitionEndTime = partitionStartTime + timeInterval;
             lastUpdateTime = Math.max(lastUpdateTime, partitionEndTime);