KAFKA-4936: Add dynamic routing in Streams

streams/src/main/java/org/apache/kafka/streams/Topology.java

@@ -519,30 +519,32 @@ public synchronized <K, V> Topology addSink(final String name,
 
     /**
      * Add a new sink that forwards records from upstream parent processor and/or source nodes to Kafka topics dynamically.
+     * The topics that it may ever send to should be pre-created.
      * The sink will use the {@link StreamsConfig#DEFAULT_KEY_SERDE_CLASS_CONFIG default key serializer} and
      * {@link StreamsConfig#DEFAULT_VALUE_SERDE_CLASS_CONFIG default value serializer} specified in the
      * {@link StreamsConfig stream configuration}.
      *
      * @param name              the unique name of the sink
-     * @param topicExtractor    the mapper to dynamically choose the name of the Kafka topic to which this sink should write per reach record
+     * @param topicExtractor    the extractor to determine the name of the Kafka topic to which this sink should write for reach record
      * @param parentNames       the name of one or more source or processor nodes whose output records this sink should consume
-     *                          and write to its topic
+     *                          and dynamically write to topics
      * @return                  itself
-     * @throws TopologyException if parent processor is not added yet, or if this processor's name is equal to the parent's name
+     * @throws TopologyException if parent processor is not added yet, or if this processor's name is not unique
      * @see #addSink(String, String, StreamPartitioner, String...)
      * @see #addSink(String, String, Serializer, Serializer, String...)
      * @see #addSink(String, String, Serializer, Serializer, StreamPartitioner, String...)
      */
     public synchronized <K, V> Topology addSink(final String name,
-                                         final TopicNameExtractor<K, V> topicExtractor,
-                                         final String... parentNames) {
+                                                final TopicNameExtractor<K, V> topicExtractor,
+                                                final String... parentNames) {
         internalTopologyBuilder.addSink(name, topicExtractor, null, null, null, parentNames);
         return this;
     }
 
     /**
-     * Add a new sink that forwards records from upstream parent processor and/or source nodes to  Kafka topics dynamically,
+     * Add a new sink that forwards records from upstream parent processor and/or source nodes to Kafka topics dynamically,
      * using the supplied partitioner.
+     * The topics that it may ever send to should be pre-created.
      * The sink will use the {@link StreamsConfig#DEFAULT_KEY_SERDE_CLASS_CONFIG default key serializer} and
      * {@link StreamsConfig#DEFAULT_VALUE_SERDE_CLASS_CONFIG default value serializer} specified in the
      * {@link StreamsConfig stream configuration}.
@@ -555,12 +557,12 @@ public synchronized <K, V> Topology addSink(final String name,
      * records among partitions using Kafka's default partitioning logic.
      *
      * @param name              the unique name of the sink
-     * @param topicExtractor    the mapper to dynamically choose the name of the Kafka topic to which this sink should write per reach record
+     * @param topicExtractor    the extractor to determine the name of the Kafka topic to which this sink should write for reach record
      * @param partitioner       the function that should be used to determine the partition for each record processed by the sink
      * @param parentNames       the name of one or more source or processor nodes whose output records this sink should consume
-     *                          and write to its topic
+     *                          and dynamically write to topics
      * @return                  itself
-     * @throws TopologyException if parent processor is not added yet, or if this processor's name is equal to the parent's name
+     * @throws TopologyException if parent processor is not added yet, or if this processor's name is not unique
      * @see #addSink(String, String, String...)
      * @see #addSink(String, String, Serializer, Serializer, String...)
      * @see #addSink(String, String, Serializer, Serializer, StreamPartitioner, String...)
@@ -575,20 +577,21 @@ public synchronized <K, V> Topology addSink(final String name,
 
     /**
      * Add a new sink that forwards records from upstream parent processor and/or source nodes to Kafka topics dynamically.
+     * The topics that it may ever send to should be pre-created.
      * The sink will use the specified key and value serializers.
      *
      * @param name              the unique name of the sink
-     * @param topicExtractor    the mapper to dynamically choose the name of the Kafka topic to which this sink should write per reach record
+     * @param topicExtractor    the extractor to determine the name of the Kafka topic to which this sink should write for reach record
      * @param keySerializer     the {@link Serializer key serializer} used when consuming records; may be null if the sink
      *                          should use the {@link StreamsConfig#DEFAULT_KEY_SERDE_CLASS_CONFIG default key serializer} specified in the
      *                          {@link StreamsConfig stream configuration}
      * @param valueSerializer   the {@link Serializer value serializer} used when consuming records; may be null if the sink
      *                          should use the {@link StreamsConfig#DEFAULT_VALUE_SERDE_CLASS_CONFIG default value serializer} specified in the
      *                          {@link StreamsConfig stream configuration}
      * @param parentNames       the name of one or more source or processor nodes whose output records this sink should consume
-     *                          and write to its topic
+     *                          and dynamically write to topics
      * @return                  itself
-     * @throws TopologyException if parent processor is not added yet, or if this processor's name is equal to the parent's name
+     * @throws TopologyException if parent processor is not added yet, or if this processor's name is not unique
      * @see #addSink(String, String, String...)
      * @see #addSink(String, String, StreamPartitioner, String...)
      * @see #addSink(String, String, Serializer, Serializer, StreamPartitioner, String...)
@@ -604,10 +607,11 @@ public synchronized <K, V> Topology addSink(final String name,
 
     /**
      * Add a new sink that forwards records from upstream parent processor and/or source nodes to Kafka topics dynamically.
+     * The topics that it may ever send to should be pre-created.
      * The sink will use the specified key and value serializers, and the supplied partitioner.
      *
      * @param name              the unique name of the sink
-     * @param topicExtractor    the mapper to dynamically choose the name of the Kafka topic to which this sink should write per reach record
+     * @param topicExtractor    the extractor to determine the name of the Kafka topic to which this sink should write for reach record
      * @param keySerializer     the {@link Serializer key serializer} used when consuming records; may be null if the sink
      *                          should use the {@link StreamsConfig#DEFAULT_KEY_SERDE_CLASS_CONFIG default key serializer} specified in the
      *                          {@link StreamsConfig stream configuration}
@@ -616,9 +620,9 @@ public synchronized <K, V> Topology addSink(final String name,
      *                          {@link StreamsConfig stream configuration}
      * @param partitioner       the function that should be used to determine the partition for each record processed by the sink
      * @param parentNames       the name of one or more source or processor nodes whose output records this sink should consume
-     *                          and write to its topic
+     *                          and dynamically write to topics
      * @return                  itself
-     * @throws TopologyException if parent processor is not added yet, or if this processor's name is equal to the parent's name
+     * @throws TopologyException if parent processor is not added yet, or if this processor's name is not unique
      * @see #addSink(String, String, String...)
      * @see #addSink(String, String, StreamPartitioner, String...)
      * @see #addSink(String, String, Serializer, Serializer, String...)

streams/src/main/java/org/apache/kafka/streams/kstream/KStream.java

@@ -471,18 +471,17 @@ void to(final String topic,
     /**
      * Dynamically materialize this stream to topics using default serializers specified in the config and producer's
      * {@link DefaultPartitioner}.
-     * The topic names for each record to send to is dynamically determined based on the {@link KeyValueMapper}.
+     * The topic names for each record to send to is dynamically determined based on the {@link TopicNameExtractor}.
      *
-     * @param topicExtractor    the dynamic topic chooser
+     * @param topicExtractor    the extractor to determine the name of the Kafka topic to write to for reach record
      */
     void to(final TopicNameExtractor<K, V> topicExtractor);
 
     /**
-     * Dynamically materialize this stream to topics using default serializers specified in the config and producer's
-     * {@link DefaultPartitioner}.
-     * The topic names for each record to send to is dynamically determined based on the {@link KeyValueMapper}.
+     * Dynamically materialize this stream to topics using default serializers specified in the config.
+     * The topic names for each record to send to is dynamically determined based on the {@link TopicNameExtractor}.
      *
-     * @param topicExtractor    the dynamic topic chooser
+     * @param topicExtractor    the extractor to determine the name of the Kafka topic to write to for reach record
      * @param produced          the options to use when producing to the topic
      */
     void to(final TopicNameExtractor<K, V> topicExtractor,

streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamImpl.java

@@ -314,6 +314,7 @@ public void to(final String topic, final Produced<K, V> produced) {
 
     @Override
     public void to(final TopicNameExtractor<K, V> topicExtractor) {
+        Objects.requireNonNull(topicExtractor, "topic extractor can't be null");
         to(topicExtractor, Produced.<K, V>with(null, null, null));
     }
 

streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java

@@ -37,6 +37,10 @@ public class WindowedStreamPartitioner<K, V> implements StreamPartitioner<Window
         this.serializer = serializer;
     }
 
+    public String topic() {
+        return this.topic;
+    }
+
     /**
      * WindowedStreamPartitioner determines the partition number for a record with the given windowed key and value
      * and the current number of partitions. The partition number id determined by the original key of the windowed key
@@ -47,10 +51,18 @@ public class WindowedStreamPartitioner<K, V> implements StreamPartitioner<Window
      * @param numPartitions the total number of partitions
      * @return an integer between 0 and {@code numPartitions-1}, or {@code null} if the default partitioning logic should be used
      */
+    @Override
     public Integer partition(final Windowed<K> windowedKey, final V value, final int numPartitions) {
         final byte[] keyBytes = serializer.serializeBaseKey(topic, windowedKey);
 
         // hash the keyBytes to choose a partition
         return toPositive(Utils.murmur2(keyBytes)) % numPartitions;
     }
+
+    public Integer partition(final String topic, final Windowed<K> windowedKey, final V value, final int numPartitions) {
+        final byte[] keyBytes = serializer.serializeBaseKey(topic, windowedKey);
+
+        // hash the keyBytes to choose a partition
+        return toPositive(Utils.murmur2(keyBytes)) % numPartitions;
+    }
 }

streams/src/main/java/org/apache/kafka/streams/processor/TopicNameExtractor.java

@@ -20,20 +20,19 @@
 import org.apache.kafka.streams.processor.internals.RecordContext;
 
 /**
- * An interface that allows to dynamically determine the topic name to send to at the sink node of the topology.
+ * An interface that allows to dynamically determine the name of the Kafka topic to send at the sink node of the topology.
  */
 @InterfaceStability.Evolving
 public interface TopicNameExtractor<K, V> {
 
     /**
      * Extracts the topic name to send to. The topic name must be pre-existed, since the Kafka Streams library will not
-     * try to automatically create the topic with the extracted name, and will fail with a timeout exception if the topic
-     * does not exist in the Kafka cluster.
+     * try to automatically create the topic with the extracted name.
      *
      * @param key           the record key
-     * @param value         the record value payload
+     * @param value         the record value
      * @param recordContext current context metadata of the record
-     * @return the topic name to send to
+     * @return              the topic name this record should be sent to
      */
     String extract(K key, V value, RecordContext recordContext);
 }

streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java

@@ -39,6 +39,8 @@
 import org.apache.kafka.streams.errors.ProductionExceptionHandler;
 import org.apache.kafka.streams.errors.ProductionExceptionHandler.ProductionExceptionHandlerResponse;
 import org.apache.kafka.streams.errors.StreamsException;
+import org.apache.kafka.streams.kstream.Windowed;
+import org.apache.kafka.streams.kstream.internals.WindowedStreamPartitioner;
 import org.apache.kafka.streams.processor.StreamPartitioner;
 import org.slf4j.Logger;
 
@@ -86,7 +88,11 @@ public <K, V> void send(final String topic,
         if (partitioner != null) {
             final List<PartitionInfo> partitions = producer.partitionsFor(topic);
             if (partitions.size() > 0) {
-                partition = partitioner.partition(key, value, partitions.size());
+                if (partitioner instanceof WindowedStreamPartitioner && ((WindowedStreamPartitioner) partitioner).topic() == null) {
+                    partition = ((WindowedStreamPartitioner) partitioner).partition(topic, (Windowed) key, value, partitions.size());
+                } else {
+                    partition = partitioner.partition(key, value, partitions.size());
+                }
             } else {
                 throw new StreamsException("Could not get partition information for topic '" + topic + "'." +
                     " This can happen if the topic does not exist.");

streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamImplTest.java

@@ -232,11 +232,11 @@ public void shouldSendDataToDynamicTopics() {
         stream.to(new TopicNameExtractor<String, String>() {
             @Override
             public String extract(String key, String value, RecordContext recordContext) {
-                return recordContext.topic() + "-" + key + "-topic";
+                return recordContext.topic() + "-" + key + "-" + value.substring(0, 1);
             }
         }, Produced.with(Serdes.String(), Serdes.String()));
-        builder.stream(input + "-a-topic", stringConsumed).process(processorSupplier);
-        builder.stream(input + "-b-topic", stringConsumed).process(processorSupplier);
+        builder.stream(input + "-a-v", stringConsumed).process(processorSupplier);
+        builder.stream(input + "-b-v", stringConsumed).process(processorSupplier);
 
         try (final TopologyTestDriver driver = new TopologyTestDriver(builder.build(), props)) {
             driver.pipeInput(recordFactory.create(input, "a", "v1"));

streams/streams-scala/src/main/scala/org/apache/kafka/streams/scala/kstream/KStream.scala

@@ -283,7 +283,7 @@ class KStream[K, V](val inner: KStreamJ[K, V]) {
     * // to the through call
     * }}}
     *
-    * @param extractor the mapper from key value to topic name
+    * @param extractor the extractor to determine the name of the Kafka topic to write to for reach record
     * @param (implicit) produced the instance of Produced that gives the serdes and `StreamPartitioner`
     * @see `org.apache.kafka.streams.kstream.KStream#to`
     */