Writing a Custom KafkaJS Partition Assigner for Local Joins

With the same partition count, encoded symbol keys, and partitioner, market-updates and notifications map BBCA to matching partition numbers. The consumer assigner determines which process owns those partitions. Matching numbers alone do not place both partitions on the same consumer.

The co partitioning post left the consumer assignment unresolved. I spent an evening implementing an assigner so my ticker could join prices and notifications within each process.

Co partitioning is data layout, not placement

Both topics have 6 partitions and use the same symbol key. The default mapping applies (murmur2(key) & 0x7fffffff) % numPartitions. Thus, matching partition numbers contain the same symbols. The consumer assignment must preserve that alignment across processes.

That is the data guarantee. Co partitioning tells you where the data lives, while consumer assignment determines which process reads it. 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.

The KafkaJS version used here includes a round robin assigner. It distributes a flat list of topic partitions across members. This can align matching partition numbers for some combinations of member and partition counts.

My setup has four pods and six partitions per topic. Its round robin assignments put some matching partitions on different pods. The records still share partition numbers, but a local join needs both topic partitions in the same process.

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 with access to cluster metadata. The returned object provides name, version, protocol(), and assign(). Members advertise supported protocols when they join the group. The elected group leader runs assign() and returns the assignment for broker distribution.

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

Round robin treats each topic and partition pair as a separate assignment item. My implementation instead chooses one owner per partition index. It assigns that index from every subscribed topic 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 explain the implementation:

• Member order: sort memberId values so identical member lists produce identical assignments. Membership changes can still move many partitions. This is not sticky assignment.
• Partition index: compute the owner with p % memberCount. That member receives partition p from each topic.
• Matching inputs: all members must subscribe to the same relevant topics. The topics must have equal partition counts and matching key mapping. With 6 partitions in market-updates and 8 in notifications, indices 6 and 7 have no match. The assigner cannot correct that data layout.

What all of this reproduces by hand is what Kafka Streams does automatically internally 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 diagram illustrates how matching partition numbers can reach different pods. It does not trace a specific run. In the custom assignment, the member that owns index p receives partition p from both topics.

Measuring whether the assignment co locates

I tested both strategies with 4 pods and 6 partitions per topic. The custom run configured partitionAssigners: [coPartitionAssigner] on every consumer. Each pod recorded its GROUP_JOIN assignment. I then counted whether each notification arrived on a pod that also owned its market partition.

With the custom assigner, each pod had matching market and notification partition sets. All 30 notifications (10 symbols times 3 notifications each) joined locally: 30/30 local and 0 remote. The default round robin run produced 0 local and 30 remote joins. The data and keys were unchanged between runs.

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 local join

An aligned assignment matters when you use it, so the second run does a partition local join.

Each pod subscribes to both topics. It builds a Map<StockSymbol, number> of the latest prices 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 arrives, the pod checks its local price map. In my test, the custom assignment let the check succeed for every notification. However, assignment alone does not guarantee that the market record has arrived first. Kafka does not order records across topics.

After startup or rebalance, the application must restore or initialize state before it relies on a local join. With round robin, localPrice.has(n.symbol) can also fail because another pod owns the market partition. The test counts that case as a miss.

A local map can avoid a network lookup during alert evaluation. Matching assignments make this possible, but state recovery, cross topic arrival order, and stale values still need explicit handling.

The Redis analog, briefly

Redis Cluster maps keys to 16384 slots with CRC16. Hash tags in market:{BBCA} and notif:{BBCA} select the same hash input and slot. Without tags, market:BBCA and notif:BBCA can use different slots.

Redis tags control data placement on nodes. The Kafka assigner controls partition ownership by consumer processes. Both can support local operations, but they operate at different boundaries.

What the assigner guarantees

• Co partitioning (same partition count, same key, same partitioner) guarantees data layout lines up, while consumer assignment determines which process reads that data.
• The built in round robin assigner can place matching partitions on different consumers. The custom assigner preserves their alignment. 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.
• Group topic partitions by index. Assign partition p from each subscribed topic to the same member.
• Measured over the same 4 pods: 30/30 notifications join locally with my assigner, 0/30 with round robin.
• Matching assignments let a consumer join notifications with its local price map. The application must still initialize and restore that map when ownership changes.

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.