[FLINK] Support native checkpoint

.github/workflows/flink.yml

@@ -70,7 +70,7 @@ jobs:
           export fmt_SOURCE=BUNDLED
           export folly_SOURCE=BUNDLED
           git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git
-          cd velox4j && git reset --hard 115edf79d265a61c30d45dfcc6ce932ad92378ca
+          cd velox4j && git reset --hard 8ceb4e5e108588d92b89498bb3f1a3c0c4b029b5
           git apply $GITHUB_WORKSPACE/gluten-flink/patches/fix-velox4j.patch
           $GITHUB_WORKSPACE/build/mvn clean install -DskipTests -Dgpg.skip -Dspotless.skip=true
           cd ..

gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/metrics/SourceTaskMetrics.java

@@ -17,13 +17,18 @@
 package org.apache.gluten.table.runtime.metrics;
 
 import io.github.zhztheplayer.velox4j.query.SerialTask;
+import io.github.zhztheplayer.velox4j.query.SerialTaskStats;
 
 import org.apache.flink.metrics.Counter;
 import org.apache.flink.metrics.groups.OperatorMetricGroup;
+import org.apache.flink.runtime.metrics.groups.InternalOperatorMetricGroup;
+import org.apache.flink.runtime.metrics.groups.TaskIOMetricGroup;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 
+import java.util.List;
+
 public class SourceTaskMetrics {
 
   private final String keyOperatorType = "operatorType";
@@ -33,31 +38,113 @@ public class SourceTaskMetrics {
   private final long metricUpdateInterval = 2000;
   private Counter sourceNumRecordsOut;
   private Counter sourceNumBytesOut;
-  private long lastUpdateTime = System.currentTimeMillis();
+  private Counter taskNumRecordsIn;
+  private Counter taskNumRecordsOut;
+  private Counter taskNumBytesIn;
+  private Counter taskNumBytesOut;
+  private long lastUpdateTime = 0;
 
   public SourceTaskMetrics(OperatorMetricGroup metricGroup) {
     sourceNumRecordsOut = metricGroup.getIOMetricGroup().getNumRecordsOutCounter();
     sourceNumBytesOut = metricGroup.getIOMetricGroup().getNumBytesOutCounter();
+    if (metricGroup instanceof InternalOperatorMetricGroup) {
+      TaskIOMetricGroup taskIOMetricGroup =
+          ((InternalOperatorMetricGroup) metricGroup).getTaskIOMetricGroup();
+      taskNumRecordsIn = taskIOMetricGroup.getNumRecordsInCounter();
+      taskNumRecordsOut = taskIOMetricGroup.getNumRecordsOutCounter();
+      taskNumBytesIn = taskIOMetricGroup.getNumBytesInCounter();
+      taskNumBytesOut = taskIOMetricGroup.getNumBytesOutCounter();
+    }
+  }
+
+  SourceTaskMetrics(Counter sourceNumRecordsOut, Counter sourceNumBytesOut) {
+    this(sourceNumRecordsOut, sourceNumBytesOut, null, null, null, null);
+  }
+
+  SourceTaskMetrics(
+      Counter sourceNumRecordsOut,
+      Counter sourceNumBytesOut,
+      Counter taskNumRecordsIn,
+      Counter taskNumRecordsOut,
+      Counter taskNumBytesIn,
+      Counter taskNumBytesOut) {
+    this.sourceNumRecordsOut = sourceNumRecordsOut;
+    this.sourceNumBytesOut = sourceNumBytesOut;
+    this.taskNumRecordsIn = taskNumRecordsIn;
+    this.taskNumRecordsOut = taskNumRecordsOut;
+    this.taskNumBytesIn = taskNumBytesIn;
+    this.taskNumBytesOut = taskNumBytesOut;
   }
 
   public boolean updateMetrics(SerialTask task, String planId) {
+    return updateMetrics(task.collectStats(), planId);
+  }
+
+  boolean updateMetrics(SerialTaskStats taskStats, String planId) {
     long currentTime = System.currentTimeMillis();
     if (currentTime - lastUpdateTime < metricUpdateInterval) {
       return false;
     }
     try {
-      ObjectNode planStats = task.collectStats().planStats(planId);
-      JsonNode jsonNode = planStats.get(keyOperatorType);
-      if (jsonNode.asText().equals(sourceOperatorName)) {
-        long numRecordsOut = planStats.get(keyInputRows).asInt();
-        long numBytesOut = planStats.get(keyInputBytes).asInt();
-        sourceNumRecordsOut.inc(numRecordsOut - sourceNumRecordsOut.getCount());
-        sourceNumBytesOut.inc(numBytesOut - sourceNumBytesOut.getCount());
+      ObjectNode planStats = findSourceStats(taskStats, planId);
+      if (planStats != null) {
+        long inputRows = planStats.get(keyInputRows).asLong();
+        long inputBytes = planStats.get(keyInputBytes).asLong();
+        syncCounter(sourceNumRecordsOut, inputRows);
+        syncCounter(sourceNumBytesOut, inputBytes);
+        syncCounter(taskNumRecordsIn, inputRows);
+        syncCounter(taskNumRecordsOut, inputRows);
+        syncCounter(taskNumBytesIn, inputBytes);
+        syncCounter(taskNumBytesOut, inputBytes);
       }
     } catch (Exception e) {
       return false;
     }
     lastUpdateTime = currentTime;
     return true;
   }
+
+  private ObjectNode findSourceStats(SerialTaskStats taskStats, String planId) {
+    try {
+      ObjectNode planStats = taskStats.planStats(planId);
+      if (isSourceStats(planStats)) {
+        return planStats;
+      }
+    } catch (Exception ignored) {
+      // Fall back to the unique TableScan below.
+    }
+
+    ObjectNode sourceStats = null;
+    List<ObjectNode> allPlanStats = taskStats.planStats();
+    for (ObjectNode planStats : allPlanStats) {
+      if (!isSourceStats(planStats)) {
+        continue;
+      }
+      if (sourceStats != null) {
+        return null;
+      }
+      sourceStats = planStats;
+    }
+    return sourceStats;
+  }
+
+  private boolean isSourceStats(ObjectNode planStats) {
+    JsonNode operatorType = planStats.get(keyOperatorType);
+    return operatorType != null
+        && sourceOperatorName.equals(operatorType.asText())
+        && planStats.has(keyInputRows)
+        && planStats.has(keyInputBytes);
+  }
+
+  private void syncCounter(Counter counter, long value) {
+    if (counter == null) {
+      return;
+    }
+    long delta = value - counter.getCount();
+    if (delta > 0) {
+      counter.inc(delta);
+    } else if (delta < 0) {
+      counter.dec(-delta);
+    }
+  }
 }

gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java

@@ -37,6 +37,8 @@
 import io.github.zhztheplayer.velox4j.type.RowType;
 
 import org.apache.flink.contrib.streaming.state.RocksDBKeyedStateBackend;
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.runtime.metrics.groups.InternalOperatorMetricGroup;
 import org.apache.flink.runtime.state.StateInitializationContext;
 import org.apache.flink.runtime.state.StateSnapshotContext;
 import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
@@ -66,6 +68,7 @@ public class GlutenOneInputOperator<IN, OUT> extends TableStreamOperator<OUT>
   private final Class<OUT> outClass;
   private transient VectorInputBridge<IN> inputBridge;
   private transient VectorOutputBridge<OUT> outputBridge;
+  private transient Counter taskNumRecordsOut;
   private final GlutenMailboxHolder mailboxHolder = new GlutenMailboxHolder();
 
   public GlutenOneInputOperator(
@@ -151,6 +154,10 @@ void initSession() {
   @Override
   public void open() throws Exception {
     super.open();
+    if (metrics instanceof InternalOperatorMetricGroup) {
+      taskNumRecordsOut =
+          ((InternalOperatorMetricGroup) metrics).getTaskIOMetricGroup().getNumRecordsOutCounter();
+    }
     if (!mailboxHolder().get().isMailboxBound()) {
       ensureMailboxInitialized(getContainingTask());
     }
@@ -201,8 +208,12 @@ private void drainTaskOutput() {
             StatefulWatermark watermark = statefulElement.asWatermark();
             output.emitWatermark(new Watermark(watermark.getTimestamp()));
           } else {
-            outputBridge.collect(
-                output, statefulElement.asRecord(), sessionResource.getAllocator(), outputType);
+            long emittedRecords =
+                outputBridge.collect(
+                    output, statefulElement.asRecord(), sessionResource.getAllocator(), outputType);
+            if (taskNumRecordsOut != null) {
+              taskNumRecordsOut.inc(emittedRecords);
+            }
           }
         } finally {
           statefulElement.close();
@@ -272,14 +283,16 @@ public String getId() {
 
   @Override
   public void prepareSnapshotPreBarrier(long checkpointId) throws Exception {
-    // TODO: notify velox
+    // Drain any in-flight data from the Velox pipeline before snapshot.
+    // Flink has already aligned barriers, so no new input will arrive.
+    drainTaskOutput();
     super.prepareSnapshotPreBarrier(checkpointId);
   }
 
   @Override
   public void snapshotState(StateSnapshotContext context) throws Exception {
     // TODO: implement it
-    task.snapshotState(0);
+    task.snapshotState(context.getCheckpointId());
     super.snapshotState(context);
   }
 

gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java

@@ -32,6 +32,8 @@
 import io.github.zhztheplayer.velox4j.stateful.StatefulWatermark;
 import io.github.zhztheplayer.velox4j.type.RowType;
 
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.state.FunctionInitializationContext;
 import org.apache.flink.runtime.state.FunctionSnapshotContext;
@@ -43,6 +45,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -65,6 +68,8 @@ public class GlutenSourceFunction<OUT> extends RichParallelSourceFunction<OUT>
   private SerialTask task;
   private SourceTaskMetrics taskMetrics;
   private final Class<OUT> outClass;
+  private transient ListState<String> checkpointState;
+  private transient String[] restoredCheckpointRecords = new String[0];
 
   public GlutenSourceFunction(
       StatefulPlanNode planNode,
@@ -205,15 +210,29 @@ public void close() throws Exception {
 
   @Override
   public void snapshotState(FunctionSnapshotContext context) throws Exception {
-    // TODO: implement it
-    this.task.snapshotState(0);
+    checkpointState.clear();
+    String[] checkpointRecords = this.task.snapshotState(context.getCheckpointId());
+    for (String checkpointRecord : checkpointRecords) {
+      checkpointState.add(checkpointRecord);
+    }
   }
 
   @Override
   public void initializeState(FunctionInitializationContext context) throws Exception {
+    checkpointState =
+        context
+            .getOperatorStateStore()
+            .getListState(
+                new ListStateDescriptor<>("gluten-native-source-checkpoint", String.class));
+    if (context.isRestored()) {
+      List<String> records = new ArrayList<>();
+      for (String checkpointRecord : checkpointState.get()) {
+        records.add(checkpointRecord);
+      }
+      restoredCheckpointRecords = records.toArray(new String[0]);
+    }
     initSession();
-    // TODO: implement it
-    this.task.initializeState(0, null);
+    this.task.initializeState(0, null, restoredCheckpointRecords);
   }
 
   public String[] notifyCheckpointComplete(long checkpointId) throws Exception {

gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenTwoInputOperator.java

@@ -33,6 +33,8 @@
 import io.github.zhztheplayer.velox4j.stateful.StatefulWatermark;
 import io.github.zhztheplayer.velox4j.type.RowType;
 
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.runtime.metrics.groups.InternalOperatorMetricGroup;
 import org.apache.flink.runtime.state.StateInitializationContext;
 import org.apache.flink.runtime.state.StateSnapshotContext;
 import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
@@ -71,6 +73,7 @@ public class GlutenTwoInputOperator<IN, OUT> extends AbstractStreamOperator<OUT>
   private final Class<OUT> outClass;
   private VectorInputBridge<IN> inputBridge;
   private VectorOutputBridge<OUT> outputBridge;
+  private transient Counter taskNumRecordsOut;
   private String description;
   private final GlutenMailboxHolder mailboxHolder = new GlutenMailboxHolder();
 
@@ -118,6 +121,10 @@ public String getDescription() {
   @Override
   public void open() throws Exception {
     super.open();
+    if (metrics instanceof InternalOperatorMetricGroup) {
+      taskNumRecordsOut =
+          ((InternalOperatorMetricGroup) metrics).getTaskIOMetricGroup().getNumRecordsOutCounter();
+    }
     if (!mailboxHolder().get().isMailboxBound()) {
       ensureMailboxInitialized(getContainingTask());
     }
@@ -181,8 +188,12 @@ private void drainTaskOutput() {
             StatefulWatermark watermark = element.asWatermark();
             output.emitWatermark(new Watermark(watermark.getTimestamp()));
           } else {
-            outputBridge.collect(
-                output, element.asRecord(), sessionResource.getAllocator(), outputType);
+            long emittedRecords =
+                outputBridge.collect(
+                    output, element.asRecord(), sessionResource.getAllocator(), outputType);
+            if (taskNumRecordsOut != null) {
+              taskNumRecordsOut.inc(emittedRecords);
+            }
           }
         } finally {
           element.close();
@@ -260,14 +271,16 @@ public String getRightId() {
 
   @Override
   public void prepareSnapshotPreBarrier(long checkpointId) throws Exception {
-    // TODO: notify velox
+    // Drain any in-flight data from the Velox pipeline before snapshot.
+    // Flink has already aligned barriers, so no new input will arrive.
+    drainTaskOutput();
     super.prepareSnapshotPreBarrier(checkpointId);
   }
 
   @Override
   public void snapshotState(StateSnapshotContext context) throws Exception {
     // TODO: implement it
-    task.snapshotState(0);
+    task.snapshotState(context.getCheckpointId());
     super.snapshotState(context);
   }
 

gluten-flink/runtime/src/main/java/org/apache/gluten/util/VectorOutputBridge.java

@@ -40,7 +40,7 @@ public interface VectorOutputBridge<OUT> extends Serializable {
    * @param allocator buffer allocator for converting RowVector
    * @param outputType the RowType schema of the output
    */
-  void collect(
+  long collect(
       Output<StreamRecord<OUT>> collector,
       StatefulRecord record,
       BufferAllocator allocator,
@@ -85,7 +85,7 @@ public RowDataOutputBridge() {
     }
 
     @Override
-    public void collect(
+    public long collect(
         Output<StreamRecord<org.apache.flink.table.data.RowData>> collector,
         StatefulRecord record,
         BufferAllocator allocator,
@@ -96,6 +96,7 @@ public void collect(
       for (org.apache.flink.table.data.RowData row : rows) {
         collector.collect(getOrCreateOutputElement().replace(row));
       }
+      return rows.size();
     }
 
     private StreamRecord<org.apache.flink.table.data.RowData> getOrCreateOutputElement() {
@@ -119,12 +120,13 @@ public StatefulRecordOutputBridge() {
     }
 
     @Override
-    public void collect(
+    public long collect(
         Output<StreamRecord<StatefulRecord>> collector,
         StatefulRecord record,
         BufferAllocator allocator,
         RowType outputType) {
       collector.collect(getOrCreateOutputElement().replace(record));
+      return 1;
     }
 
     private StreamRecord<StatefulRecord> getOrCreateOutputElement() {