[FLINK] Bind native callback target for Gluten operators

Author: lgbo-ustc

.github/workflows/flink.yml

@@ -69,8 +69,8 @@ jobs:
           export VELOX_DEPENDENCY_SOURCE=BUNDLED
           export fmt_SOURCE=BUNDLED
           export folly_SOURCE=BUNDLED
-          git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git
-          cd velox4j && git reset --hard 115edf79d265a61c30d45dfcc6ce932ad92378ca
+          git clone -b feature/native-callback-bridge https://github.com/bigo-sg/velox4j.git
+          cd velox4j && git reset --hard c2c193ead7bf88f632838bca58921ef699c89c11
           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/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java

@@ -31,6 +31,7 @@
 import io.github.zhztheplayer.velox4j.query.Query;
 import io.github.zhztheplayer.velox4j.query.SerialTask;
 import io.github.zhztheplayer.velox4j.serde.Serde;
+import io.github.zhztheplayer.velox4j.stateful.NativeCallbackTarget;
 import io.github.zhztheplayer.velox4j.stateful.StatefulElement;
 import io.github.zhztheplayer.velox4j.stateful.StatefulRecord;
 import io.github.zhztheplayer.velox4j.stateful.StatefulWatermark;
@@ -49,7 +50,7 @@
 
 /** Calculate operator in gluten, which will call Velox to run. */
 public class GlutenOneInputOperator<IN, OUT> extends TableStreamOperator<OUT>
-    implements OneInputStreamOperator<IN, OUT>, GlutenOperator {
+    implements OneInputStreamOperator<IN, OUT>, GlutenOperator, NativeCallbackTarget {
 
   private final StatefulPlanNode glutenPlan;
   private final String id;
@@ -143,6 +144,7 @@ void initSession() {
             VeloxQueryConfig.getConfig(getRuntimeContext()),
             VeloxConnectorConfig.getConfig(getRuntimeContext()));
     task = sessionResource.getSession().queryOps().execute(query);
+    task.bindNativeCallbackTarget(this);
     task.addSplit(
         id, new ExternalStreamConnectorSplit("connector-external-stream", inputQueue.id()));
     task.noMoreSplits(id);
@@ -186,6 +188,11 @@ public void scheduleProcessElementOnMailbox() {
     scheduleDrainOnMailbox(this::drainTaskOutput);
   }
 
+  @Override
+  public void onProcessElement() {
+    scheduleProcessElementOnMailbox();
+  }
+
   @Override
   public void processElementInternal() {
     drainOutput(this::drainTaskOutput);

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

@@ -28,6 +28,7 @@
 import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode;
 import io.github.zhztheplayer.velox4j.query.Query;
 import io.github.zhztheplayer.velox4j.query.SerialTask;
+import io.github.zhztheplayer.velox4j.stateful.NativeCallbackTarget;
 import io.github.zhztheplayer.velox4j.stateful.StatefulElement;
 import io.github.zhztheplayer.velox4j.stateful.StatefulRecord;
 import io.github.zhztheplayer.velox4j.stateful.StatefulWatermark;
@@ -50,7 +51,7 @@
  * instead of flink RowData.
  */
 public class GlutenTwoInputOperator<IN, OUT> extends AbstractStreamOperator<OUT>
-    implements TwoInputStreamOperator<IN, IN, OUT>, GlutenOperator {
+    implements TwoInputStreamOperator<IN, IN, OUT>, GlutenOperator, NativeCallbackTarget {
 
   private static final Logger LOG = LoggerFactory.getLogger(GlutenTwoInputOperator.class);
 
@@ -139,6 +140,11 @@ public void scheduleProcessElementOnMailbox() {
     scheduleDrainOnMailbox(this::drainTaskOutput);
   }
 
+  @Override
+  public void onProcessElement() {
+    scheduleProcessElementOnMailbox();
+  }
+
   @Override
   public void processElement1(StreamRecord<IN> element) {
     StatefulRecord statefulRecord =
@@ -296,6 +302,7 @@ private void initSession() {
             VeloxQueryConfig.getConfig(getRuntimeContext()),
             VeloxConnectorConfig.getConfig(getRuntimeContext()));
     task = sessionResource.getSession().queryOps().execute(query);
+    task.bindNativeCallbackTarget(this);
 
     ExternalStreamConnectorSplit leftSplit =
         new ExternalStreamConnectorSplit("connector-external-stream", leftInputQueue.id());