import org.apache.kafka.clients.producer.Partitioner;

import org.apache.kafka.common.Cluster;

import org.apache.kafka.common.PartitionInfo;

import org.apache.kafka.common.InvalidRecordException;

import org.apache.kafka.common.utils.Utils;

​

import java.util.List;

import java.util.Map;

​

// Class must implement the org.apache.kafka.clients.producer.Partitioner interface

public class CustomPartitioner implements Partitioner {

​

    // Pass any configurations to use in this method

    public void configure(Map<String, ?> configs) {

    }

​

    // The logic to exercise custom partitioning

    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes,

                         Cluster cluster) {

​

        // Retrieve the number of partitions for the topic

        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);

        int numPartitions = partitions.size();

​

        // Throw an exception if key isn't specified

        if ((keyBytes == null) || (!(key instanceof String)))

            throw new InvalidRecordException("Message is missing the key");

​

        // Messages with key=USA will always land in the last partition

        if (((String) key).equals("USA"))

            return numPartitions;

​

        // Other records get hashed to the rest of the partitions excluding the last one

        return (Math.abs(Utils.murmur2(keyBytes)) % (numPartitions - 1));

    }

​

    public void close() {

    }

}