python基于redis的布隆过滤器示例

"""
布隆过滤器基础类
"""
import mmh3
from redis import Redis
import math
from typing import List, Union, Tuple
from loguru import logger
import hashlib
​
​
class BaseBloom:
    # 内置种子,用于对值做哈希计算
    SEEDS = [
        543, 460, 171, 876, 796, 607, 650, 81, 837, 545, 591, 946, 846, 521, 913, 636, 878, 735, 414, 372,
        344, 324, 223, 180, 327, 891, 798, 933, 493, 293, 836, 10, 6, 544, 924, 849, 438, 41, 862, 648, 338,
        465, 562, 693, 979, 52, 763, 103, 387, 374, 349, 94, 384, 680, 574, 480, 307, 580, 71, 535, 300, 53,
        481, 519, 644, 219, 686, 236, 424, 326, 244, 212, 909, 202, 951, 56, 812, 901, 926, 250, 507, 739, 371,
        63, 584, 154, 7, 284, 617, 332, 472, 140, 605, 262, 355, 526, 647, 923, 199, 518
    ]
​
    def __init__(
            self,
            capacity: int,
            error_rate: float,
            redis_conn: Redis,
            piece_size_byte=1024 * 2,  # 分片bitmap的大小
            name: str = "bloom",
            tag: str = "default"
    ):
        self.capacity = capacity
        self.error_rate = error_rate
        self.redis_conn = redis_conn
        self.piece_size_byte = piece_size_byte  # bit map片的大小
        self.piece_size_bit = piece_size_byte * 8
        self.name = name
        self.tag = tag
        self._m = math.ceil(capacity * math.log2(math.e) * math.log2(1 / error_rate))  # 至少需要位数
        self._k = math.ceil(math.log1p(2) * self._m / capacity)  # 填充时哈希计算次数，对应种子个数
        self._seed_list = self._gen_seed_list()  # 此过滤器使用的种子列表
        self.redis_key = f"{name}:{tag}"
        self._N = 2 ** 31 - 1
        if self._m > self._N * 2 + 1:
            raise Exception("当前布隆过滤器最大支持位数无法满足要求，请更换mmh3哈希计算方法及\"_N\"值")
​
        logger.info(
            f"布隆过滤器{self.name}:{self.tag}已生成,\n"
            f"最大容量:{self.capacity},\n"
            f"误报率:{self.error_rate * 100}%,\n"
            f"最大内存占用:{(self._N * 2 + 2) / 8 / 1024 / 1024}MB,\n"
            f"使用种子数:{len(self._seed_list)},\n"
            f"单片bitmap大小{self.piece_size_byte}Bytes\n"
            f"bitmap最大创建数量:{(self._N * 2 + 2) / self.piece_size_bit}"
        )
​
    @staticmethod
    def _sha1_hash(s: str) -> str:
        """
        sha1哈希计算
        :param s:
        :return:
        """
        sha1_obj = hashlib.sha1()
        sha1_obj.update(s.encode())
        return sha1_obj.hexdigest()
​
    def _gen_seed_list(self) -> List[int]:
        """
        根据内置的种子列表生成布隆过滤器的种子列表
        :return:
        """
        free_seed_num = len(self.SEEDS) - self._k
        if free_seed_num < 0:
            raise Exception(f"内置哈希计算种子不足,请扩大种子列表，至少新增{abs(free_seed_num)}个种子")
        return self.SEEDS[0:self._k]
​
    def _hash_calculate(self, value: str) -> List[int]:
        """
        哈希计算
        :param value:
        :return:
        """
        sha1_hash_v = self._sha1_hash(value)
        hash_list = []
        for s in self._seed_list:
            h = mmh3.hash(sha1_hash_v, s)
            if h >= 0:
                hash_list.append(h)
            else:
                hash_list.append(self._N - h)
        return hash_list
​
    def add(self, value: Union[str, int, float]):
        """
        填充
        :param value:
        :return:
        """
        hash_list = self._hash_calculate(str(value))
        for h in hash_list:
            key_index, bit_map_index = self.search_index(h)
            redis_key = f"{self.redis_key}:{key_index}"
            self.redis_conn.setbit(redis_key, bit_map_index, 1)
​
    def is_existed(self, value: Union[str, int, float]) -> bool:
        """
        查询值是否在过滤器中填充过
        :param value:
        :return:
        """
        hash_list = self._hash_calculate(str(value))
        flag = 1
        for h in hash_list:
            key_index, bit_map_index = self.search_index(h)
            redis_key = f"{self.redis_key}:{key_index}"
            flag = flag & self.redis_conn.getbit(redis_key, bit_map_index)
            if not flag:
                return False
​
        return True
​
    def search_index(self, h: int) -> Tuple[int, int]:
        """
        查找key索引和对应key下的bitmap索引
        :param h:
        :return:
        """
        key_index = h // self.piece_size_bit
        bit_map_index = h % self.piece_size_bit
        return key_index, bit_map_index
​