[AMORO-3359] only set parallelism on parallelismSourceFunciton (#3360)

issue_3359 only set parallelism on parallelismSourceFunciton
apache · Dec 13, 2024 · 6ef1be7 · 6ef1be7
1 parent 24cec41
commit 6ef1be7
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/...link/amoro-mixed-flink-common/src/main/java/org/apache/amoro/flink/table/FlinkSource.java b/...link/amoro-mixed-flink-common/src/main/java/org/apache/amoro/flink/table/FlinkSource.java
@@ -38,6 +38,7 @@
 import org.apache.flink.streaming.api.datastream.DataStreamSource;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.streaming.api.functions.source.InputFormatSourceFunction;
+import org.apache.flink.streaming.api.functions.source.ParallelSourceFunction;
 import org.apache.flink.streaming.api.functions.source.SourceFunction;
 import org.apache.flink.streaming.api.operators.OneInputStreamOperatorFactory;
 import org.apache.flink.streaming.api.operators.StreamSource;
@@ -289,12 +290,13 @@ private DataStream<RowData> wrapKrb(DataStream<RowData> ds) {
         DataStreamSource sourceStream =
             env.addSource(functionProxy, tfSource.getName(), tfSource.getOutputType());
         context.generateUid(MIXED_FORMAT_FILE_TRANSFORMATION).ifPresent(sourceStream::uid);
-        return sourceStream
-            .setParallelism(scanParallelism)
-            .transform(
-                tf.getName(),
-                tf.getOutputType(),
-                new UnkeyedInputFormatOperatorFactory(inputFormatProxyFactory));
+        if (sourceStream instanceof ParallelSourceFunction) {
+          sourceStream.setParallelism(scanParallelism);
+        }
+        return sourceStream.transform(
+            tf.getName(),
+            tf.getOutputType(),
+            new UnkeyedInputFormatOperatorFactory(inputFormatProxyFactory));
       }
 
       LegacySourceTransformation tfSource = (LegacySourceTransformation) origin;