Writing a Custom KafkaJS Partition Assigner for Local Joins

Co-partitioning gets you halfway there. If market-updates and notifications have the same partition count and are both keyed by symbol, BBCA lands on the same partition number in both topics. The partitioner is deterministic, so that much is guaranteed. Getting partition 3 of both topics onto the same consumer, the same pod, inside the same process, is a separate problem, and KafkaJS does not solve it for you.

I left the co-partitioning post sitting on exactly that gap, then spent an evening closing it, because my little ticker system wanted a join and I did not want that join crossing the network.

Co-partitioning is data layout, not placement

Quick recap, since everything below builds on it: two topics are co-partitioned when they share the same partition count and the same key with the same partitioner. Kafka's default partitioner is murmur2(key) % numPartitions, which is deterministic, so a given symbol always maps to the same partition number regardless of which topic you are looking at. Both of my topics have 6 partitions and both are keyed by symbol, so partition 3 of one lines up with partition 3 of the other.

That is the data guarantee. It says nothing about which consumer process reads that data. Co-partitioning tells you where the data lives. It does not tell you where the consumer runs. The second half is co-localization, and it depends entirely on how your consumer group's partition assignment works.

Why KafkaJS's round-robin assigner breaks this

Kafka's Java client ships a RangeAssignor that co-localizes same-numbered partitions across the topics a consumer subscribes to. On the JVM, join-friendly assignment comes almost for free.

KafkaJS has no such assignor. It ships one built-in strategy: round-robin. Round-robin takes every topic-partition across every subscribed topic, lays them all out in one flat list, and deals them out to members in turn. With some combinations of pod count and partition count, that happens to co-locate matching partition numbers by coincidence. With others, it does not. My setup, 4 pods over 6 partitions per topic, is one of the ones where it does not: the notification owners end up shifted off the market-data owners, and co-location breaks even though the data underneath is still perfectly co-partitioned.

I lost more time than I want to admit assuming the JVM behavior carried across. It does not. The fix has to happen at the assignment layer, in TypeScript, since no library assigner exists to reach for.

The group protocol assigner interface

KafkaJS exposes a PartitionAssigner factory. It gets called once per consumer, gets access to the cluster metadata, and has to return an object with three things: a name, a version, and an assign function. The name and version are what get negotiated during the consumer group's join protocol, so all members in the group need to agree on both. The assign function is where the actual assignment logic lives, and it only runs on whichever member the broker elects as group leader; that member computes the assignment for everyone and the broker distributes the result.

import { AssignerProtocol, type PartitionAssigner } from "kafkajs";

const NAME = "CoPartitionAssigner";
const VERSION = 1;

export const coPartitionAssigner: PartitionAssigner = ({ cluster }) => ({
  name: NAME,
  version: VERSION,

  async assign({ members, topics }) {
    // implementation below
  },

  protocol({ topics }) {
    return {
      name: NAME,
      metadata: AssignerProtocol.MemberMetadata.encode({
        version: VERSION,
        topics,
        userData: Buffer.alloc(0),
      }),
    };
  },
});

protocol() is what each member advertises when it joins the group, telling the broker which assigner it supports and which topics it subscribed to. assign() computes who gets what, and it has to return a memberAssignment encoded through AssignerProtocol.MemberAssignment.encode for every member, including itself.

Grouping by partition number, not by flat list

The round-robin assigner's mistake is treating "partition 3 of market-updates" and "partition 3 of notifications" as two unrelated items in the same bag. My fix is to stop flattening. Iterate by partition index instead, and for each index, hand every subscribed topic's copy of that index to the same member:

async assign({ members, topics }) {
  // Deterministic member order so every member computes the same assignment.
  const sortedMembers = members.map((m) => m.memberId).sort();
  const memberCount = sortedMembers.length;

  const partitionsByTopic: Record<string, number[]> = {};
  let maxPartitions = 0;
  for (const topic of topics) {
    const ids = cluster
      .findTopicPartitionMetadata(topic)
      .map((p) => p.partitionId)
      .sort((a, b) => a - b);
    partitionsByTopic[topic] = ids;
    if (ids.length > maxPartitions) maxPartitions = ids.length;
  }

  const assignment: Record<string, Record<string, number[]>> = {};
  for (const m of sortedMembers) assignment[m] = {};

  // The crux: every topic's partition `p` is owned by the same member.
  for (let p = 0; p < maxPartitions; p++) {
    const owner = sortedMembers[p % memberCount];
    for (const topic of topics) {
      if (partitionsByTopic[topic].includes(p)) {
        (assignment[owner][topic] ??= []).push(p);
      }
    }
  }

  return sortedMembers.map((memberId) => ({
    memberId,
    memberAssignment: AssignerProtocol.MemberAssignment.encode({
      version: VERSION,
      assignment: assignment[memberId],
      userData: Buffer.alloc(0),
    }),
  }));
}

A few details worth calling out:

• Sorting memberId first. Only the elected leader runs assign, so members are not each computing this independently. The leader still wants a deterministic, reproducible mapping from partition index to member: sorting gives the same output for the same input on every run, which keeps ownership stable across rebalances instead of shuffling for no reason.
• The loop runs over partition index, not over topic-partition pairs. p % memberCount decides the owner once per index, and that owner then claims partition p on every topic that has it. The entire idea fits in those two lines.
• Equal partition counts across subscribed topics are required. If market-updates had 6 partitions and notifications had 8, indices 6 and 7 would only ever appear on one topic, and the co-partitioning guarantee the whole scheme depends on would already be broken upstream. The assigner cannot fix a topic mismatch, only preserve whatever co-partitioning already exists.

What all of this reproduces by hand is what Kafka Streams does automatically under the hood when it demands co-partitioned inputs for a join. Without that library, you write the behavior yourself.

graph TD
    subgraph "Round-robin (flat list, 4 pods)"
        RRList["market-p0, market-p1, ..., notif-p0, notif-p1, ..."]
        RRList --> RP1["pod-1: market-p0, market-p4, notif-p1, notif-p5"]
        RRList --> RP2["pod-2: market-p1, market-p5, notif-p2"]
        RRList --> RP3["pod-3: market-p2, notif-p0, notif-p3"]
        RRList --> RP4["pod-4: market-p3, notif-p4"]
    end
    subgraph "Co-partition assigner (by index, 4 pods)"
        CP1["pod-1 owns index 0: market-p0 AND notif-p0"]
        CP2["pod-2 owns index 1: market-p1 AND notif-p1"]
        CP3["pod-3 owns index 2: market-p2 AND notif-p2"]
        CP4["pod-4 owns index 3: market-p3 AND notif-p3"]
    end

The round-robin side of that diagram illustrates the failure mode, a flat deal misaligning topic-partition pairs across pods, rather than tracing one specific run. The co-partition side is the guarantee itself: whichever pod owns index p, it owns p on both topics, with no exceptions.

Measuring whether the assignment co-locates

Claiming an assigner works is easy, so I measured it. Same 4-pod consumer group, same 6-partition topics, run twice: once with partitionAssigners: [coPartitionAssigner] set on every consumer, once with the KafkaJS default. Each pod listened for GROUP_JOIN events and recorded which partitions it ended up owning on market-updates and on notifications. Because the partitioner is deterministic, I already knew which partition every symbol's notifications would land on, so I could tally whether a notification's partition matched a market-updates partition the same pod already owned, a local join, or not, a remote one.

With my assigner, every pod's market partitions and notification partitions came out as identical sets. All 30 notifications produced (10 symbols times 3 notifications each) joined locally: 30/30 local, 0 remote. With the default round-robin assigner over the same 4 pods and the same 6 partitions, the notification owners got shifted off the market-data owners entirely: 0 local, 30 remote. Same data, same keys, same co-partitioning, and the only variable that changed was which assigner ran the group protocol.

Co-partitioning alone was never going to be enough. Data lining up on paper does not help if the runtime scatters it across different processes.

Spending that assignment on a real join

An aligned assignment is worth something only once you use it, so the second run does a partition-local join.

Each pod in the group subscribes to both topics and keeps a Map<StockSymbol, number> of the latest price it has seen, built entirely from the market-updates messages it consumes:

consumer.run({
  eachMessage: async ({ topic, message }) => {
    if (topic === MARKET) {
      const tick = tickFromKafka(message.value);
      pod.localPrice.set(tick.symbol, tick.price); // local state for owned symbols
    } else {
      const n = notifFromKafka(message.value);
      if (pod.localPrice.has(n.symbol)) pod.localJoins++;
      else pod.remoteNeeded++;
    }
  },
});

When a notification for a symbol arrives, the pod checks its own local map. With my assigner active, the check always succeeds: the pod that owns a symbol's notifications partition is, by construction, the same pod that owns its market-updates partition, so the price is already sitting in memory. No cache lookup, no RPC to another pod, no shared state store, just a map read. With round-robin, notifications for a given symbol can land on a pod that never subscribed to that symbol's slice of market-updates, so localPrice.has(n.symbol) comes back false. My run counts that as a miss; in a production system it is the moment you have to reach across the network for state you need right now.

Detecting a price alert crossing needs the latest price at the instant a tick arrives. If that lookup stays a local map read instead of a network hop, the whole enrichment path gets simpler and faster, with no shared cache to keep consistent across pods. Co-partitioning plus co-localization turns "join two Kafka topics" from a distributed systems problem into an in-memory one.

The Redis analog, briefly

Redis Cluster does the same kind of thing with hash tags. Keys hash to one of 16384 slots via CRC16, and slots map to nodes. Wrapping a common substring in braces, market:{BBCA} and notif:{BBCA}, forces only that substring to be hashed, so both keys land on the same slot and therefore the same node. Without the braces, market:BBCA and notif:BBCA usually end up on different slots. Same underlying idea as Kafka's partition-index trick: force related keys onto the same physical location so a cross-stream operation never has to leave one machine. Kafka does it by partition number; Redis Cluster does it by hash slot.

The short version

• Co-partitioning (same partition count, same key, same partitioner) guarantees data layout lines up. It says nothing about which consumer process reads that data.
• KafkaJS only ships a round-robin assigner, which can break co-location even when the underlying topics are correctly co-partitioned. My 4-pod, 6-partition case is one where it does.
• A custom PartitionAssigner needs name, version, and an assign function that returns an encoded memberAssignment per member. Only the elected group leader runs assign.
• The core move is grouping by partition index across topics instead of flattening all topic-partitions into one list: partition p of every subscribed topic goes to the same member.
• Measured over the same 4 pods: 30/30 notifications join locally with my assigner, 0/30 with round-robin.
• That alignment pays for itself in the join, where notifications get enriched from an in-memory price map that stays populated for the right symbols only because of co-location.

Everything so far has run against a single broker that never went down. Next up: three brokers, a leader I stop on purpose, and Redis Cluster's very different reason for wanting three nodes.