[FLINK] Support checkpoint

src/main/cpp/CMakeLists.txt

@@ -84,7 +84,7 @@ else()
     GIT_REPOSITORY
     ${VELOX_SOURCE_URL}
     GIT_TAG
-    dd42008e5a4d05ca25eb1827b235f6a17fb8e35a)
+    ${VELOX_REF})
 endif()
 if(TARGET gflags::gflags_shared)
   get_target_property(GFLAGS_SHARED_LIBRARY gflags::gflags_shared IMPORTED_LOCATION_RELEASE)
@@ -111,7 +111,6 @@ foreach(ROCKSDB_TARGET rocksdb rocksdb-shared RocksDB::rocksdb RocksDB::rocksdb-
   endif()
 endforeach()
 
-
 # Import JniHelpers.
 set(JNI_HELPERS_REPO zhztheplayer/JniHelpersCpp)
 set(JNI_HELPERS_REF 7c58a9a24c87758fe94f33bec6a4463971af1e23)

src/main/cpp/main/velox4j/jni/JniWrapper.cc

@@ -262,19 +262,43 @@ void initializeState(
     jobject javaThis,
     jlong itrId,
     jlong context,
-    jstring keyedStateBackendConfigString) {
+    jstring keyedStateBackendConfigString,
+    jobjectArray checkpointRecords) {
   JNI_METHOD_START
   auto itr = ObjectStore::retrieve<StatefulSerialTask>(itrId);
   spotify::jni::JavaString jTypeJson{env, keyedStateBackendConfigString};
-  itr->initializeState(context, jTypeJson.get());
+  std::vector<std::string> records;
+  if (checkpointRecords != nullptr) {
+    const auto numRecords = env->GetArrayLength(checkpointRecords);
+    records.reserve(numRecords);
+    for (jsize i = 0; i < numRecords; ++i) {
+      auto record = static_cast<jstring>(
+          env->GetObjectArrayElement(checkpointRecords, i));
+      spotify::jni::JavaString jRecord{env, record};
+      records.push_back(jRecord.get());
+      env->DeleteLocalRef(record);
+    }
+  }
+  itr->initializeState(context, jTypeJson.get(), std::move(records));
   JNI_METHOD_END()
 }
 
-void snapshotState(JNIEnv* env, jobject javaThis, jlong itrId, jlong context) {
+jobject
+snapshotState(JNIEnv* env, jobject javaThis, jlong itrId, jlong context) {
   JNI_METHOD_START
   auto itr = ObjectStore::retrieve<StatefulSerialTask>(itrId);
-  itr->snapshotState(context);
-  JNI_METHOD_END()
+  std::vector<std::string> snapshots = itr->snapshotState(context);
+  jclass stringClass = env->FindClass("java/lang/String");
+  jobjectArray result =
+      env->NewObjectArray(snapshots.size(), stringClass, nullptr);
+  for (size_t i = 0; i < snapshots.size(); ++i) {
+    jstring str = env->NewStringUTF(snapshots[i].c_str());
+    env->SetObjectArrayElement(result, i, str);
+    env->DeleteLocalRef(str);
+  }
+  env->DeleteLocalRef(stringClass);
+  return result;
+  JNI_METHOD_END(nullptr)
 }
 
 jobject snapshotSourceState(JNIEnv* env, jobject javaThis, jlong itrId) {
@@ -618,11 +642,12 @@ void JniWrapper::initialize(JNIEnv* env) {
       kTypeLong,
       kTypeLong,
       kTypeString,
+      "[Ljava/lang/String;",
       nullptr);
   addNativeMethod(
       "snapshotState",
       (void*)snapshotState,
-      kTypeVoid,
+      "[Ljava/lang/String;",
       kTypeLong,
       kTypeLong,
       nullptr);

src/main/cpp/main/velox4j/query/StatefulQueryExecutor.cc

@@ -147,23 +147,25 @@ void StatefulSerialTask::notifyWatermark(long watermark) {
 
 void StatefulSerialTask::initializeState(
     long checkpointId,
-    std::string keyedStateBackendConfigString) {
+    std::string keyedStateBackendConfigString,
+    std::vector<std::string> checkpointRecords) {
   folly::dynamic obj = folly::parseJson(keyedStateBackendConfigString);
   std::shared_ptr<const stateful::KeyedStateBackendParameters> params =
       stateful::KeyedStateBackendParameters::create(obj, nullptr);
   if (params &&
       params->getBackendType() == stateful::StateBackendType::ROCKSDB) {
     auto rocksdbParams =
         stateful::RocksDBKeyedStateBackendParameters::create(obj, nullptr);
-    task_->initializeState(rocksdbParams);
+    task_->initializeState(rocksdbParams, checkpointRecords);
   } else {
     // params maybe null, then initialize by using default heap state backend.
-    task_->initializeState(params);
+    task_->initializeState(params, checkpointRecords);
   }
 }
 
-void StatefulSerialTask::snapshotState(long checkpointId) {
+std::vector<std::string> StatefulSerialTask::snapshotState(long checkpointId) {
   sourceStateSnapshots_ = task_->snapshotState(checkpointId);
+  return sourceStateSnapshots_;
 }
 
 std::vector<std::string> StatefulSerialTask::snapshotSourceState() {

src/main/cpp/main/velox4j/query/StatefulQueryExecutor.h

@@ -20,6 +20,7 @@
 #include <velox/experimental/stateful/StatefulTask.h>
 #include <velox/experimental/stateful/StreamElement.h>
 #include <string>
+#include <vector>
 #include "Query.h"
 #include "velox4j/iterator/UpIterator.h"
 #include "velox4j/memory/MemoryManager.h"
@@ -53,9 +54,10 @@ class StatefulSerialTask : public UpIterator {
 
   void initializeState(
       long checkpointId,
-      std::string keyedStateBackendConfigString);
+      std::string keyedStateBackendConfigString,
+      std::vector<std::string> checkpointRecords);
 
-  void snapshotState(long checkpointId);
+  std::vector<std::string> snapshotState(long checkpointId);
 
   std::vector<std::string> snapshotSourceState();
 

src/main/java/io/github/zhztheplayer/velox4j/jni/JniApi.java

@@ -114,12 +114,16 @@ public void notifyWatermark(UpIterator itr, long watermark) {
     jni.notifyWatermark(itr.id(), watermark);
   }
 
-  public void initializeState(UpIterator itr, long context, String keyedStateBackendConfigString) {
-    jni.initializeState(itr.id(), context, keyedStateBackendConfigString);
+  public void initializeState(
+      UpIterator itr,
+      long context,
+      String keyedStateBackendConfigString,
+      String[] checkpointRecords) {
+    jni.initializeState(itr.id(), context, keyedStateBackendConfigString, checkpointRecords);
   }
 
-  public void snapshotState(UpIterator itr, long context) {
-    jni.snapshotState(itr.id(), context);
+  public String[] snapshotState(UpIterator itr, long context) {
+    return jni.snapshotState(itr.id(), context);
   }
 
   public String[] snapshotSourceState(UpIterator itr) {

src/main/java/io/github/zhztheplayer/velox4j/jni/JniWrapper.java

@@ -65,9 +65,10 @@ public long sessionId() {
   // For Flink.
   native void notifyWatermark(long id, long watermark);
 
-  native void initializeState(long id, long context, String keyedStateBackendConfigString);
+  native void initializeState(
+      long id, long context, String keyedStateBackendConfigString, String[] checkpointRecords);
 
-  native void snapshotState(long id, long context);
+  native String[] snapshotState(long id, long context);
 
   native String[] snapshotSourceState(long id);
 

src/main/java/io/github/zhztheplayer/velox4j/query/SerialTask.java

@@ -94,13 +94,20 @@ public SerialTaskStats collectStats() {
 
   public void initializeState(
       long context, KeyedStateBackendParameters keyedStateBackendParameters) {
+    initializeState(context, keyedStateBackendParameters, new String[0]);
+  }
+
+  public void initializeState(
+      long context,
+      KeyedStateBackendParameters keyedStateBackendParameters,
+      String[] checkpointRecords) {
     String keyedStateBackendJsonParameters =
         keyedStateBackendParameters != null ? Serde.toJson(keyedStateBackendParameters) : "{}";
-    jniApi.initializeState(this, context, keyedStateBackendJsonParameters);
+    jniApi.initializeState(this, context, keyedStateBackendJsonParameters, checkpointRecords);
   }
 
-  public void snapshotState(long context) {
-    jniApi.snapshotState(this, context);
+  public String[] snapshotState(long context) {
+    return jniApi.snapshotState(this, context);
   }
 
   public String[] snapshotSourceState() {

src/main/java/io/github/zhztheplayer/velox4j/query/SerialTaskStats.java

@@ -62,4 +62,12 @@ public ObjectNode planStats(String planNodeId) {
     }
     return out.get(0);
   }
+
+  public List<ObjectNode> planStats() {
+    final List<ObjectNode> out = new ArrayList<>();
+    for (JsonNode each : planStatsDynamic) {
+      out.add((ObjectNode) each);
+    }
+    return out;
+  }
 }

src/test/java/io/github/zhztheplayer/velox4j/query/SerialTaskStatsTest.java

@@ -0,0 +1,39 @@
+/*
+* 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 io.github.zhztheplayer.velox4j.query;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class SerialTaskStatsTest {
+
+  @Test
+  public void testPlanStatsReturnsAllPlanStats() {
+    final SerialTaskStats stats =
+        SerialTaskStats.fromJson(
+            "{"
+                + "\"planStats\":["
+                + "{\"planNodeId\":\"scan-1\",\"operatorType\":\"TableScan\"},"
+                + "{\"planNodeId\":\"project-1\",\"operatorType\":\"FilterProject\"}"
+                + "]"
+                + "}");
+
+    Assert.assertEquals(2, stats.planStats().size());
+    Assert.assertEquals("scan-1", stats.planStats().get(0).get("planNodeId").asText());
+    Assert.assertEquals("project-1", stats.planStats().get(1).get("planNodeId").asText());
+  }
+}