Repository: cunla/fakeredis-py

Issue:
Implement support for `TOPK.RESERVE` (topk command)
Implement support for command `TOPK.RESERVE` in topk_mixin.py.

Initializes a TopK with specified parameters.

Here is the [Official documentation](https://redis.io/commands/topk.reserve/)

Discussion:

Patch:
diff --git a/fakeredis/_fakesocket.py b/fakeredis/_fakesocket.py
index 7ffd29e..b988145 100644
--- a/fakeredis/_fakesocket.py
+++ b/fakeredis/_fakesocket.py
@@ -1,4 +1,4 @@
-from fakeredis.stack import JSONCommandsMixin, BFCommandsMixin, CFCommandsMixin, CMSCommandsMixin
+from fakeredis.stack import JSONCommandsMixin, BFCommandsMixin, CFCommandsMixin, CMSCommandsMixin, TopkCommandsMixin
 from ._basefakesocket import BaseFakeSocket
 from .commands_mixins.bitmap_mixin import BitmapCommandsMixin
 from .commands_mixins.connection_mixin import ConnectionCommandsMixin
@@ -43,6 +43,7 @@ class FakeSocket(
     BFCommandsMixin,
     CFCommandsMixin,
     CMSCommandsMixin,
+    TopkCommandsMixin,
 ):
     def __init__(self, server, db):
         super(FakeSocket, self).__init__(server, db)
diff --git a/fakeredis/stack/__init__.py b/fakeredis/stack/__init__.py
index 605e3da..37e5513 100644
--- a/fakeredis/stack/__init__.py
+++ b/fakeredis/stack/__init__.py
@@ -1,3 +1,5 @@
+from ._topk_mixin import TopkCommandsMixin  # noqa: F401
+
 try:
     from jsonpath_ng.ext import parse  # noqa: F401
     from redis.commands.json.path import Path  # noqa: F401
diff --git a/fakeredis/stack/_topk_mixin.py b/fakeredis/stack/_topk_mixin.py
new file mode 100644
index 0000000..c38172e
--- /dev/null
+++ b/fakeredis/stack/_topk_mixin.py
@@ -0,0 +1,221 @@
+"""Command mixin for emulating `redis-py`'s top-k functionality."""
+import heapq
+import random
+import time
+from collections import Counter
+from typing import Any, List, Optional, Tuple
+
+from fakeredis import _msgs as msgs
+from fakeredis._command_args_parsing import extract_args
+from fakeredis._commands import Key, Int, Float, command, CommandItem
+from fakeredis._helpers import OK, SimpleError
+
+random.seed(time.time())
+
+
+class Bucket(object):
+
+    def __init__(self, counter: int, fingerprint: int):
+        self.counter = counter
+        self.fingerprint = fingerprint
+
+    def add(self, fingerprint: int, incr: int, decay: float) -> int:
+        if self.fingerprint == fingerprint:
+            self.counter += incr
+            return self.counter
+        elif self._decay(decay):
+            self.counter += incr
+            self.fingerprint = fingerprint
+            return self.counter
+        return 0
+
+    def count(self, fingerprint: int) -> int:
+        if self.fingerprint == fingerprint:
+            return self.counter
+        return 0
+
+    def _decay(self, decay: float) -> bool:
+        if self.counter > 0:
+            probability = decay ** self.counter
+            if probability >= 1 or random.random() < probability:
+                self.counter -= 1
+        return self.counter == 0
+
+
+class HashArray(object):
+    def __init__(self, width: int, decay: float):
+        self.width = width
+        self.decay = decay
+        self.array = [Bucket(0, 0) for _ in range(width)]
+        self._seed = random.getrandbits(32)
+
+    def count(self, item: bytes) -> int:
+        return self.get_bucket(item).count(self._hash(item))
+
+    def add(self, item: bytes, incr: int) -> int:
+        bucket = self.get_bucket(item)
+        return bucket.add(self._hash(item), incr, self.decay)
+
+    def get_bucket(self, item: bytes) -> Bucket:
+        return self.array[self._hash(item) % self.width]
+
+    def _hash(self, item: bytes) -> int:
+        return hash(item) ^ self._seed
+
+
+class HeavyKeeper(object):
+    def __init__(self, k: int, width: int = 1024, depth: int = 5, decay: float = 0.9):
+        self.k = k
+        self.width = width
+        self.depth = depth
+        self.decay = decay
+        self.hash_arrays = [HashArray(width, decay) for _ in range(depth)]
+        self.min_heap = list()
+
+    def _index(self, val: bytes) -> int:
+        for ind, item in enumerate(self.min_heap):
+            if item[1] == val:
+                return ind
+        return -1
+
+    def add(self, item: bytes, incr: int) -> Optional[bytes]:
+        max_count = 0
+        for i in range(self.depth):
+            count = self.hash_arrays[i].add(item, incr)
+            max_count = max(max_count, count)
+        if len(self.min_heap) < self.k:
+            heapq.heappush(self.min_heap, (max_count, item))
+            return None
+        ind = self._index(item)
+        if ind >= 0:
+            self.min_heap[ind] = (max_count, item)
+            heapq.heapify(self.min_heap)
+            return None
+        if max_count > self.min_heap[0][0]:
+            expelled = heapq.heapreplace(self.min_heap, (max_count, item))
+            return expelled[1]
+
+
+    def count(self, item: bytes) -> int:
+        ind = self._index(item)
+        if ind > 0:
+            return self.min_heap[ind][0]
+        return max([ha.count(item) for ha in self.hash_arrays])
+
+    def list(self, k: Optional[int] = None) -> List[Tuple[int, bytes]]:
+        sorted_list = sorted(self.min_heap, key=lambda x: x[0], reverse=True)
+        if k is None:
+            return sorted_list
+        return sorted_list[:k]
+
+
+class TopkCommandsMixin:
+    """`CommandsMixin` for enabling TopK compatibility in `fakeredis`."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+    @command(name="TOPK.ADD", fixed=(Key(HeavyKeeper), bytes,), repeat=(bytes,), flags=msgs.FLAG_NO_INITIATE, )
+    def topk_add(self, key: CommandItem, *args: bytes) -> List[Optional[bytes]]:
+        if key.value is None:
+            raise SimpleError("TOPK: key does not exist")
+        if not isinstance(key.value, HeavyKeeper):
+            raise SimpleError("TOPK: key is not a HeavyKeeper")
+        res = [key.value.add(_item, 1) for _item in args]
+        key.updated()
+        return res
+
+    @command(name="TOPK.COUNT",
+             fixed=(Key(HeavyKeeper), bytes,),
+             repeat=(bytes,),
+             flags=msgs.FLAG_NO_INITIATE, )
+    def topk_count(self, key: CommandItem, *args: bytes) -> List[Optional[bytes]]:
+        if key.value is None:
+            raise SimpleError("TOPK: key does not exist")
+        if not isinstance(key.value, HeavyKeeper):
+            raise SimpleError("TOPK: key is not a HeavyKeeper")
+        res = [key.value.count(_item) for _item in args]
+        return res
+
+    @command(name="TOPK.QUERY",
+             fixed=(Key(HeavyKeeper), bytes,),
+             repeat=(bytes,),
+             flags=msgs.FLAG_NO_INITIATE, )
+    def topk_query(self, key: CommandItem, *args: bytes) -> List[Optional[bytes]]:
+        if key.value is None:
+            raise SimpleError("TOPK: key does not exist")
+        if not isinstance(key.value, HeavyKeeper):
+            raise SimpleError("TOPK: key is not a HeavyKeeper")
+        topk = {item[1] for item in key.value.list()}
+        res = [1 if _item in topk else 0 for _item in args]
+        return res
+
+    @command(name="TOPK.INCRBY", fixed=(Key(), bytes, Int,), repeat=(bytes, Int), flags=msgs.FLAG_NO_INITIATE, )
+    def topk_incrby(self, key, *args):
+        if key.value is None:
+            raise SimpleError("TOPK: key does not exist")
+        if not isinstance(key.value, HeavyKeeper):
+            raise SimpleError("TOPK: key is not a HeavyKeeper")
+        if len(args) % 2 != 0:
+            raise SimpleError("TOPK: number of arguments must be even")
+        res = list()
+        for i in range(0, len(args), 2):
+            val, count = args[i], int(args[i + 1])
+            res.append(key.value.add(val, count))
+        key.updated()
+        return res
+
+    @command(name="TOPK.INFO", fixed=(Key(),), repeat=(), flags=msgs.FLAG_NO_INITIATE, )
+    def topk_info(self, key):
+        if key.value is None:
+            raise SimpleError("TOPK: key does not exist")
+        if not isinstance(key.value, HeavyKeeper):
+            raise SimpleError("TOPK: key is not a HeavyKeeper")
+        return [
+            b"k", key.value.k,
+            b"width", key.value.width,
+            b"depth", key.value.depth,
+            b"decay", key.value.decay,
+        ]
+
+    @command(name="TOPK.LIST", fixed=(Key(),), repeat=(bytes,), flags=msgs.FLAG_NO_INITIATE, )
+    def topk_list(self, key, *args):
+        (withcount,), _ = extract_args(args, ("withcount",))
+        if key.value is None:
+            raise SimpleError("TOPK: key does not exist")
+        if not isinstance(key.value, HeavyKeeper):
+            raise SimpleError("TOPK: key is not a HeavyKeeper")
+        res = key.value.list()
+        if not withcount:
+            res = [item[1] for item in res]
+        else:
+            res = [[item[1], item[0]] for item in res]
+            res = [item for sublist in res for item in sublist]
+        return res
+
+    @command(name="TOPK.RESERVE", fixed=(Key(), Int,), repeat=(Int, Int, Float,), flags=msgs.FLAG_NO_INITIATE, )
+    def topk_reserve(self, key, topk, *args):
+        if len(args) == 3:
+            width, depth, decay = args
+        else:
+            width, depth, decay = 8, 7, 0.9
+        if key.value is not None:
+            raise SimpleError("TOPK: key already set")
+        key.update(HeavyKeeper(topk, width, depth, decay))
+        return OK
+
+
+if __name__ == '__main__':
+    real = Counter()
+    heavy_keeper = HeavyKeeper(10)
+    random.seed(10)
+    for i in range(100000):
+        ip_addr = f'{random.randint(0, 100)}'.encode()
+        real[ip_addr] += 1
+        heavy_keeper.add(ip_addr, 1)
+        if i % 10000 == 0:
+            print(heavy_keeper.list(5))
+    items = heavy_keeper.list(5)
+    for item in items:
+        print(f"{item[1]}: count={item[0]} real_count={real[item[1]]}")
+    print(real)


Test patch:
diff --git a/test/test_stack/test_topk.py b/test/test_stack/test_topk.py
new file mode 100644
index 0000000..b230ea4
--- /dev/null
+++ b/test/test_stack/test_topk.py
@@ -0,0 +1,39 @@
+import pytest
+import redis
+
+
+def test_topk_incrby(r: redis.Redis):
+    assert r.topk().reserve("topk", 3, 10, 3, 1)
+    assert [None, None, None] == r.topk().incrby(
+        "topk", ["bar", "baz", "42"], [3, 6, 2]
+    )
+    assert [None, "bar"] == r.topk().incrby("topk", ["42", "xyzzy"], [8, 4])
+    with pytest.deprecated_call():
+        assert [3, 6, 10, 4, 0] == r.topk().count(
+            "topk", "bar", "baz", "42", "xyzzy", 4
+        )
+
+
+def test_topk(r: redis.Redis):
+    # test list with empty buckets
+    assert r.topk().reserve("topk", 3, 50, 4, 0.9)
+    ret = r.topk().add("topk", "A", "B", "C", "D", "D", "E", "A", "A", "B", "C", "G", "D", "B", "D", "A", "E", "E", 1)
+    assert len(ret) == 18
+
+    with pytest.deprecated_call():
+        assert r.topk().count("topk", "A", "B", "C", "D", "E", "F", "G") == [4, 3, 2, 4, 3, 0, 1]
+    ret = r.topk().query("topk", "A", "B", "C", "D", "E", "F", "G")
+    assert (ret == [1, 0, 0, 1, 1, 0, 0]) or (ret == [1, 1, 0, 1, 0, 0, 0])
+    # test full list
+    assert r.topk().reserve("topklist", 3, 50, 3, 0.9)
+    assert r.topk().add("topklist", "A", "B", "D", "E", "A", "A", "B", "C", "G", "D", "B", "A", "B", "E", "E",
+                        )
+    with pytest.deprecated_call():
+        assert r.topk().count("topklist", "A", "B", "C", "D", "E", "F", "G") == [4, 4, 1, 2, 3, 0, 1]
+    assert r.topk().list("topklist") == ["A", "B", "E"]
+    assert r.topk().list("topklist", withcount=True) == ["A", 4, "B", 4, "E", 3]
+    info = r.topk().info("topklist")
+    assert 3 == info["k"]
+    assert 50 == info["width"]
+    assert 3 == info["depth"]
+    assert 0.9 == round(float(info["decay"]), 1)