[GLUTEN-12206][FLINK] Support native checkpoint snapshot and restore for source operator

.github/workflows/flink.yml

@@ -87,8 +87,8 @@ jobs:
           export VELOX_DEPENDENCY_SOURCE=BUNDLED
           export fmt_SOURCE=BUNDLED
           export folly_SOURCE=BUNDLED
-          git clone -b feature/watermark-status-propagation https://github.com/bigo-sg/velox4j.git
-          cd velox4j && git reset --hard 6e2046f
+          git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git
+          cd velox4j && git reset --hard 3ccc4e61dec4d569dc3dd9408b78ee1d4facf3ba
           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/docs/Flink.md

@@ -48,7 +48,7 @@ As some features have not been committed to upstream, you have to use the follow
 ## fetch velox4j code
 git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git
 cd velox4j
-git reset --hard 1b1ad0833220476be6a042b83a474cf58f8f076c
+git reset --hard 3ccc4e61dec4d569dc3dd9408b78ee1d4facf3ba
 mvn clean install -DskipTests -Dgpg.skip -Dspotless.skip=true
 ```
 **Get gluten**

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

@@ -24,7 +24,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 
-public class SourceTaskMetrics {
+public class SourceOperatorMetrics {
 
   private final String keyOperatorType = "operatorType";
   private final String sourceOperatorName = "TableScan";
@@ -35,7 +35,7 @@ public class SourceTaskMetrics {
   private Counter sourceNumBytesOut;
   private long lastUpdateTime = System.currentTimeMillis();
 
-  public SourceTaskMetrics(OperatorMetricGroup metricGroup) {
+  public SourceOperatorMetrics(OperatorMetricGroup metricGroup) {
     sourceNumRecordsOut = metricGroup.getIOMetricGroup().getNumRecordsOutCounter();
     sourceNumBytesOut = metricGroup.getIOMetricGroup().getNumBytesOutCounter();
   }

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

@@ -326,14 +326,13 @@ public String getId() {
 
   @Override
   public void prepareSnapshotPreBarrier(long checkpointId) throws Exception {
-    // TODO: notify velox
     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

@@ -18,7 +18,7 @@
 
 import org.apache.gluten.table.runtime.config.VeloxConnectorConfig;
 import org.apache.gluten.table.runtime.config.VeloxQueryConfig;
-import org.apache.gluten.table.runtime.metrics.SourceTaskMetrics;
+import org.apache.gluten.table.runtime.metrics.SourceOperatorMetrics;
 import org.apache.gluten.vectorized.FlinkRowToVLVectorConvertor;
 
 import io.github.zhztheplayer.velox4j.connector.ConnectorSplit;
@@ -33,6 +33,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;
@@ -44,6 +46,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -64,8 +67,10 @@ public class GlutenSourceFunction<OUT> extends RichParallelSourceFunction<OUT>
   private GlutenSessionResource sessionResource;
   private Query query;
   private SerialTask task;
-  private SourceTaskMetrics taskMetrics;
+  private SourceOperatorMetrics metrics;
   private final Class<OUT> outClass;
+  private transient ListState<String> checkpointState;
+  private transient String[] restoredCheckpointRecords = new String[0];
 
   public GlutenSourceFunction(
       StatefulPlanNode planNode,
@@ -115,7 +120,7 @@ public void run(SourceContext<OUT> sourceContext) throws Exception {
         default:
           return;
       }
-      taskMetrics.updateMetrics(task, id);
+      metrics.updateMetrics(task, id);
     }
   }
 
@@ -208,15 +213,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 {
@@ -252,6 +271,6 @@ private void initSession() {
     task = session.queryOps().execute(query);
     task.addSplit(id, activeSplit);
     task.noMoreSplits(id);
-    taskMetrics = new SourceTaskMetrics(getRuntimeContext().getMetricGroup());
+    metrics = new SourceOperatorMetrics(getRuntimeContext().getMetricGroup());
   }
 }

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

@@ -41,7 +41,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,
@@ -91,7 +91,7 @@ public RowDataOutputBridge() {
     }
 
     @Override
-    public void collect(
+    public long collect(
         Output<StreamRecord<org.apache.flink.table.data.RowData>> collector,
         StatefulRecord record,
         BufferAllocator allocator,
@@ -102,6 +102,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() {
@@ -125,12 +126,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() {
@@ -144,7 +146,7 @@ private StreamRecord<StatefulRecord> getOrCreateOutputElement() {
   class PartitionCommitInfoOutputBridge implements VectorOutputBridge<PartitionCommitInfo> {
 
     @Override
-    public void collect(
+    public long collect(
         Output<StreamRecord<PartitionCommitInfo>> collector,
         StatefulRecord record,
         BufferAllocator allocator,