diff --git a/.github/workflows/flink.yml b/.github/workflows/flink.yml index b66bddf4279..e32cc1b7fcb 100644 --- a/.github/workflows/flink.yml +++ b/.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 .. diff --git a/gluten-flink/docs/Flink.md b/gluten-flink/docs/Flink.md index 9c1900f06b1..3e3adde696a 100644 --- a/gluten-flink/docs/Flink.md +++ b/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** diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/metrics/SourceTaskMetrics.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/metrics/SourceOperatorMetrics.java similarity index 96% rename from gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/metrics/SourceTaskMetrics.java rename to gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/metrics/SourceOperatorMetrics.java index 5bbbb028b6a..8dcc80438f0 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/metrics/SourceTaskMetrics.java +++ b/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(); } diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java index 5729a60b41c..ba9329e67b8 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java +++ b/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); } diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java index 5dfc8447a90..a47b6c7c17b 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java +++ b/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 extends RichParallelSourceFunction private GlutenSessionResource sessionResource; private Query query; private SerialTask task; - private SourceTaskMetrics taskMetrics; + private SourceOperatorMetrics metrics; private final Class outClass; + private transient ListState checkpointState; + private transient String[] restoredCheckpointRecords = new String[0]; public GlutenSourceFunction( StatefulPlanNode planNode, @@ -115,7 +120,7 @@ public void run(SourceContext 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 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()); } } diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/util/VectorOutputBridge.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/util/VectorOutputBridge.java index bbc4ea796ce..0e749daa829 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/util/VectorOutputBridge.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/util/VectorOutputBridge.java @@ -41,7 +41,7 @@ public interface VectorOutputBridge extends Serializable { * @param allocator buffer allocator for converting RowVector * @param outputType the RowType schema of the output */ - void collect( + long collect( Output> collector, StatefulRecord record, BufferAllocator allocator, @@ -91,7 +91,7 @@ public RowDataOutputBridge() { } @Override - public void collect( + public long collect( Output> 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 getOrCreateOutputElement() { @@ -125,12 +126,13 @@ public StatefulRecordOutputBridge() { } @Override - public void collect( + public long collect( Output> collector, StatefulRecord record, BufferAllocator allocator, RowType outputType) { collector.collect(getOrCreateOutputElement().replace(record)); + return 1; } private StreamRecord getOrCreateOutputElement() { @@ -144,7 +146,7 @@ private StreamRecord getOrCreateOutputElement() { class PartitionCommitInfoOutputBridge implements VectorOutputBridge { @Override - public void collect( + public long collect( Output> collector, StatefulRecord record, BufferAllocator allocator,