75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
"""Jump consistent hash (Lamping & Veach, 2014).
|
|
|
|
Maps a 64-bit key to a bucket in [0, num_buckets) with minimal
|
|
rebalancing when the bucket count changes: only ~1/(k+1) of keys move
|
|
when going from k to k+1 buckets.
|
|
|
|
The reference algorithm uses 64-bit wrapping arithmetic and a double
|
|
for the jump computation. This implementation reproduces it exactly
|
|
using Python ints masked to 64 bits, so results match the C reference.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import struct
|
|
from typing import Any, Hashable
|
|
|
|
|
|
_MASK64 = (1 << 64) - 1
|
|
_MULT = 2862933555777941757
|
|
_SCALE = 1 << 31 # 2^31
|
|
|
|
|
|
def _to_u64(key: Any) -> int:
|
|
if isinstance(key, int):
|
|
return key & _MASK64
|
|
if isinstance(key, bytes):
|
|
raw = key
|
|
else:
|
|
raw = str(key).encode("utf-8")
|
|
# md5 -> 128 bits -> fold to 64
|
|
digest = hashlib.md5(raw).digest()
|
|
(hi, lo) = struct.unpack(">QQ", digest)
|
|
return (hi ^ lo) & _MASK64
|
|
|
|
|
|
def jump_hash(key: Any, num_buckets: int) -> int:
|
|
"""Return bucket index in [0, num_buckets) for key."""
|
|
if not isinstance(num_buckets, int) or isinstance(num_buckets, bool):
|
|
raise TypeError("num_buckets must be an int")
|
|
if num_buckets < 0:
|
|
raise ValueError("num_buckets must be >= 0")
|
|
if num_buckets == 0:
|
|
raise ValueError("num_buckets must be >= 1")
|
|
if num_buckets == 1:
|
|
return 0
|
|
|
|
k = _to_u64(key)
|
|
b = -1
|
|
j = 0
|
|
while j < num_buckets:
|
|
b = j
|
|
k = (k * _MULT + 1) & _MASK64
|
|
# j = (b+1) * 2^31 / ((key >> 33) + 1)
|
|
denom = (k >> 33) + 1
|
|
j = (b + 1) * _SCALE // denom
|
|
return b
|
|
|
|
|
|
class JumpConsistentHash:
|
|
"""Convenience wrapper keeping a bucket count."""
|
|
|
|
def __init__(self, num_buckets: int):
|
|
if num_buckets < 1:
|
|
raise ValueError("num_buckets must be >= 1")
|
|
self.num_buckets = num_buckets
|
|
|
|
def bucket(self, key: Any) -> int:
|
|
return jump_hash(key, self.num_buckets)
|
|
|
|
def resize(self, new_num_buckets: int) -> "JumpConsistentHash":
|
|
if new_num_buckets < 1:
|
|
raise ValueError("num_buckets must be >= 1")
|
|
self.num_buckets = new_num_buckets
|
|
return self
|