Repository: Parsely/pykafka
Pull request #593 (2016-08-18 15:20:59)

Adding GroupHashingPartitioner
 This is a specialized HashingPartitioner which can divide incoming messages between a configurable subset of the available partitions. For example, given a total of sixteen partitions, one could specify a group size of four to smooth the distribution of incoming messages across a group of four partitions, thus avoiding hot spots in one partition.


Patch:
diff --git a/pykafka/partitioners.py b/pykafka/partitioners.py
index b46c0fa..207fa6b 100644
--- a/pykafka/partitioners.py
+++ b/pykafka/partitioners.py
@@ -17,7 +17,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 __all__ = ["random_partitioner", "BasePartitioner", "HashingPartitioner",
-           "hashing_partitioner"]
+           "hashing_partitioner", "GroupHashingPartitioner"]
 import random
 
 from hashlib import sha1
@@ -83,3 +83,53 @@ class HashingPartitioner(BasePartitioner):
 
 
 hashing_partitioner = HashingPartitioner()
+
+
+class GroupHashingPartitioner(BasePartitioner):
+    """
+    Messages published with the identical keys will be directed to a consistent subset of 'n'
+    partitions from the set of available partitions. For example, if there are 16 partitions and group_size=4,
+    messages with the identical keys will be shared equally between a subset of four partitions, instead of always
+    being directed to the same partition.
+
+    The same guarantee caveats apply as to the :class:`pykafka.base.HashingPartitioner`.
+    """
+    def __init__(self, hash_func, group_size=1):
+        """
+        :param hash_func: A hash function
+        :type hash_func: function
+        :param group_size: Size of the partition group to assign to. For example, if there are 16 partitions, and we
+            want to smooth the distribution of identical keys between a set of 4, use 4 as the group_size.
+        :type group_size: Integer value between (0, total_partition_count)
+        """
+        self.hash_func = hash_func
+        self.group_size = group_size
+        if self.hash_func is None:
+            raise ValueError(
+                'hash_func must be specified when using GroupHashingPartitioner'
+            )
+        if self.group_size < 1:
+            raise ValueError(
+                'group_size cannot be < 1 when using GroupHashingPartitioner'
+            )
+
+    def __call__(self, partitions, key):
+        """
+        :param partitions: The partitions from which to choose
+        :type partitions: sequence of :class:`pykafka.base.BasePartition`
+        :param key: Key used for routing
+        :type key: Any hashable type if using the default :func:`hash`
+            implementation, any valid value for your custom hash function
+        :returns: A partition
+        :rtype: :class:`pykafka.base.BasePartition`
+        """
+        if key is None:
+            raise ValueError(
+                'key cannot be `None` when using hashing partitioner'
+            )
+        if self.group_size > len(partitions):
+            raise ValueError(
+                'group_size cannot be > available partitions'
+            )
+        partitions = sorted(partitions)
+        return partitions[abs(self.hash_func(key) + random.randrange(0, self.group_size)) % len(partitions)]


Test patch:
diff --git a/tests/pykafka/test_partitioners.py b/tests/pykafka/test_partitioners.py
new file mode 100644
index 0000000..3ff61ec
--- /dev/null
+++ b/tests/pykafka/test_partitioners.py
@@ -0,0 +1,52 @@
+import unittest2
+
+from hashlib import sha1
+from pykafka.partitioners import GroupHashingPartitioner
+
+
+class TestGroupHashingPartitioner(unittest2.TestCase):
+
+    def test_valid_inputs_success(self):
+        # Test data; 1st element is group size, 2nd is total partition count, 3rd is the key
+        key = 'foo'.encode('utf-8')
+        data = [[1, 16, key],
+                [2, 16, key],
+                [4, 16, key],
+                [16, 16, key],
+                [1, 1, key]]
+        for row in data:
+            self._run_test(row[0], row[1], row[2])
+
+    def test_invalid_inputs_error(self):
+        key = 'foo'.encode('utf-8')
+        data = [[0, 16, key],
+                [17, 16, key]]
+        for row in data:
+            with self.assertRaises(ValueError):
+                self._run_test(row[0], row[1], row[2])
+                continue
+
+    def test_create_with_zero_group_size_raises_error(self):
+        with self.assertRaises(ValueError):
+            GroupHashingPartitioner(hash_func=None, group_size=0)
+
+    def test_create_with_negative_group_size_raises_error(self):
+        with self.assertRaises(ValueError):
+            GroupHashingPartitioner(hash_func=None, group_size=-1)
+
+    def test_missing_hash_function_raises_error(self):
+        with self.assertRaises(ValueError):
+            GroupHashingPartitioner(hash_func=None, group_size=1)
+
+    def _run_test(self, group_size, total_partition_count, key):
+        def hash_func(k): return int(sha1(k).hexdigest(), 16)
+        # Instead of one hash for each key, generate a list of n possible hashes, where n=group_size
+        hashed_keys = [abs(hash_func(key) + x) for x in range(group_size)]
+        # Obtain a list of valid partitions
+        valid_partitions = list(map(lambda x: x % total_partition_count, hashed_keys))
+        self.assertEquals(len(valid_partitions), group_size)
+
+        # Call the partitioner and check that the returned partition is in the valid list
+        partitioner = GroupHashingPartitioner(hash_func, group_size)
+        x = partitioner.__call__(list(range(total_partition_count)), key)
+        self.assertTrue(x in valid_partitions)