Support watermark status propagation

src/main/cpp/CMakeLists.txt

@@ -82,7 +82,7 @@ else()
     GIT_REPOSITORY
     ${VELOX_SOURCE_URL}
     GIT_TAG
-    c9efdf308a94ea1cb846e2440aeb975572497d87)
+    ${VELOX_REF})
 endif()
 if(TARGET gflags::gflags_shared)
   get_target_property(GFLAGS_SHARED_LIBRARY gflags::gflags_shared

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

@@ -180,6 +180,22 @@ jobject statefulTaskGet(JNIEnv* env, jobject javaThis, jlong itrId) {
     jobject result =
         env->NewObject(resultClass, constructor, id, watermark->timestamp());
 
+    env->DeleteLocalRef(id);
+    env->DeleteLocalRef(resultClass);
+    return result;
+  } else if (element->isWatermarkStatus()) {
+    auto watermarkStatus =
+        std::static_pointer_cast<stateful::WatermarkStatus>(element);
+
+    jclass resultClass = env->FindClass(
+        "io/github/zhztheplayer/velox4j/stateful/StatefulWatermarkStatus");
+    jmethodID constructor =
+        env->GetMethodID(resultClass, "<init>", "(Ljava/lang/String;Z)V");
+
+    jstring id = env->NewStringUTF(watermarkStatus->nodeId().c_str());
+    jobject result =
+        env->NewObject(resultClass, constructor, id, watermarkStatus->idle());
+
     env->DeleteLocalRef(id);
     env->DeleteLocalRef(resultClass);
     return result;
@@ -257,6 +273,29 @@ void notifyWatermark(
   JNI_METHOD_END()
 }
 
+void notifyIndexedWatermarkStatus(
+    JNIEnv* env,
+    jobject javaThis,
+    jlong itrId,
+    jboolean idle,
+    jint index) {
+  JNI_METHOD_START
+  auto itr = ObjectStore::retrieve<StatefulSerialTask>(itrId);
+  itr->notifyWatermarkStatus(idle, index);
+  JNI_METHOD_END()
+}
+
+void notifyWatermarkStatus(
+    JNIEnv* env,
+    jobject javaThis,
+    jlong itrId,
+    jboolean idle) {
+  JNI_METHOD_START
+  auto itr = ObjectStore::retrieve<StatefulSerialTask>(itrId);
+  itr->notifyWatermarkStatus(idle);
+  JNI_METHOD_END()
+}
+
 void initializeState(
     JNIEnv* env,
     jobject javaThis,
@@ -611,6 +650,21 @@ void JniWrapper::initialize(JNIEnv* env) {
       kTypeLong,
       kTypeLong,
       nullptr);
+  addNativeMethod(
+      "notifyIndexedWatermarkStatus",
+      (void*)notifyIndexedWatermarkStatus,
+      kTypeVoid,
+      kTypeLong,
+      kTypeBool,
+      kTypeInt,
+      nullptr);
+  addNativeMethod(
+      "notifyWatermarkStatus",
+      (void*)notifyWatermarkStatus,
+      kTypeVoid,
+      kTypeLong,
+      kTypeBool,
+      nullptr);
   addNativeMethod(
       "initializeState",
       (void*)initializeState,

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

@@ -137,14 +137,22 @@ void StatefulSerialTask::setNativeCallbackBridge(
   task_->setNativeCallbackBridge(std::move(callbackBridge));
 }
 
-void StatefulSerialTask::notifyWatermark(long watermark, int index) {
+void StatefulSerialTask::notifyWatermark(long watermark, int32_t index) {
   task_->notifyWatermark(watermark, index);
 }
 
 void StatefulSerialTask::notifyWatermark(long watermark) {
   task_->notifyWatermark(watermark);
 }
 
+void StatefulSerialTask::notifyWatermarkStatus(bool idle, int32_t index) {
+  task_->notifyWatermarkStatus(idle, index);
+}
+
+void StatefulSerialTask::notifyWatermarkStatus(bool idle) {
+  task_->notifyWatermarkStatus(idle);
+}
+
 void StatefulSerialTask::initializeState(
     long checkpointId,
     std::string keyedStateBackendConfigString) {

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

@@ -19,6 +19,7 @@
 
 #include <velox/experimental/stateful/StatefulTask.h>
 #include <velox/experimental/stateful/StreamElement.h>
+#include <cstdint>
 #include <string>
 #include "Query.h"
 #include "velox4j/iterator/UpIterator.h"
@@ -47,10 +48,14 @@ class StatefulSerialTask : public UpIterator {
       std::shared_ptr<facebook::velox::stateful::NativeCallbackBridge>
           callbackBridge);
 
-  void notifyWatermark(long watermark, int index);
+  void notifyWatermark(long watermark, int32_t index);
 
   void notifyWatermark(long watermark);
 
+  void notifyWatermarkStatus(bool idle, int32_t index);
+
+  void notifyWatermarkStatus(bool idle);
+
   void initializeState(
       long checkpointId,
       std::string keyedStateBackendConfigString);

src/main/cpp/velox-ref.txt

@@ -1 +1 @@
-c9efdf308a94ea1cb846e2440aeb975572497d87
+d267ca9bf825ed8b366d9c2e90eb50f719d6479c

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

@@ -114,6 +114,14 @@ public void notifyWatermark(UpIterator itr, long watermark) {
     jni.notifyWatermark(itr.id(), watermark);
   }
 
+  public void notifyWatermarkStatus(UpIterator itr, boolean idle, int index) {
+    jni.notifyIndexedWatermarkStatus(itr.id(), idle, index);
+  }
+
+  public void notifyWatermarkStatus(UpIterator itr, boolean idle) {
+    jni.notifyWatermarkStatus(itr.id(), idle);
+  }
+
   public void initializeState(UpIterator itr, long context, String keyedStateBackendConfigString) {
     jni.initializeState(itr.id(), context, keyedStateBackendConfigString);
   }

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

@@ -65,6 +65,12 @@ public long sessionId() {
   // For Flink.
   native void notifyWatermark(long id, long watermark);
 
+  // For Flink.
+  native void notifyIndexedWatermarkStatus(long id, boolean idle, int index);
+
+  // For Flink.
+  native void notifyWatermarkStatus(long id, boolean idle);
+
   native void initializeState(long id, long context, String keyedStateBackendConfigString);
 
   native void snapshotState(long id, long context);

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

@@ -75,6 +75,16 @@ public void notifyWatermark(long watermark) {
     jniApi.notifyWatermark(this, watermark);
   }
 
+  // This method is for Flink. idle == true means IDLE; false means ACTIVE.
+  public void notifyWatermarkStatus(boolean idle, int index) {
+    jniApi.notifyWatermarkStatus(this, idle, index);
+  }
+
+  // This method is for Flink. idle == true means IDLE; false means ACTIVE.
+  public void notifyWatermarkStatus(boolean idle) {
+    jniApi.notifyWatermarkStatus(this, idle);
+  }
+
   @Override
   public long id() {
     return id;

src/main/java/io/github/zhztheplayer/velox4j/stateful/StatefulElement.java

@@ -31,10 +31,18 @@ public String getNodeId() {
 
   public abstract boolean isRecord();
 
+  public boolean isWatermarkStatus() {
+    return false;
+  }
+
   public StatefulWatermark asWatermark() {
     return (StatefulWatermark) this;
   }
 
+  public StatefulWatermarkStatus asWatermarkStatus() {
+    return (StatefulWatermarkStatus) this;
+  }
+
   public StatefulRecord asRecord() {
     return (StatefulRecord) this;
   }

src/main/java/io/github/zhztheplayer/velox4j/stateful/StatefulWatermarkStatus.java

@@ -0,0 +1,46 @@
+/*
+* 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.stateful;
+
+public class StatefulWatermarkStatus extends StatefulElement {
+  // true means WatermarkStatus.IDLE; false means WatermarkStatus.ACTIVE.
+  private final boolean idle;
+
+  public StatefulWatermarkStatus(String nodeId, boolean idle) {
+    super(nodeId);
+    this.idle = idle;
+  }
+
+  public boolean isIdle() {
+    return idle;
+  }
+
+  @Override
+  public boolean isWatermark() {
+    return false;
+  }
+
+  @Override
+  public boolean isRecord() {
+    return false;
+  }
+
+  @Override
+  public boolean isWatermarkStatus() {
+    return true;
+  }
+}