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 ba9329e67b..eedc7c0e5c 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 @@ -331,20 +331,27 @@ public void prepareSnapshotPreBarrier(long checkpointId) throws Exception { @Override public void snapshotState(StateSnapshotContext context) throws Exception { - // TODO: implement it - task.snapshotState(context.getCheckpointId()); + snapshotNativeState(context); super.snapshotState(context); } + protected void snapshotNativeState(StateSnapshotContext context) throws Exception { + task.snapshotState(context.getCheckpointId()); + } + @Override public void initializeState(StateInitializationContext context) throws Exception { + initializeNativeState(context); + super.initializeState(context); + } + + protected void initializeNativeState(StateInitializationContext context) throws Exception { if (task == null) { initSession(); } if (!(getKeyedStateBackend() instanceof RocksDBKeyedStateBackend)) { task.initializeState(0, null); } - super.initializeState(context); } @Override diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java index e508fb80c3..0f4eee8411 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java @@ -22,15 +22,31 @@ import io.github.zhztheplayer.velox4j.type.RowType; import org.apache.flink.api.common.TaskInfo; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.connector.file.table.stream.PartitionCommitInfo; +import org.apache.flink.runtime.state.StateInitializationContext; +import org.apache.flink.runtime.state.StateSnapshotContext; import org.apache.flink.streaming.api.watermark.Watermark; import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; import java.util.Map; public class GlutenStreamingFileWriterOperator extends GlutenOneInputOperator { + private static final Logger LOG = + LoggerFactory.getLogger(GlutenStreamingFileWriterOperator.class); + + private transient ListState checkpointState; + private transient String[] restoredCheckpointRecords = new String[0]; + public GlutenStreamingFileWriterOperator( StatefulPlanNode plan, String id, @@ -56,6 +72,44 @@ public void processWatermark(Watermark mark) throws Exception { output.emitWatermark(mark); } + @Override + protected void snapshotNativeState(StateSnapshotContext context) throws Exception { + checkpointState.clear(); + String[] checkpointRecords = task.snapshotState(context.getCheckpointId()); + for (String checkpointRecord : checkpointRecords) { + checkpointState.add(checkpointRecord); + } + } + + @Override + protected void initializeNativeState(StateInitializationContext context) throws Exception { + checkpointState = + context + .getOperatorStateStore() + .getListState( + new ListStateDescriptor<>("gluten-native-file-writer-checkpoint", String.class)); + if (context.isRestored()) { + List records = new ArrayList<>(); + for (String checkpointRecord : checkpointState.get()) { + records.add(checkpointRecord); + } + restoredCheckpointRecords = records.toArray(new String[0]); + } + if (restoredCheckpointRecords != null) { + LOG.info( + "Restore native file writer state for operator {}, restored {}, records {}, contents {}", + getDescription(), + context.isRestored(), + restoredCheckpointRecords.length, + Arrays.toString( + restoredCheckpointRecords != null ? restoredCheckpointRecords : new String[0])); + } + if (task == null) { + initSession(); + } + task.initializeState(0, null, restoredCheckpointRecords); + } + @Override public void notifyCheckpointComplete(long checkpointId) throws Exception { String[] committed = task.notifyCheckpointComplete(checkpointId); diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperatorTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperatorTest.java new file mode 100644 index 0000000000..6f3b0145c1 --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperatorTest.java @@ -0,0 +1,123 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.table.runtime.operators; + +import io.github.zhztheplayer.velox4j.plan.EmptyNode; +import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; +import io.github.zhztheplayer.velox4j.query.SerialTask; +import io.github.zhztheplayer.velox4j.stateful.KeyedStateBackendParameters; +import io.github.zhztheplayer.velox4j.type.IntegerType; +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.runtime.state.StateInitializationContext; +import org.apache.flink.runtime.state.StateInitializationContextImpl; +import org.apache.flink.runtime.state.StateSnapshotContextSynchronousImpl; +import org.apache.flink.streaming.api.operators.collect.utils.MockOperatorStateStore; +import org.apache.flink.table.data.RowData; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +class GlutenStreamingFileWriterOperatorTest { + private static final ListStateDescriptor CHECKPOINT_STATE_DESCRIPTOR = + new ListStateDescriptor<>("gluten-native-file-writer-checkpoint", String.class); + + @Test + void snapshotNativeStateStoresNativeCheckpointRecords() throws Exception { + MockOperatorStateStore stateStore = new MockOperatorStateStore(); + TestSerialTask task = new TestSerialTask(new String[] {"bucket-a", "bucket-b"}); + GlutenStreamingFileWriterOperator operator = createOperator(task); + + operator.initializeNativeState(createInitializationContext(null, stateStore)); + operator.snapshotNativeState(new StateSnapshotContextSynchronousImpl(42L, 0L)); + + ListState checkpointState = stateStore.getListState(CHECKPOINT_STATE_DESCRIPTOR); + assertThat(task.snapshotCheckpointId).isEqualTo(42L); + assertThat(checkpointState.get()).containsExactly("bucket-a", "bucket-b"); + } + + @Test + void initializeNativeStateRestoresCheckpointRecordsToNativeTask() throws Exception { + MockOperatorStateStore stateStore = new MockOperatorStateStore(); + ListState checkpointState = stateStore.getListState(CHECKPOINT_STATE_DESCRIPTOR); + checkpointState.add("bucket-a"); + checkpointState.add("bucket-b"); + TestSerialTask task = new TestSerialTask(new String[0]); + GlutenStreamingFileWriterOperator operator = createOperator(task); + + operator.initializeNativeState(createInitializationContext(1L, stateStore)); + + assertThat(task.initializeContext).isZero(); + assertThat(task.initializeParameters).isNull(); + assertThat(task.restoredCheckpointRecords).containsExactly("bucket-a", "bucket-b"); + } + + private static StateInitializationContext createInitializationContext( + Long restoredCheckpointId, MockOperatorStateStore stateStore) { + return new StateInitializationContextImpl( + restoredCheckpointId, stateStore, null, Collections.emptyList(), Collections.emptyList()); + } + + private static GlutenStreamingFileWriterOperator createOperator(TestSerialTask task) { + RowType rowType = new RowType(List.of("id"), List.of(new IntegerType())); + StatefulPlanNode plan = new StatefulPlanNode("scan", new EmptyNode(rowType)); + GlutenStreamingFileWriterOperator operator = + new GlutenStreamingFileWriterOperator<>( + plan, + "scan", + rowType, + Map.of("scan", rowType), + RowData.class, + "test native file writer"); + operator.task = task; + return operator; + } + + private static class TestSerialTask extends SerialTask { + private final String[] checkpointRecords; + private long snapshotCheckpointId = -1L; + private long initializeContext = -1L; + private KeyedStateBackendParameters initializeParameters; + private String[] restoredCheckpointRecords; + + private TestSerialTask(String[] checkpointRecords) { + super(null, 0L); + this.checkpointRecords = checkpointRecords; + } + + @Override + public String[] snapshotState(long checkpointId) { + snapshotCheckpointId = checkpointId; + return checkpointRecords; + } + + @Override + public void initializeState( + long context, KeyedStateBackendParameters parameters, String[] checkpointRecords) { + initializeContext = context; + initializeParameters = parameters; + restoredCheckpointRecords = checkpointRecords; + } + } +}