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

mjsax pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git


The following commit(s) were added to refs/heads/trunk by this push:
     new cba250ead18 KAFKA-20197: Headers aware StreamPartitioner (#22570)
cba250ead18 is described below

commit cba250ead18e8fa72dbf3b7539e1bed76fb6fce2
Author: Uladzislau Blok <[email protected]>
AuthorDate: Sun Jul 26 03:15:28 2026 +0200

    KAFKA-20197: Headers aware StreamPartitioner (#22570)
    
    Implementation for header-aware partitioners (KIP-1321).
    
    This PR allows to pass in a records `Headers` object into
    `StreamPartitioner`.
    
    Reviewers: Alieh Saeedi <[email protected]>, Matthias J. Sax
    <[email protected]>
---
 .../org/apache/kafka/streams/KafkaStreams.java     |  34 +++-
 .../streams/kstream/internals/KTableImpl.java      |  85 ++++++++-
 .../internals/WindowedStreamPartitioner.java       |  19 +-
 .../kafka/streams/processor/StreamPartitioner.java |  17 +-
 .../internals/DefaultStreamPartitioner.java        |   9 +-
 .../processor/internals/RecordCollectorImpl.java   |   2 +-
 .../processor/internals/StreamsMetadataState.java  |  27 ++-
 .../KafkaStreamsNamedTopologyWrapper.java          |   6 +-
 .../org/apache/kafka/streams/KafkaStreamsTest.java |  68 ++++++-
 .../kstream/internals/KStreamRepartitionTest.java  |  10 +-
 .../KTableKTableForeignKeyJoinScenarioTest.java    |  91 ++++++++++
 .../internals/WindowedStreamPartitionerTest.java   |  47 +++--
 .../internals/DefaultStreamPartitionerTest.java    |  86 +++++++++
 .../processor/internals/ProcessorTopologyTest.java |   6 +
 .../processor/internals/RecordCollectorTest.java   |  65 ++++++-
 .../internals/StreamsMetadataStateTest.java        | 197 +++++++++++++++++----
 16 files changed, 681 insertions(+), 88 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java 
b/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java
index ad9aada7c0d..aecdf3f67bb 100644
--- a/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java
+++ b/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java
@@ -30,6 +30,8 @@ import org.apache.kafka.common.annotation.InterfaceAudience;
 import org.apache.kafka.common.annotation.InterfaceStability.Evolving;
 import org.apache.kafka.common.annotation.SuppressKafkaInternalApiUsage;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.metrics.KafkaMetricsContext;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
@@ -1784,11 +1786,24 @@ public class KafkaStreams implements AutoCloseable {
         return streamsMetadataState.allMetadataForStore(storeName);
     }
 
+    /**
+     * Finds the metadata containing the active hosts and standby hosts where 
the key being queried would reside,
+     * without requiring record headers to be provided.
+     * <p>
+     * If your partitioner or serializer makes use of headers, use the {@link 
#queryMetadataForKey(String, Object, Headers, Serializer) Headers} overload, 
otherwise the returned metadata may not match where the key actually resides.
+     */
+    public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
+                                                    final K key,
+                                                    final Serializer<K> 
keySerializer) {
+        return queryMetadataForKey(storeName, key, new RecordHeaders(), 
keySerializer);
+    }
+
     /**
      * Finds the metadata containing the active hosts and standby hosts where 
the key being queried would reside.
      *
      * @param storeName     the {@code storeName} to find metadata for
      * @param key           the key to find metadata for
+     * @param headers       the record headers
      * @param keySerializer serializer for the key
      * @param <K>           key type
      * Returns {@link KeyQueryMetadata} containing all metadata about hosting 
the given key for the given store,
@@ -1796,9 +1811,22 @@ public class KafkaStreams implements AutoCloseable {
      */
     public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
                                                     final K key,
+                                                    final Headers headers,
                                                     final Serializer<K> 
keySerializer) {
         validateIsRunningOrRebalancing();
-        return streamsMetadataState.keyQueryMetadataForKey(storeName, key, 
keySerializer);
+        return streamsMetadataState.keyQueryMetadataForKey(storeName, key, 
headers, keySerializer);
+    }
+
+    /**
+     * Finds the metadata containing the active hosts and standby hosts where 
the key being queried would reside,
+     * using the supplied partitioner and without requiring record headers to 
be provided.
+     * <p>
+     * If your partitioner or serializer makes use of headers, use the {@link 
#queryMetadataForKey(String, Object, Headers, StreamPartitioner) Headers} 
overload, otherwise the returned metadata may not match where the key actually 
resides.
+     */
+    public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
+                                                    final K key,
+                                                    final StreamPartitioner<? 
super K, ?> partitioner) {
+        return queryMetadataForKey(storeName, key, new RecordHeaders(), 
partitioner);
     }
 
     /**
@@ -1806,6 +1834,7 @@ public class KafkaStreams implements AutoCloseable {
      *
      * @param storeName     the {@code storeName} to find metadata for
      * @param key           the key to find metadata for
+     * @param headers       the record headers
      * @param partitioner   the partitioner to be used to locate the host for 
the key
      * @param <K>           key type
      * Returns {@link KeyQueryMetadata} containing all metadata about hosting 
the given key for the given store, using
@@ -1813,9 +1842,10 @@ public class KafkaStreams implements AutoCloseable {
      */
     public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
                                                     final K key,
+                                                    final Headers headers,
                                                     final StreamPartitioner<? 
super K, ?> partitioner) {
         validateIsRunningOrRebalancing();
-        return streamsMetadataState.keyQueryMetadataForKey(storeName, key, 
partitioner);
+        return streamsMetadataState.keyQueryMetadataForKey(storeName, key, 
headers, partitioner);
     }
 
     /**
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
 
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
index 5f0d043c7e4..1a190037141 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.kstream.internals;
 
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.serialization.Serde;
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.streams.KeyValue;
@@ -1279,7 +1280,7 @@ public class KTableImpl<K, S, V> extends 
AbstractStream<K, V> implements KTable<
         final StreamPartitioner<KO, SubscriptionWrapper<K>> 
subscriptionSinkPartitioner =
                 tableJoinedInternal.otherPartitioner() == null
                         ? null
-                        : (topic, key, val, numPartitions) -> 
getPartition.apply(tableJoinedInternal.otherPartitioner().partitions(topic, 
key, null, numPartitions));
+                        : new 
SubscriptionSinkPartitioner<>(tableJoinedInternal.otherPartitioner(), 
getPartition);
 
         final StreamSinkNode<KO, SubscriptionWrapper<K>> subscriptionSink = 
new StreamSinkNode<>(
             renamed.suffixWithOrElseGet("-subscription-registration-sink", 
builder, SINK_NAME),
@@ -1345,16 +1346,12 @@ public class KTableImpl<K, S, V> extends 
AbstractStream<K, V> implements KTable<
         final String finalRepartitionTopicName = 
renamed.suffixWithOrElseGet("-subscription-response", builder, 
SUBSCRIPTION_RESPONSE) + TOPIC_SUFFIX;
         
builder.internalTopologyBuilder.addInternalTopic(finalRepartitionTopicName, 
InternalTopicProperties.empty());
 
-        final StreamPartitioner<K, SubscriptionResponseWrapper<VO>> 
defaultForeignResponseSinkPartitioner =
-                (topic, key, subscriptionResponseWrapper, numPartitions) -> {
-                    final Integer partition = 
subscriptionResponseWrapper.primaryPartition();
-                    return partition == null ? Optional.empty() : 
Optional.of(Collections.singleton(partition));
-                };
+        final StreamPartitioner<K, SubscriptionResponseWrapper<VO>> 
defaultForeignResponseSinkPartitioner = new 
DefaultForeignResponseSinkPartitioner<>();
 
         final StreamPartitioner<K, SubscriptionResponseWrapper<VO>> 
foreignResponseSinkPartitioner =
                 tableJoinedInternal.partitioner() == null
                         ? defaultForeignResponseSinkPartitioner
-                        : (topic, key, val, numPartitions) -> 
getPartition.apply(tableJoinedInternal.partitioner().partitions(topic, key, 
null, numPartitions));
+                        : new 
ForeignResponseSinkPartitioner<>(tableJoinedInternal.partitioner(), 
getPartition);
 
         final StreamSinkNode<K, SubscriptionResponseWrapper<VO>> 
foreignResponseSink =
             new StreamSinkNode<>(
@@ -1438,4 +1435,78 @@ public class KTableImpl<K, S, V> extends 
AbstractStream<K, V> implements KTable<
             tableNode.setOutputVersioned(materializedInternal.storeSupplier() 
instanceof VersionedBytesStoreSupplier);
         }
     }
+
+    private static class SubscriptionSinkPartitioner<K, KO> implements 
StreamPartitioner<KO, SubscriptionWrapper<K>> {
+        private final StreamPartitioner<KO, Void> otherPartitioner;
+        private final Function<Optional<Set<Integer>>, Optional<Set<Integer>>> 
getPartition;
+
+        private SubscriptionSinkPartitioner(
+            final StreamPartitioner<KO, Void> otherPartitioner,
+            final Function<Optional<Set<Integer>>, Optional<Set<Integer>>> 
getPartition
+        ) {
+            this.otherPartitioner = otherPartitioner;
+            this.getPartition = getPartition;
+        }
+
+        @SuppressWarnings("removal")
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final KO 
key, final SubscriptionWrapper<K> value, final int numPartitions) {
+            throw new UnsupportedOperationException("This method is deprecated 
and should not be called.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic,
+                                                 final KO key,
+                                                 final SubscriptionWrapper<K> 
value,
+                                                 final Headers headers,
+                                                 final int numPartitions) {
+            return getPartition.apply(otherPartitioner.partitions(topic, key, 
null, headers, numPartitions));
+        }
+    }
+
+    private static class ForeignResponseSinkPartitioner<K, VO> implements 
StreamPartitioner<K, SubscriptionResponseWrapper<VO>> {
+        private final StreamPartitioner<K, Void> partitioner;
+        private final Function<Optional<Set<Integer>>, Optional<Set<Integer>>> 
getPartition;
+
+        private ForeignResponseSinkPartitioner(
+            final StreamPartitioner<K, Void> partitioner,
+            final Function<Optional<Set<Integer>>, Optional<Set<Integer>>> 
getPartition
+        ) {
+            this.partitioner = partitioner;
+            this.getPartition = getPartition;
+        }
+
+        @SuppressWarnings("removal")
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final K 
key, final SubscriptionResponseWrapper<VO> value, final int numPartitions) {
+            throw new UnsupportedOperationException("This method is deprecated 
and should not be called.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic,
+                                                 final K key,
+                                                 final 
SubscriptionResponseWrapper<VO> value,
+                                                 final Headers headers,
+                                                 final int numPartitions) {
+            return getPartition.apply(partitioner.partitions(topic, key, null, 
headers, numPartitions));
+        }
+    }
+
+    private static class DefaultForeignResponseSinkPartitioner<K, VO> 
implements StreamPartitioner<K, SubscriptionResponseWrapper<VO>> {
+        @SuppressWarnings("removal")
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final K 
key, final SubscriptionResponseWrapper<VO> value, final int numPartitions) {
+            throw new UnsupportedOperationException("This method is deprecated 
and should not be called.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic,
+                                                 final K key,
+                                                 final 
SubscriptionResponseWrapper<VO> value,
+                                                 final Headers headers,
+                                                 final int numPartitions) {
+            final Integer partition = value.primaryPartition();
+            return partition == null ? Optional.empty() : 
Optional.of(Collections.singleton(partition));
+        }
+    }
 }
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
 
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
index 54735b98bb1..b2386e31321 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
@@ -17,7 +17,7 @@
 package org.apache.kafka.streams.kstream.internals;
 
 import org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
-import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.streams.kstream.Windowed;
 import org.apache.kafka.streams.processor.StreamPartitioner;
 
@@ -33,21 +33,28 @@ public class WindowedStreamPartitioner<K, V> implements 
StreamPartitioner<Window
         this.serializer = serializer;
     }
 
+    @SuppressWarnings({"removal"})
+    @Override
+    public Optional<Set<Integer>> partitions(final String topic, final 
Windowed<K> windowedKey, final V value, final int numPartitions) {
+        throw new UnsupportedOperationException("This method is deprecated and 
should not be called.");
+    }
+
     /**
      * 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
      * using the same logic as DefaultPartitioner so that the topic is 
partitioned by the original key.
      *
-     * @param topic the topic name this record is sent to
-     * @param windowedKey the key of the record
-     * @param value the value of the record
+     * @param topic         the topic name this record is sent to
+     * @param windowedKey   the key of the record
+     * @param value         the value of the record
+     * @param headers       the record headers
      * @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 Optional<Set<Integer>> partitions(final String topic, final 
Windowed<K> windowedKey, final V value, final int numPartitions) {
+    public Optional<Set<Integer>> partitions(final String topic, final 
Windowed<K> windowedKey, final V value, final Headers headers, final int 
numPartitions) {
         // for windowed key, the key bytes should never be null
-        final byte[] keyBytes = serializer.serializeBaseKey(topic, new 
RecordHeaders(), windowedKey);
+        final byte[] keyBytes = serializer.serializeBaseKey(topic, headers, 
windowedKey);
 
         // stick with the same built-in partitioner util functions that 
producer used
         // to make sure its behavior is consistent with the producer
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
index 20e6cfb5d81..bb81affa811 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.streams.processor;
 
 import org.apache.kafka.common.annotation.InterfaceAudience;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.streams.Topology;
 
 import java.util.Optional;
@@ -44,7 +45,7 @@ import java.util.Set;
  * for that topic.
  * <p>
  * All StreamPartitioner implementations should be stateless and a pure 
function so they can be shared across topic and sink nodes.
- * 
+ *
  * @param <K> the type of keys
  * @param <V> the type of values
  * @see Topology#addSink(String, String, 
org.apache.kafka.common.serialization.Serializer,
@@ -56,16 +57,26 @@ import java.util.Set;
 public interface StreamPartitioner<K, V> {
 
     /**
-     * Determine the number(s) of the partition(s) to which a record with the 
given key and value should be sent, 
+     * @deprecated Since 4.4. Use {@link #partitions(String, Object, Object, 
Headers, int)} instead.
+     * This method is planned to be removed in 5.0.
+     */
+    @Deprecated(since = "4.4", forRemoval = true)
+    Optional<Set<Integer>> partitions(String topic, K key, V value, int 
numPartitions);
+
+    /**
+     * Determine the number(s) of the partition(s) to which a record with the 
given key and value should be sent,
      * for the given topic and current partition count
      * @param topic the topic name this record is sent to
      * @param key the key of the record
      * @param value the value of the record
+     * @param headers the record headers
      * @param numPartitions the total number of partitions
      * @return an Optional of Set of integers between 0 and {@code 
numPartitions-1},
      * Empty optional means using default partitioner
      * Optional of an empty set means the record won't be sent to any 
partitions i.e drop it.
      * Optional of Set of integers means the partitions to which the record 
should be sent to.
      * */
-    Optional<Set<Integer>> partitions(String topic, K key, V value, int 
numPartitions);
+    default Optional<Set<Integer>> partitions(final String topic, final K key, 
final V value, final Headers headers, final int numPartitions) {
+        return partitions(topic, key, value, numPartitions);
+    }
 }
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
index 439eb379e00..868ed183236 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.streams.processor.internals;
 
 import org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.streams.processor.StreamPartitioner;
 
@@ -32,9 +33,15 @@ public class DefaultStreamPartitioner<K, V> implements 
StreamPartitioner<K, V> {
         this.keySerializer = keySerializer;
     }
 
+    @SuppressWarnings({"removal"})
     @Override
     public Optional<Set<Integer>> partitions(final String topic, final K key, 
final V value, final int numPartitions) {
-        final byte[] keyBytes = keySerializer.serialize(topic, key);
+        throw new UnsupportedOperationException("This method is deprecated and 
should not be called.");
+    }
+
+    @Override
+    public Optional<Set<Integer>> partitions(final String topic, final K key, 
final V value, final Headers headers, final int numPartitions) {
+        final byte[] keyBytes = keySerializer.serialize(topic, headers, key);
 
         // if the key bytes are not available, we just return empty optional 
to let the producer decide
         // which partition to send internally; otherwise stick with the same 
built-in partitioner
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
index dd9baeb6743..19f57dd7995 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
@@ -159,7 +159,7 @@ public class RecordCollectorImpl implements RecordCollector 
{
                 );
             }
             if (!partitions.isEmpty()) {
-                final Optional<Set<Integer>> maybeMulticastPartitions = 
partitioner.partitions(topic, key, value, partitions.size());
+                final Optional<Set<Integer>> maybeMulticastPartitions = 
partitioner.partitions(topic, key, value, headers, partitions.size());
                 if (maybeMulticastPartitions.isEmpty()) {
                     // A null//empty partition indicates we should use the 
default partitioner
                     send(topic, key, value, headers, null, timestamp, 
keySerializer, valueSerializer, processorNodeId, context);
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
index 5d84c599cf4..ea2bc70b336 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
@@ -18,6 +18,7 @@ package org.apache.kafka.streams.processor.internals;
 
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.utils.internals.LogContext;
 import org.apache.kafka.streams.KafkaStreams;
@@ -190,7 +191,7 @@ public class StreamsMetadataState {
     /**
      * Find the {@link KeyQueryMetadata}s for a given storeName and key. This 
method will use the
      * {@link DefaultStreamPartitioner} to locate the store. If a custom 
partitioner has been used
-     * please use {@link StreamsMetadataState#keyQueryMetadataForKey(String, 
Object, StreamPartitioner)} instead.
+     * please use {@link StreamsMetadataState#keyQueryMetadataForKey(String, 
Object, Headers, StreamPartitioner)} instead.
      *
      * Note: the key may not exist in the {@link 
org.apache.kafka.streams.processor.StateStore},
      * this method provides a way of finding which {@link KeyQueryMetadata} it 
would exist on.
@@ -205,6 +206,7 @@ public class StreamsMetadataState {
      */
     public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final 
String storeName,
                                                                     final K 
key,
+                                                                    final 
Headers headers,
                                                                     final 
Serializer<K> keySerializer) {
         Objects.requireNonNull(keySerializer, "keySerializer can't be null");
         if (topologyMetadata.hasNamedTopologies()) {
@@ -214,19 +216,22 @@ public class StreamsMetadataState {
         }
         return keyQueryMetadataForKey(storeName,
                                       key,
+                                      headers,
                                       new 
DefaultStreamPartitioner<>(keySerializer));
     }
 
     /**
-     * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object, 
Serializer)}
+     * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object, 
Headers, Serializer)}
      */
     public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final 
String storeName,
                                                                     final K 
key,
+                                                                    final 
Headers headers,
                                                                     final 
Serializer<K> keySerializer,
                                                                     final 
String topologyName) {
         Objects.requireNonNull(keySerializer, "keySerializer can't be null");
         return keyQueryMetadataForKey(storeName,
                                       key,
+                                      headers,
                                       new 
DefaultStreamPartitioner<>(keySerializer),
                                       topologyName);
     }
@@ -239,6 +244,7 @@ public class StreamsMetadataState {
      *
      * @param storeName   Name of the store
      * @param key         Key to use
+     * @param headers     the record headers
      * @param partitioner partitioner to use to find correct partition for key
      * @param <K>         key type
      * @return The {@link KeyQueryMetadata} for the storeName and key or 
{@link KeyQueryMetadata#NOT_AVAILABLE}
@@ -246,9 +252,11 @@ public class StreamsMetadataState {
      */
     public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final 
String storeName,
                                                                     final K 
key,
+                                                                    final 
Headers headers,
                                                                     final 
StreamPartitioner<? super K, ?> partitioner) {
         Objects.requireNonNull(storeName, "storeName can't be null");
         Objects.requireNonNull(key, "key can't be null");
+        Objects.requireNonNull(headers, "headers can't be null");
         Objects.requireNonNull(partitioner, "partitioner can't be null");
         if (topologyMetadata.hasNamedTopologies()) {
             throw new IllegalArgumentException("Cannot invoke the 
keyQueryMetadataForKey(storeName, key, partitioner)"
@@ -273,22 +281,23 @@ public class StreamsMetadataState {
         if (sourceTopicsInfo == null) {
             return null;
         }
-        return keyQueryMetadataForKey(storeName, key, partitioner, 
sourceTopicsInfo);
+        return keyQueryMetadataForKey(storeName, key, headers, partitioner, 
sourceTopicsInfo);
     }
 
     /**
-     * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object, 
StreamPartitioner)}
+     * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object, 
Headers, StreamPartitioner)}
      */
     public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final 
String storeName,
                                                                     final K 
key,
+                                                                    final 
Headers headers,
                                                                     final 
StreamPartitioner<? super K, ?> partitioner,
                                                                     final 
String topologyName) {
         Objects.requireNonNull(storeName, "storeName can't be null");
         Objects.requireNonNull(key, "key can't be null");
+        Objects.requireNonNull(headers, "headers can't be null");
         Objects.requireNonNull(partitioner, "partitioner can't be null");
         Objects.requireNonNull(topologyName, "topologyName can't be null");
 
-
         if (!isInitialized()) {
             return KeyQueryMetadata.NOT_AVAILABLE;
         }
@@ -297,7 +306,7 @@ public class StreamsMetadataState {
         if (sourceTopicsInfo == null) {
             return null;
         }
-        return keyQueryMetadataForKey(storeName, key, partitioner, 
sourceTopicsInfo, topologyName);
+        return keyQueryMetadataForKey(storeName, key, headers, partitioner, 
sourceTopicsInfo, topologyName);
     }
 
     /**
@@ -474,10 +483,11 @@ public class StreamsMetadataState {
 
     private <K> KeyQueryMetadata keyQueryMetadataForKey(final String storeName,
                                                         final K key,
+                                                        final Headers headers,
                                                         final 
StreamPartitioner<? super K, ?> partitioner,
                                                         final SourceTopicsInfo 
sourceTopicsInfo) {
 
-        final Integer partition = 
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
 key, null, sourceTopicsInfo.maxPartitions));
+        final Integer partition = 
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
 key, null, headers, sourceTopicsInfo.maxPartitions));
         final Set<TopicPartition> matchingPartitions = new HashSet<>();
         for (final String sourceTopic : sourceTopicsInfo.sourceTopics) {
             matchingPartitions.add(new TopicPartition(sourceTopic, partition));
@@ -507,11 +517,12 @@ public class StreamsMetadataState {
 
     private <K> KeyQueryMetadata keyQueryMetadataForKey(final String storeName,
                                                         final K key,
+                                                        final Headers headers,
                                                         final 
StreamPartitioner<? super K, ?> partitioner,
                                                         final SourceTopicsInfo 
sourceTopicsInfo,
                                                         final String 
topologyName) {
         Objects.requireNonNull(topologyName, "topology name must not be null");
-        final Integer partition = 
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
 key, null, sourceTopicsInfo.maxPartitions));
+        final Integer partition = 
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
 key, null, headers, sourceTopicsInfo.maxPartitions));
         final Set<TopicPartition> matchingPartitions = new HashSet<>();
         for (final String sourceTopic : sourceTopicsInfo.sourceTopics) {
             matchingPartitions.add(new TopicPartition(sourceTopic, partition));
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
index b0a2a3ff07c..c9feee9e246 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
@@ -20,6 +20,8 @@ import 
org.apache.kafka.clients.admin.DeleteConsumerGroupOffsetsResult;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.GroupIdNotFoundException;
 import org.apache.kafka.common.errors.GroupSubscribedToTopicException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.internals.KafkaFutureImpl;
 import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.utils.internals.LogContext;
@@ -418,7 +420,7 @@ public class KafkaStreamsNamedTopologyWrapper extends 
KafkaStreams {
     }
 
     /**
-     * See {@link KafkaStreams#queryMetadataForKey(String, Object, Serializer)}
+     * See {@link KafkaStreams#queryMetadataForKey(String, Object, Headers, 
Serializer)}
      */
     public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
                                                     final K key,
@@ -426,7 +428,7 @@ public class KafkaStreamsNamedTopologyWrapper extends 
KafkaStreams {
                                                     final String topologyName) 
{
         verifyTopologyStateStore(topologyName, storeName);
         validateIsRunningOrRebalancing();
-        return streamsMetadataState.keyQueryMetadataForKey(storeName, key, 
keySerializer, topologyName);
+        return streamsMetadataState.keyQueryMetadataForKey(storeName, key, new 
RecordHeaders(), keySerializer, topologyName);
     }
 
     /**
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java 
b/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java
index e1d8af5d318..8558aaf0d88 100644
--- a/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java
@@ -25,11 +25,14 @@ import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.internals.KafkaFutureImpl;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.MetricsReporter;
 import org.apache.kafka.common.metrics.Sensor.RecordingLevel;
 import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.common.utils.LogCaptureAppender;
@@ -45,6 +48,7 @@ import 
org.apache.kafka.streams.internals.metrics.ClientMetrics;
 import org.apache.kafka.streams.kstream.Materialized;
 import org.apache.kafka.streams.processor.StandbyUpdateListener;
 import org.apache.kafka.streams.processor.StateRestoreListener;
+import org.apache.kafka.streams.processor.StreamPartitioner;
 import org.apache.kafka.streams.processor.api.Processor;
 import org.apache.kafka.streams.processor.api.ProcessorContext;
 import org.apache.kafka.streams.processor.api.Record;
@@ -1109,13 +1113,60 @@ public class KafkaStreamsTest {
         final AtomicReference<StreamThread.State> state2 = 
prepareStreamThread(streamThreadTwo, 2);
         prepareThreadState(streamThreadOne, state1);
         prepareThreadState(streamThreadTwo, state2);
+        final StreamPartitioner<String, Object> simplePartitioner = new 
SimplePartitioner();
         try (final KafkaStreams streams = new 
KafkaStreams(getBuilderWithSource().build(), props, supplier, time)) {
-            assertThrows(StreamsNotStartedException.class, () -> 
streams.queryMetadataForKey("store", "key", (topic, key, value, numPartitions) 
-> Optional.of(Collections.singleton(0))));
+            assertThrows(StreamsNotStartedException.class, () -> 
streams.queryMetadataForKey("store", "key", simplePartitioner));
             streams.start();
             waitForApplicationState(Collections.singletonList(streams), 
KafkaStreams.State.RUNNING, DEFAULT_DURATION);
             streams.close();
             waitForApplicationState(Collections.singletonList(streams), 
KafkaStreams.State.NOT_RUNNING, DEFAULT_DURATION);
-            assertThrows(IllegalStateException.class, () -> 
streams.queryMetadataForKey("store", "key", (topic, key, value, numPartitions) 
-> Optional.of(Collections.singleton(0))));
+            assertThrows(IllegalStateException.class, () -> 
streams.queryMetadataForKey("store", "key", simplePartitioner));
+        }
+    }
+
+    @Test
+    public void shouldPropagateSerializerAndHeadersToStreamsMetadataState() {
+        prepareStreams();
+        prepareStreamThread(streamThreadOne, 1);
+        prepareStreamThread(streamThreadTwo, 2);
+
+        try (final MockedConstruction<StreamsMetadataState> 
metadataStateMockedConstruction = mockConstruction(StreamsMetadataState.class)) 
{
+            try (final KafkaStreams streams = new 
KafkaStreams(getBuilderWithSource().build(), props, supplier, time)) {
+                streams.start();
+                final StreamsMetadataState mockMetadataState = 
metadataStateMockedConstruction.constructed().get(0);
+
+                final Headers headers = new RecordHeaders();
+                headers.add("key", "value".getBytes());
+                final Serializer<String> serializer = new StringSerializer();
+
+                streams.queryMetadataForKey("store", "key", headers, 
serializer);
+
+                verify(mockMetadataState).keyQueryMetadataForKey("store", 
"key", headers, serializer);
+            }
+        }
+    }
+
+    @Test
+    public void shouldPropagatePartitionerAndHeadersToStreamsMetadataState() {
+        prepareStreams();
+        prepareStreamThread(streamThreadOne, 1);
+        prepareStreamThread(streamThreadTwo, 2);
+
+        try (final MockedConstruction<StreamsMetadataState> 
metadataStateMockedConstruction = mockConstruction(StreamsMetadataState.class)) 
{
+            try (final KafkaStreams streams = new 
KafkaStreams(getBuilderWithSource().build(), props, supplier, time)) {
+                streams.start();
+                final StreamsMetadataState mockMetadataState = 
metadataStateMockedConstruction.constructed().get(0);
+
+                final Headers headers = new RecordHeaders();
+                headers.add("key", "value".getBytes());
+
+                @SuppressWarnings("unchecked")
+                final StreamPartitioner<String, Object> partitioner = 
mock(StreamPartitioner.class);
+
+                streams.queryMetadataForKey("store", "key", headers, 
partitioner);
+
+                verify(mockMetadataState).keyQueryMetadataForKey("store", 
"key", headers, partitioner);
+            }
         }
     }
 
@@ -2056,4 +2107,17 @@ public class KafkaStreamsTest {
             }
         }
     }
+
+    private static class SimplePartitioner implements 
StreamPartitioner<String, Object> {
+        @SuppressWarnings("removal")
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+            throw new AssertionError("Deprecated 4-argument partitions method 
was called instead of 5-argument method containing headers.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
+            return Optional.of(Collections.singleton(0));
+        }
+    }
 }
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
index 7a47632d449..ecb733a2a7d 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.kstream.internals;
 
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.serialization.IntegerDeserializer;
 import org.apache.kafka.common.serialization.IntegerSerializer;
 import org.apache.kafka.common.serialization.Serdes;
@@ -57,6 +58,7 @@ import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -85,8 +87,8 @@ public class KStreamRepartitionTest {
         @SuppressWarnings("unchecked")
         final StreamPartitioner<Integer, String> streamPartitionerMock = 
mock(StreamPartitioner.class);
 
-        when(streamPartitionerMock.partitions(anyString(), eq(0), eq("X0"), 
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
-        when(streamPartitionerMock.partitions(anyString(), eq(1), eq("X1"), 
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
+        when(streamPartitionerMock.partitions(anyString(), eq(0), eq("X0"), 
any(Headers.class), 
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
+        when(streamPartitionerMock.partitions(anyString(), eq(1), eq("X1"), 
any(Headers.class), 
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
 
         final String repartitionOperationName = "test";
         final Repartitioned<Integer, String> repartitioned = Repartitioned
@@ -118,8 +120,8 @@ public class KStreamRepartitionTest {
             assertTrue(testOutputTopic.readRecordsToList().isEmpty());
         }
 
-        verify(streamPartitionerMock).partitions(anyString(), eq(0), eq("X0"), 
anyInt());
-        verify(streamPartitionerMock).partitions(anyString(), eq(1), eq("X1"), 
anyInt());
+        verify(streamPartitionerMock).partitions(anyString(), eq(0), eq("X0"), 
any(Headers.class), anyInt());
+        verify(streamPartitionerMock).partitions(anyString(), eq(1), eq("X1"), 
any(Headers.class), anyInt());
     }
 
     @Test
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
index 1ccff8098c5..3f1ef503115 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.streams.kstream.internals;
 
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.serialization.IntegerDeserializer;
 import org.apache.kafka.common.serialization.IntegerSerializer;
 import org.apache.kafka.common.serialization.Serdes;
@@ -33,25 +35,34 @@ import org.apache.kafka.streams.kstream.Consumed;
 import org.apache.kafka.streams.kstream.KTable;
 import org.apache.kafka.streams.kstream.Materialized;
 import org.apache.kafka.streams.kstream.Produced;
+import org.apache.kafka.streams.kstream.TableJoined;
+import org.apache.kafka.streams.processor.StreamPartitioner;
 import org.apache.kafka.streams.state.KeyValueStore;
+import org.apache.kafka.streams.test.TestRecord;
 import org.apache.kafka.streams.utils.UniqueTopicSerdeScope;
 import org.apache.kafka.test.StreamsTestUtils;
 import org.apache.kafka.test.TestUtils;
 
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.apache.kafka.common.utils.Utils.mkEntry;
 import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.apache.kafka.common.utils.Utils.mkProperties;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class KTableKTableForeignKeyJoinScenarioTest {
 
@@ -394,4 +405,84 @@ public class KTableKTableForeignKeyJoinScenarioTest {
             assertThat(x, is(Collections.singletonMap(1, 
"(999-alpha,(999-alpha,beta))")));
         }
     }
+
+    @Test
+    public void shouldPropagateHeadersToCustomForeignKeyJoinPartitioners() {
+        final StreamsBuilder builder = new StreamsBuilder();
+        final KTable<Integer, String> aTable = builder.table("A");
+        final KTable<Integer, String> bTable = builder.table("B");
+
+        final AtomicBoolean leftPartitionerCalled = new AtomicBoolean(false);
+        final AtomicBoolean rightPartitionerCalled = new AtomicBoolean(false);
+
+        final StreamPartitioner<Integer, Void> leftPartitioner = new 
StreamPartitioner<>() {
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
Integer key, final Void value, final Headers headers, final int numPartitions) {
+                leftPartitionerCalled.set(true);
+                assertNotNull(headers);
+                assertNotNull(headers.lastHeader("test-header"));
+                assertEquals("test-value", new 
String(headers.lastHeader("test-header").value()));
+                return Optional.of(Collections.singleton(0));
+            }
+
+            @SuppressWarnings("removal")
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
Integer key, final Void value, final int numPartitions) {
+                throw new AssertionError("Deprecated 4-argument partitions 
method was called instead of 5-argument method containing headers.");
+            }
+        };
+
+        final StreamPartitioner<Integer, Void> rightPartitioner = new 
StreamPartitioner<>() {
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
Integer key, final Void value, final Headers headers, final int numPartitions) {
+                rightPartitionerCalled.set(true);
+                assertNotNull(headers);
+                assertNotNull(headers.lastHeader("test-header"));
+                assertEquals("test-value", new 
String(headers.lastHeader("test-header").value()));
+                return Optional.of(Collections.singleton(0));
+            }
+
+            @SuppressWarnings("removal")
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
Integer key, final Void value, final int numPartitions) {
+                throw new AssertionError("Deprecated 4-argument partitions 
method was called instead of 5-argument method containing headers.");
+            }
+        };
+
+        final KTable<Integer, String> fkJoinResult = aTable.join(
+            bTable,
+            value -> Integer.parseInt(value.split("-")[0]),
+            (aVal, bVal) -> "(" + aVal + "," + bVal + ")",
+            TableJoined.with(leftPartitioner, rightPartitioner),
+            Materialized.as("asdf")
+        );
+
+        fkJoinResult.toStream().to("output");
+
+        final Properties config = new Properties();
+        config.setProperty(StreamsConfig.APPLICATION_ID_CONFIG, "test-app");
+        config.setProperty(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, 
"dummy:1234");
+        config.setProperty(StreamsConfig.STATE_DIR_CONFIG, 
TestUtils.tempDirectory().getAbsolutePath());
+        config.setProperty(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, 
Serdes.Integer().getClass().getName());
+        config.setProperty(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, 
Serdes.String().getClass().getName());
+
+        try (final TopologyTestDriver driver = new 
TopologyTestDriver(builder.build(), config)) {
+            final TestInputTopic<Integer, String> aTopic = 
driver.createInputTopic("A", new IntegerSerializer(), new StringSerializer());
+            final TestInputTopic<Integer, String> bTopic = 
driver.createInputTopic("B", new IntegerSerializer(), new StringSerializer());
+
+            final Headers headers = new RecordHeaders();
+            headers.add("test-header", "test-value".getBytes());
+
+            // Pipe input to A with headers. This triggers 
SubscriptionSinkPartitioner,
+            // which delegates to rightPartitioner.
+            aTopic.pipeInput(new TestRecord<>(1, "999-alpha", headers));
+
+            // Pipe input to B with headers. This triggers 
ForeignResponseSinkPartitioner,
+            // which delegates to leftPartitioner.
+            bTopic.pipeInput(new TestRecord<>(999, "beta", headers));
+        }
+
+        assertTrue(leftPartitionerCalled.get());
+        assertTrue(rightPartitionerCalled.get());
+    }
 }
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
index 25a15efdd4c..c8d27458290 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
@@ -20,8 +20,9 @@ import 
org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.serialization.IntegerSerializer;
-import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.streams.kstream.TimeWindowedSerializer;
 import org.apache.kafka.streams.kstream.Windowed;
 
@@ -35,14 +36,17 @@ import java.util.Random;
 import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class WindowedStreamPartitionerTest {
 
     private final String topicName = "topic";
 
     private final IntegerSerializer intSerializer = new IntegerSerializer();
-    private final StringSerializer stringSerializer = new StringSerializer();
 
     private final List<PartitionInfo> infos = Arrays.asList(
             new PartitionInfo(topicName, 0, Node.noNode(), new Node[0], new 
Node[0]),
@@ -57,28 +61,37 @@ public class WindowedStreamPartitionerTest {
             Collections.emptySet(), Collections.emptySet());
 
     @Test
-    public void testCopartitioning() {
-        final Random rand = new Random();
+    public void shouldThrowUnsupportedOperationExceptionForDeprecatedMethod() {
         final WindowedSerializer<Integer> timeWindowedSerializer = new 
TimeWindowedSerializer<>(intSerializer);
         final WindowedStreamPartitioner<Integer, String> streamPartitioner = 
new WindowedStreamPartitioner<>(timeWindowedSerializer);
+        final Windowed<Integer> windowedKey = new Windowed<>(1, new 
TimeWindow(10, 20));
+
+        assertThrows(UnsupportedOperationException.class, () -> 
streamPartitioner.partitions(topicName, windowedKey, "value", infos.size()));
+    }
 
-        for (int k = 0; k < 10; k++) {
-            final Integer key = rand.nextInt();
-            final byte[] keyBytes = intSerializer.serialize(topicName, key);
+    @Test
+    public void testCopartitioningWithHeaders() {
+        final Random rand = new Random();
+        final Headers headers = new RecordHeaders();
+        headers.add("key", "value".getBytes());
 
-            final String value = key.toString();
+        @SuppressWarnings("unchecked")
+        final WindowedSerializer<Integer> mockSerializer = 
mock(WindowedSerializer.class);
+        final WindowedStreamPartitioner<Integer, String> streamPartitioner = 
new WindowedStreamPartitioner<>(mockSerializer);
 
-            final Set<Integer> expected = 
Set.of(BuiltInPartitioner.partitionForKey(keyBytes, 
cluster.partitionsForTopic(topicName).size()));
+        final Integer key = rand.nextInt();
+        final String value = key.toString();
+        final TimeWindow window = new TimeWindow(10, 20);
+        final Windowed<Integer> windowedKey = new Windowed<>(key, window);
+        final byte[] keyBytes = intSerializer.serialize(topicName, key);
+        final Set<Integer> expected = 
Set.of(BuiltInPartitioner.partitionForKey(keyBytes, 
cluster.partitionsForTopic(topicName).size()));
 
-            for (int w = 1; w < 10; w++) {
-                final TimeWindow window = new TimeWindow(10 * w, 20 * w);
+        when(mockSerializer.serializeBaseKey(topicName, headers, 
windowedKey)).thenReturn(keyBytes);
 
-                final Windowed<Integer> windowedKey = new Windowed<>(key, 
window);
-                final Optional<Set<Integer>> actual = 
streamPartitioner.partitions(topicName, windowedKey, value, infos.size());
+        final Optional<Set<Integer>> actual = 
streamPartitioner.partitions(topicName, windowedKey, value, headers, 
infos.size());
 
-                assertTrue(actual.isPresent());
-                assertEquals(expected, actual.get());
-            }
-        }
+        assertTrue(actual.isPresent());
+        assertEquals(expected, actual.get());
+        verify(mockSerializer).serializeBaseKey(topicName, headers, 
windowedKey);
     }
 }
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitionerTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitionerTest.java
new file mode 100644
index 00000000000..57864c656cc
--- /dev/null
+++ 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitionerTest.java
@@ -0,0 +1,86 @@
+/*
+ * 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.kafka.streams.processor.internals;
+
+import org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.serialization.Serializer;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.Optional;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class DefaultStreamPartitionerTest {
+
+    private static final String TOPIC = "topic";
+    private static final String KEY = "key";
+    private static final String VALUE = "value";
+    private static final int NUM_PARTITIONS = 5;
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void shouldPropagateHeadersToSerializer() {
+        final Serializer<String> keySerializer = mock(Serializer.class);
+        final DefaultStreamPartitioner<String, String> 
defaultStreamPartitioner = new DefaultStreamPartitioner<>(keySerializer);
+        final Headers headers = new RecordHeaders();
+        headers.add("key", "value".getBytes());
+        final byte[] expectedKeyBytes = "serializedKey".getBytes();
+
+        when(keySerializer.serialize(TOPIC, headers, 
KEY)).thenReturn(expectedKeyBytes);
+
+        final Optional<Set<Integer>> partition = 
defaultStreamPartitioner.partitions(TOPIC, KEY, VALUE, headers, NUM_PARTITIONS);
+
+        verify(keySerializer).serialize(TOPIC, headers, KEY);
+        assertTrue(partition.isPresent());
+        
assertEquals(Collections.singleton(BuiltInPartitioner.partitionForKey(expectedKeyBytes,
 NUM_PARTITIONS)), partition.get());
+    }
+
+    @SuppressWarnings({"unchecked"})
+    @Test
+    public void shouldThrowUnsupportedOperationExceptionForDeprecatedMethod() {
+        final Serializer<String> keySerializer = mock(Serializer.class);
+        final DefaultStreamPartitioner<String, String> 
defaultStreamPartitioner = new DefaultStreamPartitioner<>(keySerializer);
+
+        assertThrows(UnsupportedOperationException.class, () -> 
defaultStreamPartitioner.partitions(TOPIC, KEY, VALUE, NUM_PARTITIONS));
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void shouldReturnEmptyOptionalWhenSerializedKeyIsNull() {
+        final Serializer<String> keySerializer = mock(Serializer.class);
+        final DefaultStreamPartitioner<String, String> 
defaultStreamPartitioner = new DefaultStreamPartitioner<>(keySerializer);
+        final Headers headers = new RecordHeaders();
+
+        when(keySerializer.serialize(TOPIC, headers, KEY)).thenReturn(null);
+
+        final Optional<Set<Integer>> partition = 
defaultStreamPartitioner.partitions(TOPIC, KEY, VALUE, headers, NUM_PARTITIONS);
+
+        verify(keySerializer).serialize(TOPIC, headers, KEY);
+        assertFalse(partition.isPresent());
+    }
+}
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
index f39d818f2b8..55e8dddc2d3 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
@@ -1103,8 +1103,14 @@ public class ProcessorTopologyTest {
     }
 
     static class DroppingPartitioner implements StreamPartitioner<String, 
String> {
+        @SuppressWarnings("removal")
         @Override
         public Optional<Set<Integer>> partitions(final String topic, final 
String key, final String value, final int numPartitions) {
+            throw new AssertionError("Deprecated 4-argument partitions method 
was called instead of 5-argument method containing headers.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final String value, final Headers headers, final int numPartitions) 
{
             final Set<Integer> partitions = new HashSet<>();
             for (int i = 1; i < numPartitions; i += 2) {
                 partitions.add(i);
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
index 61c125058d5..b2cf4acdec2 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
@@ -109,6 +109,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
@@ -137,8 +138,18 @@ public class RecordCollectorTest {
     private final StringSerializer stringSerializer = new StringSerializer();
     private final ByteArraySerializer byteArraySerializer = new 
ByteArraySerializer();
 
-    private final StreamPartitioner<String, Object> streamPartitioner =
-        (topic, key, value, numPartitions) -> 
Optional.of(Collections.singleton(Integer.parseInt(key) % numPartitions));
+    private final StreamPartitioner<String, Object> streamPartitioner = new 
StreamPartitioner<String, Object>() {
+        @SuppressWarnings("removal")
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+            throw new AssertionError("Deprecated 4-argument partitions method 
was called instead of 5-argument method containing headers.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
+            return Optional.of(Collections.singleton(Integer.parseInt(key) % 
numPartitions));
+        }
+    };
 
     private final MockProducer<byte[], byte[]> mockProducer
         = new MockProducer<>(cluster, true, new 
org.apache.kafka.clients.producer.RoundRobinPartitioner(), new 
ByteArraySerializer(), new ByteArraySerializer());
@@ -304,8 +315,14 @@ public class RecordCollectorTest {
     @Test
     public void shouldSendOnlyToEvenPartitions() {
         class EvenPartitioner implements StreamPartitioner<String, Object> {
+            @SuppressWarnings("removal")
             @Override
             public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+                throw new AssertionError("Deprecated 4-argument partitions 
method was called instead of 5-argument method containing headers.");
+            }
+
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
                 final Set<Integer> partitions = new HashSet<>();
                 for (int i = 0; i < numPartitions; i += 2) {
                     partitions.add(i);
@@ -369,8 +386,14 @@ public class RecordCollectorTest {
     public void shouldBroadcastToAllPartitions() {
 
         class BroadcastingPartitioner implements StreamPartitioner<String, 
Object> {
+            @SuppressWarnings("removal")
             @Override
             public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+                throw new AssertionError("Deprecated 4-argument partitions 
method was called instead of 5-argument method containing headers.");
+            }
+
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
                 return Optional.of(IntStream.range(0, 
numPartitions).boxed().collect(Collectors.toSet()));
             }
         }
@@ -430,8 +453,14 @@ public class RecordCollectorTest {
     public void shouldDropAllRecords() {
 
         class DroppingPartitioner implements StreamPartitioner<String, Object> 
{
+            @SuppressWarnings("removal")
             @Override
             public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+                throw new AssertionError("Deprecated 4-argument partitions 
method was called instead of 5-argument method containing headers.");
+            }
+
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
                 return Optional.of(Collections.emptySet());
             }
         }
@@ -503,8 +532,14 @@ public class RecordCollectorTest {
     public void shouldUseDefaultPartitionerViaPartitions() {
 
         class DefaultPartitioner implements StreamPartitioner<String, Object> {
+            @SuppressWarnings("removal")
             @Override
             public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+                throw new AssertionError("Deprecated 4-argument partitions 
method was called instead of 5-argument method containing headers.");
+            }
+
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
                 return Optional.empty();
             }
         }
@@ -562,8 +597,18 @@ public class RecordCollectorTest {
     @Test
     public void shouldUseDefaultPartitionerAsPartitionReturnsEmptyOptional() {
 
-        final StreamPartitioner<String, Object> streamPartitioner =
-                (topic, key, value, numPartitions) -> Optional.empty();
+        final StreamPartitioner<String, Object> streamPartitioner = new 
StreamPartitioner<String, Object>() {
+            @SuppressWarnings("removal")
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+                throw new AssertionError("Deprecated 4-argument partitions 
method was called instead of 5-argument method containing headers.");
+            }
+
+            @Override
+            public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
+                return Optional.empty();
+            }
+        };
 
         final SinkNode<?, ?> sinkNode = new SinkNode<>(
                 sinkNodeName,
@@ -613,6 +658,18 @@ public class RecordCollectorTest {
         assertEquals(9, mockProducer.history().size());
     }
 
+    @Test
+    public void shouldPropagateHeadersToPartitioner() {
+        final Headers headers = new RecordHeaders(new Header[] {new 
RecordHeader("key", "value".getBytes())});
+        @SuppressWarnings("unchecked")
+        final StreamPartitioner<String, Object> mockPartitioner = 
mock(StreamPartitioner.class);
+        when(mockPartitioner.partitions(topic, "3", "0", headers, 
3)).thenReturn(Optional.of(Collections.singleton(0)));
+
+        collector.send(topic, "3", "0", headers, null, stringSerializer, 
stringSerializer, null, context, mockPartitioner);
+
+        verify(mockPartitioner).partitions(topic, "3", "0", headers, 3);
+    }
+
     @Test
     public void shouldUseDefaultPartitionerAsStreamPartitionerIsNull() {
 
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
index 4ea73fcbf03..ae30308380d 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
@@ -18,6 +18,8 @@ package org.apache.kafka.streams.processor.internals;
 
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.serialization.Serdes;
 import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.utils.internals.LogContext;
@@ -53,6 +55,11 @@ import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class StreamsMetadataStateTest {
 
@@ -129,20 +136,10 @@ public class StreamsMetadataStateTest {
         topologyMetadata.buildAndRewriteTopology();
         metadataState = new StreamsMetadataState(topologyMetadata, hostOne, 
logContext);
         metadataState.onChange(hostToActivePartitions, 
hostToStandbyPartitions, partitionInfos);
-        partitioner = (topic, key, value, numPartitions) -> 
Optional.of(Collections.singleton(1));
+        partitioner = new FixedPartitionPartitioner(1);
         storeNames = Set.of("table-one", "table-two", "merged-table", 
globalTable);
     }
 
-    static class MultiValuedPartitioner implements StreamPartitioner<String, 
Object> {
-        @Override
-        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
-            final Set<Integer> partitions = new HashSet<>();
-            partitions.add(0);
-            partitions.add(1);
-            return Optional.of(partitions);
-        }
-    }
-
     @Test
     public void shouldNotThrowExceptionWhenOnChangeNotCalled() {
         final Collection<StreamsMetadata> metadata = new StreamsMetadataState(
@@ -237,9 +234,13 @@ public class StreamsMetadataStateTest {
             Collections.singletonMap(tp4, new PartitionInfo("topic-three", 1, 
null, null, null)));
 
         final KeyQueryMetadata expected = new KeyQueryMetadata(hostThree, 
Set.of(hostTwo), 0);
-        final KeyQueryMetadata actual = 
metadataState.keyQueryMetadataForKey("table-three",
-                                                                    "the-key",
-                                                                    
Serdes.String().serializer());
+        final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+                "table-three",
+                "the-key",
+                new RecordHeaders(),
+                Serdes.String().serializer()
+        );
+
         assertEquals(expected, actual);
     }
 
@@ -253,9 +254,12 @@ public class StreamsMetadataStateTest {
 
         final KeyQueryMetadata expected = new KeyQueryMetadata(hostTwo, 
Collections.emptySet(), 1);
 
-        final KeyQueryMetadata actual = 
metadataState.keyQueryMetadataForKey("table-three",
+        final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+                "table-three",
                 "the-key",
-                partitioner);
+                new RecordHeaders(),
+                partitioner
+        );
         assertEquals(expected, actual);
         assertEquals(1, actual.partition());
     }
@@ -269,15 +273,23 @@ public class StreamsMetadataStateTest {
                 Collections.singletonMap(tp4, new PartitionInfo("topic-three", 
1, null, null, null)));
 
 
-        assertThrows(IllegalArgumentException.class, () -> 
metadataState.keyQueryMetadataForKey("table-three",
+        assertThrows(IllegalArgumentException.class, () -> 
metadataState.keyQueryMetadataForKey(
+                "table-three",
                 "the-key",
-                new MultiValuedPartitioner()));
+                new RecordHeaders(),
+                new MultiValuedPartitioner()
+        ));
     }
 
     @Test
     public void shouldReturnNotAvailableWhenClusterIsEmpty() {
         metadataState.onChange(Collections.emptyMap(), Collections.emptyMap(), 
Collections.emptyMap());
-        final KeyQueryMetadata result = 
metadataState.keyQueryMetadataForKey("table-one", "a", 
Serdes.String().serializer());
+        final KeyQueryMetadata result = metadataState.keyQueryMetadataForKey(
+                "table-one",
+                "a",
+                new RecordHeaders(),
+                Serdes.String().serializer()
+        );
         assertEquals(KeyQueryMetadata.NOT_AVAILABLE, result);
     }
 
@@ -291,39 +303,76 @@ public class StreamsMetadataStateTest {
 
         final KeyQueryMetadata expected = new KeyQueryMetadata(hostTwo, 
Set.of(hostOne), 2);
 
-        final KeyQueryMetadata actual = 
metadataState.keyQueryMetadataForKey("merged-table",  "the-key",
-            (topic, key, value, numPartitions) -> 
Optional.of(Collections.singleton(2)));
+        final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+                "merged-table",
+                "the-key",
+                new RecordHeaders(),
+                new FixedPartitionPartitioner(2)
+        );
 
         assertEquals(expected, actual);
     }
 
     @Test
     public void shouldReturnNullOnGetWithKeyWhenStoreDoesntExist() {
-        final KeyQueryMetadata actual = 
metadataState.keyQueryMetadataForKey("not-a-store",
+        final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+                "not-a-store",
                 "key",
-                Serdes.String().serializer());
+                new RecordHeaders(),
+                Serdes.String().serializer()
+        );
         assertNull(actual);
     }
 
     @Test
     public void shouldThrowWhenKeyIsNull() {
-        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey("table-three", null, 
Serdes.String().serializer()));
+        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey(
+                "table-three",
+                null,
+                new RecordHeaders(),
+                Serdes.String().serializer()
+        ));
     }
 
     @Test
     public void shouldThrowWhenSerializerIsNull() {
-        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey("table-three", "key", (Serializer<Object>) 
null));
+        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey(
+                "table-three",
+                "key",
+                new RecordHeaders(),
+                (Serializer<Object>) null
+        ));
     }
 
     @Test
     public void shouldThrowIfStoreNameIsNull() {
-        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey(null, "key", 
Serdes.String().serializer()));
+        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey(
+                null,
+                "key",
+                new RecordHeaders(),
+                Serdes.String().serializer()
+        ));
     }
 
     @SuppressWarnings("unchecked")
     @Test
     public void shouldThrowIfStreamPartitionerIsNull() {
-        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey(null, "key", (StreamPartitioner) null));
+        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey(
+                null,
+                "key",
+                new RecordHeaders(),
+                (StreamPartitioner) null
+        ));
+    }
+
+    @Test
+    public void shouldThrowIfHeadersIsNull() {
+        assertThrows(NullPointerException.class, () -> 
metadataState.keyQueryMetadataForKey(
+                "table-three",
+                "key",
+                null,
+                Serdes.String().serializer()
+        ));
     }
 
     @Test
@@ -347,7 +396,12 @@ public class StreamsMetadataStateTest {
 
     @Test
     public void shouldGetQueryMetadataForGlobalStoreWithKey() {
-        final KeyQueryMetadata metadata = 
metadataState.keyQueryMetadataForKey(globalTable, "key", 
Serdes.String().serializer());
+        final KeyQueryMetadata metadata = metadataState.keyQueryMetadataForKey(
+                globalTable,
+                "key",
+                new RecordHeaders(),
+                Serdes.String().serializer()
+        );
         assertEquals(hostOne, metadata.activeHost());
         assertTrue(metadata.standbyHosts().isEmpty());
     }
@@ -360,12 +414,22 @@ public class StreamsMetadataStateTest {
             logContext
         );
         streamsMetadataState.onChange(hostToActivePartitions, 
hostToStandbyPartitions, partitionInfos);
-        assertNotNull(streamsMetadataState.keyQueryMetadataForKey(globalTable, 
"key", Serdes.String().serializer()));
+        assertNotNull(streamsMetadataState.keyQueryMetadataForKey(
+                globalTable,
+                "key",
+                new RecordHeaders(),
+                Serdes.String().serializer()
+        ));
     }
 
     @Test
     public void shouldGetQueryMetadataForGlobalStoreWithKeyAndPartitioner() {
-        final KeyQueryMetadata metadata = 
metadataState.keyQueryMetadataForKey(globalTable, "key", partitioner);
+        final KeyQueryMetadata metadata = metadataState.keyQueryMetadataForKey(
+                globalTable,
+                "key",
+                new RecordHeaders(),
+                partitioner
+        );
         assertEquals(hostOne, metadata.activeHost());
         assertTrue(metadata.standbyHosts().isEmpty());
     }
@@ -378,7 +442,43 @@ public class StreamsMetadataStateTest {
             logContext
         );
         streamsMetadataState.onChange(hostToActivePartitions, 
hostToStandbyPartitions, partitionInfos);
-        assertNotNull(streamsMetadataState.keyQueryMetadataForKey(globalTable, 
"key", partitioner));
+        assertNotNull(streamsMetadataState.keyQueryMetadataForKey(
+                globalTable,
+                "key",
+                new RecordHeaders(),
+                partitioner)
+        );
+    }
+
+    @Test
+    public void shouldPropagateHeadersToPartitioner() {
+        final Headers headers = new RecordHeaders();
+        headers.add("key", "value".getBytes());
+
+        @SuppressWarnings("unchecked")
+        final StreamPartitioner<String, Object> mockPartitioner = 
mock(StreamPartitioner.class);
+        when(mockPartitioner.partitions(
+                eq("topic-three"),
+                eq("the-key"),
+                eq(null),
+                eq(headers),
+                anyInt()
+        )).thenReturn(Optional.of(Collections.singleton(0)));
+
+        metadataState.keyQueryMetadataForKey(
+                "table-three",
+                "the-key",
+                headers,
+                mockPartitioner
+        );
+
+        verify(mockPartitioner).partitions(
+                "topic-three",
+                "the-key",
+                null,
+                headers,
+                1
+        );
     }
 
     @Test
@@ -404,4 +504,39 @@ public class StreamsMetadataStateTest {
 
         assertFalse(metadataState.allMetadata().isEmpty(), "encapsulation 
broken");
     }
+
+    private static class MultiValuedPartitioner implements 
StreamPartitioner<String, Object> {
+        @SuppressWarnings("removal")
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+            throw new AssertionError("Deprecated 4-argument partitions method 
was called instead of 5-argument method containing headers.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
+            final Set<Integer> partitions = new HashSet<>();
+            partitions.add(0);
+            partitions.add(1);
+            return Optional.of(partitions);
+        }
+    }
+
+    private static class FixedPartitionPartitioner implements 
StreamPartitioner<String, Object> {
+        private final int partition;
+
+        FixedPartitionPartitioner(final int partition) {
+            this.partition = partition;
+        }
+
+        @SuppressWarnings("removal")
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final int numPartitions) {
+            throw new AssertionError("Deprecated 4-argument partitions method 
was called instead of 5-argument method containing headers.");
+        }
+
+        @Override
+        public Optional<Set<Integer>> partitions(final String topic, final 
String key, final Object value, final Headers headers, final int numPartitions) 
{
+            return Optional.of(Collections.singleton(partition));
+        }
+    }
 }

Reply via email to