kafka parallelism

Author: todd5167

kafka09/kafka09-sink/src/main/java/com/dtstack/flink/sql/sink/kafka/KafkaSink.java

@@ -53,6 +53,9 @@ public class KafkaSink implements RetractStreamTableSink<Row>, IStreamSinkGener<
 	/** Serialization schema for encoding records to Kafka. */
 	protected SerializationSchema serializationSchema;
 
+	protected int parallelism;
+
+
 	@Override
 	public KafkaSink genStreamSink(TargetTableInfo targetTableInfo) {
 		KafkaSinkTableInfo kafka09SinkTableInfo = (KafkaSinkTableInfo) targetTableInfo;
@@ -64,6 +67,11 @@ public KafkaSink genStreamSink(TargetTableInfo targetTableInfo) {
 		}
 		this.fieldTypes = types;
 
+		Integer parallelism = kafka09SinkTableInfo.getParallelism();
+		if (parallelism != null) {
+			this.parallelism = parallelism;
+		}
+
 		properties = new Properties();
 		for (String key : kafka09SinkTableInfo.getKafkaParamKeys()) {
 			properties.setProperty(key, kafka09SinkTableInfo.getKafkaParam(key));
@@ -90,7 +98,7 @@ public void emitDataStream(DataStream<Tuple2<Boolean, Row>> dataStream) {
 
 		DataStream<Row> ds = dataStream.map((Tuple2<Boolean, Row> record) -> {
 			return record.f1;
-		}).returns(getOutputType().getTypeAt(1));
+		}).returns(getOutputType().getTypeAt(1)).setParallelism(parallelism);
 
 		kafkaTableSink.emitDataStream(ds);
 	}

kafka09/kafka09-sink/src/main/java/com/dtstack/flink/sql/sink/kafka/table/KafkaSinkParser.java

@@ -42,6 +42,10 @@ public TableInfo getTableInfo(String tableName, String fieldsInfo, Map<String, O
 
         kafka09SinkTableInfo.setBootstrapServers(MathUtil.getString(props.get(KafkaSinkTableInfo.BOOTSTRAPSERVERS_KEY.toLowerCase())));
         kafka09SinkTableInfo.setTopic(MathUtil.getString(props.get(KafkaSinkTableInfo.TOPIC_KEY.toLowerCase())));
+
+        Integer parallelism = MathUtil.getIntegerVal(props.get(KafkaSinkTableInfo.PARALLELISM_KEY.toLowerCase()));
+        kafka09SinkTableInfo.setParallelism(parallelism);
+
         for (String key : props.keySet()) {
             if (!key.isEmpty() && key.startsWith("kafka.")) {
                 kafka09SinkTableInfo.addKafkaParam(key.substring(6), props.get(key).toString());

kafka09/kafka09-source/src/main/java/com/dtstack/flink/sql/source/kafka/KafkaSource.java

@@ -30,6 +30,7 @@
 import org.apache.flink.api.common.functions.RuntimeContext;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.streaming.api.datastream.DataStreamSource;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer09;
 import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition;
@@ -115,6 +116,12 @@ public Table genStreamSource(SourceTableInfo sourceTableInfo, StreamExecutionEnv
 
         String fields = StringUtils.join(kafka09SourceTableInfo.getFields(), ",");
         String sourceOperatorName = SOURCE_OPERATOR_NAME_TPL.replace("${topic}", topicName).replace("${table}", sourceTableInfo.getName());
-        return tableEnv.fromDataStream(env.addSource(kafkaSrc, sourceOperatorName, typeInformation), fields);
+
+        DataStreamSource kafkaSource = env.addSource(kafkaSrc, sourceOperatorName, typeInformation);
+        Integer parallelism = kafka09SourceTableInfo.getParallelism();
+        if (parallelism != null) {
+            kafkaSource.setParallelism(parallelism);
+        }
+        return tableEnv.fromDataStream(kafkaSource, fields);
     }
 }

kafka10/kafka10-sink/src/main/java/com/dtstack/flink/sql/sink/kafka/KafkaSink.java

@@ -55,6 +55,8 @@ public class KafkaSink implements RetractStreamTableSink<Row>, IStreamSinkGener<
 
     protected Properties properties;
 
+    protected int parallelism;
+
     /** Serialization schema for encoding records to Kafka. */
     protected SerializationSchema serializationSchema;
 
@@ -69,6 +71,11 @@ public KafkaSink genStreamSink(TargetTableInfo targetTableInfo) {
         }
         this.fieldTypes = types;
 
+        Integer parallelism = kafka10SinkTableInfo.getParallelism();
+        if (parallelism != null) {
+            this.parallelism = parallelism;
+        }
+
         properties = new Properties();
         for (String key : kafka10SinkTableInfo.getKafkaParamKeys()) {
             properties.setProperty(key, kafka10SinkTableInfo.getKafkaParam(key));
@@ -95,7 +102,7 @@ public void emitDataStream(DataStream<Tuple2<Boolean, Row>> dataStream) {
 
         DataStream<Row> ds = dataStream.map((Tuple2<Boolean, Row> record) -> {
             return record.f1;
-        }).returns(getOutputType().getTypeAt(1));
+        }).returns(getOutputType().getTypeAt(1)).setParallelism(parallelism);
 
         kafkaTableSink.emitDataStream(ds);
     }

kafka10/kafka10-sink/src/main/java/com/dtstack/flink/sql/sink/kafka/table/KafkaSinkParser.java

@@ -42,6 +42,10 @@ public TableInfo getTableInfo(String tableName, String fieldsInfo, Map<String, O
 
         kafka10SinkTableInfo.setBootstrapServers(MathUtil.getString(props.get(KafkaSinkTableInfo.BOOTSTRAPSERVERS_KEY.toLowerCase())));
         kafka10SinkTableInfo.setTopic(MathUtil.getString(props.get(KafkaSinkTableInfo.TOPIC_KEY.toLowerCase())));
+
+        Integer parallelism = MathUtil.getIntegerVal(props.get(KafkaSinkTableInfo.PARALLELISM_KEY.toLowerCase()));
+        kafka10SinkTableInfo.setParallelism(parallelism);
+
         for (String key : props.keySet()) {
             if (!key.isEmpty() && key.startsWith("kafka.")) {
                 kafka10SinkTableInfo.addKafkaParam(key.substring(6), props.get(key).toString());

kafka10/kafka10-source/src/main/java/com/dtstack/flink/sql/source/kafka/KafkaSource.java

@@ -28,6 +28,7 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.streaming.api.datastream.DataStreamSource;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer010;
 import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition;
@@ -116,6 +117,12 @@ public Table genStreamSource(SourceTableInfo sourceTableInfo, StreamExecutionEnv
 
 		String fields = StringUtils.join(kafka010SourceTableInfo.getFields(), ",");
 		String sourceOperatorName = SOURCE_OPERATOR_NAME_TPL.replace("${topic}", topicName).replace("${table}", sourceTableInfo.getName());
-		return tableEnv.fromDataStream(env.addSource(kafkaSrc, sourceOperatorName, typeInformation), fields);
+
+		DataStreamSource kafkaSource = env.addSource(kafkaSrc, sourceOperatorName, typeInformation);
+		Integer parallelism = kafka010SourceTableInfo.getParallelism();
+		if (parallelism != null) {
+			kafkaSource.setParallelism(parallelism);
+		}
+		return tableEnv.fromDataStream(kafkaSource, fields);
 	}
 }

kafka11/kafka11-sink/src/main/java/com/dtstack/flink/sql/sink/kafka/KafkaSink.java

@@ -51,6 +51,8 @@ public class KafkaSink  implements RetractStreamTableSink<Row>, IStreamSinkGener
 
     protected String topic;
 
+    protected int parallelism;
+
     protected Properties properties;
 
     /** Serialization schema for encoding records to Kafka. */
@@ -67,6 +69,11 @@ public KafkaSink genStreamSink(TargetTableInfo targetTableInfo) {
         }
         this.fieldTypes = types;
 
+        Integer parallelism = kafka11SinkTableInfo.getParallelism();
+        if (parallelism != null) {
+            this.parallelism = parallelism;
+        }
+
         properties = new Properties();
         for (String key : kafka11SinkTableInfo.getKafkaParamKeys()) {
             properties.setProperty(key, kafka11SinkTableInfo.getKafkaParam(key));
@@ -89,7 +96,6 @@ public void emitDataStream(DataStream<Tuple2<Boolean, Row>> dataStream) {
                 serializationSchema
         );
 
-
         DataStream<Row> ds = dataStream.map((Tuple2<Boolean, Row> record) -> {
             return record.f1;
         }).returns(getOutputType().getTypeAt(1));

kafka11/kafka11-sink/src/main/java/com/dtstack/flink/sql/sink/kafka/table/KafkaSinkParser.java

@@ -42,6 +42,10 @@ public TableInfo getTableInfo(String tableName, String fieldsInfo, Map<String, O
 
         kafka11SinkTableInfo.setBootstrapServers(MathUtil.getString(props.get(KafkaSinkTableInfo.BOOTSTRAPSERVERS_KEY.toLowerCase())));
         kafka11SinkTableInfo.setTopic(MathUtil.getString(props.get(KafkaSinkTableInfo.TOPIC_KEY.toLowerCase())));
+
+        Integer parallelism = MathUtil.getIntegerVal(props.get(KafkaSinkTableInfo.PARALLELISM_KEY.toLowerCase()));
+        kafka11SinkTableInfo.setParallelism(parallelism);
+
         for (String key : props.keySet()) {
             if (!key.isEmpty() && key.startsWith("kafka.")) {
                 kafka11SinkTableInfo.addKafkaParam(key.substring(6), props.get(key).toString());

kafka11/kafka11-source/src/main/java/com/dtstack/flink/sql/source/kafka/KafkaSource.java

@@ -28,6 +28,7 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.streaming.api.datastream.DataStreamSource;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer011;
 import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition;
@@ -117,6 +118,12 @@ public Table genStreamSource(SourceTableInfo sourceTableInfo, StreamExecutionEnv
 
 		String fields = StringUtils.join(kafka011SourceTableInfo.getFields(), ",");
 		String sourceOperatorName = SOURCE_OPERATOR_NAME_TPL.replace("${topic}", topicName).replace("${table}", sourceTableInfo.getName());
-		return tableEnv.fromDataStream(env.addSource(kafkaSrc, sourceOperatorName, typeInformation), fields);
+
+		DataStreamSource kafkaSource = env.addSource(kafkaSrc, sourceOperatorName, typeInformation);
+		Integer parallelism = kafka011SourceTableInfo.getParallelism();
+		if (parallelism != null) {
+			kafkaSource.setParallelism(parallelism);
+		}
+		return tableEnv.fromDataStream(kafkaSource, fields);
 	}
 }