You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by al...@apache.org on 2019/06/28 07:30:05 UTC

[flink] 03/04: [FLINK-11693] Support FlinkKafkaPartitioner functionality with new KafkaSerializationSchema

This is an automated email from the ASF dual-hosted git repository.

aljoscha pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit d9c71e323d7765d03ecc0359692ec33d79abf426
Author: Aljoscha Krettek <al...@apache.org>
AuthorDate: Wed Jun 26 12:01:02 2019 +0200

    [FLINK-11693] Support FlinkKafkaPartitioner functionality with new KafkaSerializationSchema
---
 .../connectors/kafka/KafkaContextAware.java        | 61 ++++++++++++++++++++++
 .../connectors/kafka/KafkaSerializationSchema.java |  9 +++-
 .../connectors/kafka/FlinkKafkaProducer.java       | 32 +++++++++++-
 3 files changed, 99 insertions(+), 3 deletions(-)

diff --git a/flink-connectors/flink-connector-kafka-base/src/main/java/org/apache/flink/streaming/connectors/kafka/KafkaContextAware.java b/flink-connectors/flink-connector-kafka-base/src/main/java/org/apache/flink/streaming/connectors/kafka/KafkaContextAware.java
new file mode 100644
index 0000000..22399c0
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka-base/src/main/java/org/apache/flink/streaming/connectors/kafka/KafkaContextAware.java
@@ -0,0 +1,61 @@
+/*
+ * 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 org.apache.flink.streaming.connectors.kafka;
+
+import org.apache.flink.annotation.PublicEvolving;
+
+/**
+ * An interface for {@link KafkaSerializationSchema KafkaSerializationSchemas} that need information
+ * about the context where the Kafka Producer is running along with information about the available
+ * partitions.
+ *
+ * <p>You only need to override the methods for the information that you need. However, {@link
+ * #getTargetTopic(Object)} is required because it is used to determine the available partitions.
+ */
+@PublicEvolving
+public interface KafkaContextAware<T> {
+
+
+	/**
+	 * Sets the number of the parallel subtask that the Kafka Producer is running on. The numbering
+	 * starts from 0 and goes up to parallelism-1 (parallelism as returned by {@link
+	 * #setNumParallelInstances(int)}).
+	 */
+	default void setParallelInstanceId(int parallelInstanceId) {
+	}
+
+	/**
+	 * Sets the parallelism with which the parallel task of the Kafka Producer runs.
+	 */
+	default void setNumParallelInstances(int numParallelInstances) {
+	}
+
+	/**
+	 * Sets the available partitions for the topic returned from {@link #getTargetTopic(Object)}.
+	 */
+	default void setPartitions(int[] partitions) {
+	}
+
+	/**
+	 * Returns the topic that the presented element should be sent to. This is not used for setting
+	 * the topic (this is done via the {@link org.apache.kafka.clients.producer.ProducerRecord} that
+	 * is returned from {@link KafkaSerializationSchema#serialize(Object, Long)}, it is only used
+	 * for getting the available partitions that are presented to {@link #setPartitions(int[])}.
+	 */
+	String getTargetTopic(T element);
+}
diff --git a/flink-connectors/flink-connector-kafka-base/src/main/java/org/apache/flink/streaming/connectors/kafka/KafkaSerializationSchema.java b/flink-connectors/flink-connector-kafka-base/src/main/java/org/apache/flink/streaming/connectors/kafka/KafkaSerializationSchema.java
index 31fda0e..37d30fc 100644
--- a/flink-connectors/flink-connector-kafka-base/src/main/java/org/apache/flink/streaming/connectors/kafka/KafkaSerializationSchema.java
+++ b/flink-connectors/flink-connector-kafka-base/src/main/java/org/apache/flink/streaming/connectors/kafka/KafkaSerializationSchema.java
@@ -26,8 +26,13 @@ import javax.annotation.Nullable;
 import java.io.Serializable;
 
 /**
- * A {@link KafkaSerializationSchema} defines how to serialize values of type {@code T} into
- * {@link ProducerRecord ProducerRecords}.
+ * A {@link KafkaSerializationSchema} defines how to serialize values of type {@code T} into {@link
+ * ProducerRecord ProducerRecords}.
+ *
+ * <p>Please also implement {@link KafkaContextAware} if your serialization schema needs
+ * information
+ * about the available partitions and the number of parallel subtasks along with the subtask ID on
+ * which the Kafka Producer is running.
  *
  * @param <T> the type of values being serialized
  */
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/FlinkKafkaProducer.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/FlinkKafkaProducer.java
index a8bff7c..e2c8bbc 100644
--- a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/FlinkKafkaProducer.java
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/FlinkKafkaProducer.java
@@ -770,8 +770,30 @@ public class FlinkKafkaProducer<IN>
 			} else {
 				record = new ProducerRecord<>(targetTopic, null, timestamp, serializedKey, serializedValue);
 			}
-		} else {
+		} else if (kafkaSchema != null) {
+			if (kafkaSchema instanceof KafkaContextAware) {
+				@SuppressWarnings("unchecked")
+				KafkaContextAware<IN> contextAwareSchema =
+						(KafkaContextAware<IN>) kafkaSchema;
+
+				String targetTopic = contextAwareSchema.getTargetTopic(next);
+				if (targetTopic == null) {
+					targetTopic = defaultTopicId;
+				}
+				int[] partitions = topicPartitionsMap.get(targetTopic);
+
+				if (null == partitions) {
+					partitions = getPartitionsByTopic(targetTopic, transaction.producer);
+					topicPartitionsMap.put(targetTopic, partitions);
+				}
+
+				contextAwareSchema.setPartitions(partitions);
+			}
 			record = kafkaSchema.serialize(next, context.timestamp());
+		} else {
+			throw new RuntimeException(
+					"We have neither KafkaSerializationSchema nor KeyedSerializationSchema, this" +
+							"is a bug.");
 		}
 
 		pendingRecords.incrementAndGet();
@@ -1101,6 +1123,14 @@ public class FlinkKafkaProducer<IN>
 			flinkKafkaPartitioner.open(ctx.getIndexOfThisSubtask(), ctx.getNumberOfParallelSubtasks());
 		}
 
+		if (kafkaSchema instanceof KafkaContextAware) {
+			KafkaContextAware<IN> contextAwareSchema =
+					(KafkaContextAware<IN>) kafkaSchema;
+
+			contextAwareSchema.setParallelInstanceId(ctx.getIndexOfThisSubtask());
+			contextAwareSchema.setNumParallelInstances(ctx.getNumberOfParallelSubtasks());
+		}
+
 		LOG.info("Starting FlinkKafkaInternalProducer ({}/{}) to produce into default topic {}",
 			ctx.getIndexOfThisSubtask() + 1, ctx.getNumberOfParallelSubtasks(), defaultTopicId);