[FLINK] Support ORC filesystem sink format

.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 f8684c84ac68729f1370be1ee6a80a0fe4738818
           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/patches/fix-velox4j.patch

@@ -1,3 +1,16 @@
+diff --git a/src/main/cpp/CMakeLists.txt b/src/main/cpp/CMakeLists.txt
+index 9ed0551..5ded913 100644
+--- a/src/main/cpp/CMakeLists.txt
++++ b/src/main/cpp/CMakeLists.txt
+@@ -82,7 +82,7 @@ else()
+     GIT_REPOSITORY
+     ${VELOX_SOURCE_URL}
+     GIT_TAG
+-    c9efdf308a94ea1cb846e2440aeb975572497d87)
++    bfa91ab98d6b53fccf2c862cdc926832835db35c)
+ endif()
+ if(TARGET gflags::gflags_shared)
+   get_target_property(GFLAGS_SHARED_LIBRARY gflags::gflags_shared
 diff --git a/src/main/cpp/main/CMakeLists.txt b/src/main/cpp/main/CMakeLists.txt
 index ada06f8..7ed80bb 100644
 --- a/src/main/cpp/main/CMakeLists.txt
@@ -15,3 +28,31 @@ index ada06f8..7ed80bb 100644
      message(\"Found all shared 3rd dependencies of velox4j: \${3rd_deps}\")
      foreach(dep IN LISTS 3rd_deps)
          get_filename_component(link_name \${dep} NAME)
+diff --git a/src/main/cpp/main/velox4j/query/StatefulQueryExecutor.cc b/src/main/cpp/main/velox4j/query/StatefulQueryExecutor.cc
+index 2357cf2..159e3d6 100644
+--- a/src/main/cpp/main/velox4j/query/StatefulQueryExecutor.cc
++++ b/src/main/cpp/main/velox4j/query/StatefulQueryExecutor.cc
+@@ -81,15 +81,8 @@ StatefulSerialTask::StatefulSerialTask(
+ }
+
+ StatefulSerialTask::~StatefulSerialTask() {
+-  if (task_ != nullptr && task_->isRunning()) {
+-    // TODO: add a method to finish the task and set state.
+-    task_->finish();
+-    // FIXME: Calling .wait() may take no effect in single thread execution
+-    //  mode.
+-    task_->requestCancel().wait();
+-  }
+   if (task_ != nullptr) {
+-    task_->setNativeCallbackBridge(nullptr);
++    task_->close();
+   }
+   task_.reset();
+ }
+diff --git a/src/main/cpp/velox-ref.txt b/src/main/cpp/velox-ref.txt
+index c24a828..47d494e 100644
+--- a/src/main/cpp/velox-ref.txt
++++ b/src/main/cpp/velox-ref.txt
+@@ -1 +1 @@
+-c9efdf308a94ea1cb846e2440aeb975572497d87
++bfa91ab98d6b53fccf2c862cdc926832835db35c

gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecExchange.java

@@ -143,7 +143,6 @@ protected Transformation<RowData> translateToPlanInternal(
           // TODO: velox's parallelism need to be set here, as some nodes need it.
           // should set it when operator init.
           parallelism = inputTransform.getParallelism();
-          keySelector = new GlutenKeySelector();
           final ExecEdge inputEdge = getInputEdges().get(0);
           io.github.zhztheplayer.velox4j.type.RowType glutenInputType =
               (io.github.zhztheplayer.velox4j.type.RowType)
@@ -181,7 +180,8 @@ protected Transformation<RowData> translateToPlanInternal(
                   parallelism,
                   false);
           partitioner =
-              new GlutenKeyGroupStreamPartitioner(keySelector, DEFAULT_LOWER_BOUND_MAX_PARALLELISM);
+              new GlutenKeyGroupStreamPartitioner(
+                  new GlutenKeySelector(), keySelector, DEFAULT_LOWER_BOUND_MAX_PARALLELISM);
         } else {
           parallelism = ExecutionConfig.PARALLELISM_DEFAULT;
           partitioner =

gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecJoin.java

@@ -41,7 +41,6 @@
 import org.apache.flink.FlinkVersion;
 import org.apache.flink.api.dag.Transformation;
 import org.apache.flink.configuration.ReadableConfig;
-import org.apache.flink.streaming.api.operators.StreamOperator;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
 import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
 import org.apache.flink.table.data.RowData;
@@ -325,30 +324,15 @@ protected Transformation<RowData> translateToPlanInternal(
     }
 
     final RowType returnType = (RowType) getOutputType();
-    final TwoInputTransformation<RowData, RowData, RowData> transform;
-    if (operator instanceof GlutenTwoInputOperator) {
-      transform =
-          ExecNodeUtil.createTwoInputTransformation(
-              leftTransform,
-              rightTransform,
-              createTransformationMeta(JOIN_TRANSFORMATION, config),
-              new GlutenTwoInputOperatorFactory<RowData, RowData, RowData>(
-                  (StreamOperator<RowData>) operator),
-              InternalTypeInfo.of(returnType),
-              leftTransform.getParallelism(),
-              0L,
-              false);
-    } else {
-      transform =
-          ExecNodeUtil.createTwoInputTransformation(
-              leftTransform,
-              rightTransform,
-              createTransformationMeta(JOIN_TRANSFORMATION, config),
-              operator,
-              InternalTypeInfo.of(returnType),
-              leftTransform.getParallelism(),
-              false);
-    }
+    final TwoInputTransformation<RowData, RowData, RowData> transform =
+        ExecNodeUtil.createTwoInputTransformation(
+            leftTransform,
+            rightTransform,
+            createTransformationMeta(JOIN_TRANSFORMATION, config),
+            new GlutenTwoInputOperatorFactory<>(operator),
+            InternalTypeInfo.of(returnType),
+            leftTransform.getParallelism(),
+            0L);
 
     // set KeyType and Selector for state
     RowDataKeySelector leftSelect =

gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecWindowJoin.java

@@ -37,7 +37,6 @@
 import org.apache.flink.FlinkVersion;
 import org.apache.flink.api.dag.Transformation;
 import org.apache.flink.configuration.ReadableConfig;
-import org.apache.flink.streaming.api.operators.StreamOperator;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
 import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
 import org.apache.flink.table.api.TableException;
@@ -247,12 +246,10 @@ protected Transformation<RowData> translateToPlanInternal(
             leftTransform,
             rightTransform,
             createTransformationMeta(WINDOW_JOIN_TRANSFORMATION, config),
-            new GlutenTwoInputOperatorFactory<RowData, RowData, RowData>(
-                (StreamOperator<RowData>) operator),
+            new GlutenTwoInputOperatorFactory<>(operator),
             InternalTypeInfo.of(returnType),
             leftTransform.getParallelism(),
-            0L,
-            false);
+            0L);
 
     // set KeyType and Selector for state
     RowDataKeySelector leftSelect =

gluten-flink/pom.xml

@@ -31,6 +31,8 @@
   <properties>
     <flink.version>1.19.2</flink.version>
     <velox4j.version>0.1.0-SNAPSHOT</velox4j.version>
+    <protobuf.version>3.25.5</protobuf.version>
+    <hadoop.version>2.7.2</hadoop.version>
     <scala.binary.version>2.12</scala.binary.version>
     <maven.compiler.source>11</maven.compiler.source>
     <maven.compiler.target>11</maven.compiler.target>

gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java

@@ -16,10 +16,8 @@
  */
 package org.apache.gluten.client;
 
-import org.apache.gluten.streaming.api.operators.GlutenOneInputOperatorFactory;
 import org.apache.gluten.streaming.api.operators.GlutenOperator;
 import org.apache.gluten.streaming.api.operators.GlutenStreamSource;
-import org.apache.gluten.streaming.api.operators.GlutenTwoInputOperatorFactory;
 import org.apache.gluten.table.runtime.keyselector.GlutenKeySelector;
 import org.apache.gluten.table.runtime.operators.GlutenOneInputOperator;
 import org.apache.gluten.table.runtime.operators.GlutenSourceFunction;
@@ -243,16 +241,16 @@ private void createOffloadedOneInputOperator(
     Class<?> outClass = supportsVectorOutput ? StatefulRecord.class : RowData.class;
     GlutenOneInputOperator<?, ?> newOneInputOp =
         sourceOperator.cloneWithInputOutputClasses(planNode, inClass, outClass);
-    // setStreamOperator would wrap this in Flink's SimpleOperatorFactory and skip Gluten-specific
-    // mailbox binding performed by GlutenOneInputOperatorFactory.
-    offloadedOpConfig.setStreamOperatorFactory(new GlutenOneInputOperatorFactory<>(newOneInputOp));
+    offloadedOpConfig.setStreamOperator(newOneInputOp);
     if (supportsVectorOutput) {
       setOffloadedOutputSerializer(offloadedOpConfig, sourceOperator);
     }
     if (supportsVectorInput) {
       setOffloadedInputSerializer(offloadedOpConfig, sourceOperator);
       setOffloadedStatePartitioner(
           sourceOpConfig, offloadedOpConfig, 0, sourceOperator.getDescription());
+    } else {
+      offloadedOpConfig.setupNetworkInputs(sourceOpConfig.getTypeSerializerIn1(userClassloader));
     }
   }
 
@@ -283,17 +281,18 @@ private void createOffloadedTwoInputOperator(
             sourceOperator.getOutputTypes(),
             inClass,
             outClass);
-    // setStreamOperator would wrap this in Flink's SimpleOperatorFactory, which only initializes
-    // ProcessingTimeService for Flink AbstractStreamOperator. GlutenTwoInputOperator uses
-    // GlutenAbstractStreamOperator so it needs the Gluten-specific factory.
-    offloadedOpConfig.setStreamOperatorFactory(new GlutenTwoInputOperatorFactory<>(newTwoInputOp));
+    offloadedOpConfig.setStreamOperator(newTwoInputOp);
     offloadedOpConfig.setStatePartitioner(0, new GlutenKeySelector());
     offloadedOpConfig.setStatePartitioner(1, new GlutenKeySelector());
     if (supportsVectorOutput) {
       setOffloadedOutputSerializer(offloadedOpConfig, sourceOperator);
     }
     if (supportsVectorInput) {
       setOffloadedInputSerializersForTwoInputOperator(offloadedOpConfig, sourceOperator);
+    } else {
+      offloadedOpConfig.setupNetworkInputs(
+          sourceOpConfig.getTypeSerializerIn1(userClassloader),
+          sourceOpConfig.getTypeSerializerIn2(userClassloader));
     }
   }
 
@@ -486,6 +485,9 @@ private void updateStatePartitioners(StreamConfig sourceConfig, StreamConfig off
 
   private boolean supportsVectorOutput(
       OperatorChainSlice chainSlice, OperatorChainSliceGraph chainSliceGraph, JobVertex jobVertex) {
+    if (isStreamExecExchange(chainSlice)) {
+      return false;
+    }
     List<Integer> downstreamSliceIds = chainSlice.getOutputs();
 
     // If chainSlice has downstream in the operator chain, check these downstream
@@ -498,11 +500,22 @@ private boolean supportsVectorOutput(
     }
   }
 
+  private boolean isStreamExecExchange(OperatorChainSlice chainSlice) {
+    for (StreamConfig opConfig : chainSlice.getOperatorConfigs()) {
+      Optional<GlutenOperator> glutenOperator = Utils.getGlutenOperator(opConfig, userClassloader);
+      if (glutenOperator.isPresent()
+          && "StreamExecExchange".equals(glutenOperator.get().getDescription())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private boolean checkDownstreamSlicesInChain(
       OperatorChainSliceGraph chainSliceGraph, List<Integer> downstreamSliceIds) {
     for (Integer downstreamSliceId : downstreamSliceIds) {
       OperatorChainSlice downstreamSlice = chainSliceGraph.getSlice(downstreamSliceId);
-      if (!downstreamSlice.isOffloadable()) {
+      if (!downstreamSlice.isOffloadable() || isStreamExecExchange(downstreamSlice)) {
         return false;
       }
       if (!downstreamSlice.getInputs().stream()
@@ -521,16 +534,7 @@ private boolean checkDownstreamVerticesFromJobGraph(JobVertex jobVertex) {
       // conversion)
       return true;
     }
-    // Check if all downstream vertices' operators are gluten operators
-    for (JobVertex downstreamVertex : downstreamVertices) {
-      StreamConfig downstreamStreamConfig = new StreamConfig(downstreamVertex.getConfiguration());
-      Optional<GlutenOperator> glutenOperator =
-          Utils.getGlutenOperator(downstreamStreamConfig, userClassloader);
-      if (!glutenOperator.isPresent()) {
-        return false;
-      }
-    }
-    return true;
+    return false;
   }
 
   private boolean supportsVectorInput(
@@ -569,16 +573,7 @@ private boolean checkUpstreamVerticesFromJobGraph(JobVertex jobVertex) {
       // If there is no upstream JobVertex, can input RowVector (no upstream needs conversion)
       return true;
     }
-    // Check if the last operator of all upstream vertices are gluten operators
-    for (JobVertex upstreamVertex : upstreamVertices) {
-      StreamConfig lastOperatorConfig = findLastOperatorInChain(upstreamVertex);
-      Optional<GlutenOperator> glutenOperator =
-          Utils.getGlutenOperator(lastOperatorConfig, userClassloader);
-      if (!glutenOperator.isPresent()) {
-        return false;
-      }
-    }
-    return true;
+    return false;
   }
 
   /** Get all downstream JobVertices of a JobVertex. */

gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/runtime/partitioner/GlutenKeyGroupStreamPartitioner.java

@@ -26,6 +26,7 @@
 import org.apache.flink.streaming.runtime.partitioner.ConfigurableStreamPartitioner;
 import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.table.data.RowData;
 import org.apache.flink.util.Preconditions;
 
 import java.util.Objects;
@@ -41,13 +42,22 @@ public class GlutenKeyGroupStreamPartitioner extends StreamPartitioner<StatefulR
   private static final long serialVersionUID = 1L;
 
   private final KeySelector<StatefulRecord, Integer> keySelector;
+  private final KeySelector<RowData, ?> rowDataKeySelector;
 
   private int maxParallelism;
 
   public GlutenKeyGroupStreamPartitioner(
       KeySelector<StatefulRecord, Integer> keySelector, int maxParallelism) {
+    this(keySelector, null, maxParallelism);
+  }
+
+  public GlutenKeyGroupStreamPartitioner(
+      KeySelector<StatefulRecord, Integer> keySelector,
+      KeySelector<RowData, ?> rowDataKeySelector,
+      int maxParallelism) {
     Preconditions.checkArgument(maxParallelism > 0, "Number of key-groups must be > 0!");
     this.keySelector = Preconditions.checkNotNull(keySelector);
+    this.rowDataKeySelector = rowDataKeySelector;
     this.maxParallelism = maxParallelism;
   }
 
@@ -57,12 +67,19 @@ public int getMaxParallelism() {
 
   @Override
   public int selectChannel(SerializationDelegate<StreamRecord<StatefulRecord>> record) {
+    Object value = record.getInstance().getValue();
     try {
-      int channel = keySelector.getKey(record.getInstance().getValue());
-      return channel;
+      if (value instanceof StatefulRecord) {
+        return keySelector.getKey((StatefulRecord) value);
+      }
+      Preconditions.checkState(
+          rowDataKeySelector != null,
+          "RowData key selector is required when partitioning non-StatefulRecord values.");
+      Object key = rowDataKeySelector.getKey((RowData) value);
+      return KeyGroupRangeAssignment.assignKeyToParallelOperator(
+          key, maxParallelism, numberOfChannels);
     } catch (Exception e) {
-      throw new RuntimeException(
-          "Could not extract key from " + record.getInstance().getValue(), e);
+      throw new RuntimeException("Could not extract key from " + value, e);
     }
   }
 
@@ -104,11 +121,13 @@ public boolean equals(Object o) {
       return false;
     }
     final GlutenKeyGroupStreamPartitioner that = (GlutenKeyGroupStreamPartitioner) o;
-    return maxParallelism == that.maxParallelism && keySelector.equals(that.keySelector);
+    return maxParallelism == that.maxParallelism
+        && keySelector.equals(that.keySelector)
+        && Objects.equals(rowDataKeySelector, that.rowDataKeySelector);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(super.hashCode(), keySelector, maxParallelism);
+    return Objects.hash(super.hashCode(), keySelector, rowDataKeySelector, maxParallelism);
   }
 }