Merge pull request redis#453 from pepijndevos/master

Implement HyperLogLog commands

Author: andymccurdy

redis/client.py

@@ -276,9 +276,9 @@ class StrictRedis(object):
         ),
         string_keys_to_dict(
             'BITCOUNT DECRBY DEL GETBIT HDEL HLEN INCRBY LINSERT LLEN LPUSHX '
-            'RPUSHX SADD SCARD SDIFFSTORE SETBIT SETRANGE SINTERSTORE SREM '
-            'STRLEN SUNIONSTORE ZADD ZCARD ZREM ZREMRANGEBYRANK '
-            'ZREMRANGEBYSCORE',
+            'PFADD PFCOUNT RPUSHX SADD SCARD SDIFFSTORE SETBIT SETRANGE '
+            'SINTERSTORE SREM STRLEN SUNIONSTORE ZADD ZCARD ZREM '
+            'ZREMRANGEBYRANK ZREMRANGEBYSCORE',
             int
         ),
         string_keys_to_dict('INCRBYFLOAT HINCRBYFLOAT', float),
@@ -290,7 +290,7 @@ class StrictRedis(object):
         string_keys_to_dict('SORT', sort_return_tuples),
         string_keys_to_dict('ZSCORE ZINCRBY', float_or_none),
         string_keys_to_dict(
-            'FLUSHALL FLUSHDB LSET LTRIM MSET RENAME '
+            'FLUSHALL FLUSHDB LSET LTRIM MSET PFMERGE RENAME '
             'SAVE SELECT SHUTDOWN SLAVEOF WATCH UNWATCH',
             lambda r: nativestr(r) == 'OK'
         ),
@@ -1527,6 +1527,22 @@ def _zaggregate(self, command, dest, keys, aggregate=None):
             pieces.append(aggregate)
         return self.execute_command(*pieces)
 
+    # HYPERLOGLOG COMMANDS
+    def pfadd(self, name, *values):
+        "Adds the specified elements to the specified HyperLogLog."
+        return self.execute_command('PFADD', name, *values)
+
+    def pfcount(self, name):
+        """
+        Return the approximated cardinality of
+        the set observed by the HyperLogLog at key.
+        """
+        return self.execute_command('PFCOUNT', name)
+
+    def pfmerge(self, dest, *sources):
+        "Merge N different HyperLogLogs into a single one."
+        return self.execute_command('PFMERGE', dest, *sources)
+
     # HASH COMMANDS
     def hdel(self, name, *keys):
         "Delete ``keys`` from hash ``name``"

tests/test_commands.py

@@ -951,6 +951,33 @@ def test_zunionstore_with_weight(self, r):
         assert r.zrange('d', 0, -1, withscores=True) == \
             [(b('a2'), 5), (b('a4'), 12), (b('a3'), 20), (b('a1'), 23)]
 
+    # HYPERLOGLOG TESTS
+    @skip_if_server_version_lt('2.8.9')
+    def test_pfadd(self, r):
+        members = set([b('1'), b('2'), b('3')])
+        assert r.pfadd('a', *members) == 1
+        assert r.pfadd('a', *members) == 0
+        assert r.pfcount('a') == len(members)
+
+    @skip_if_server_version_lt('2.8.9')
+    def test_pfcount(self, r):
+        members = set([b('1'), b('2'), b('3')])
+        r.pfadd('a', *members)
+        assert r.pfcount('a') == len(members)
+
+    @skip_if_server_version_lt('2.8.9')
+    def test_pfmerge(self, r):
+        mema = set([b('1'), b('2'), b('3')])
+        memb = set([b('2'), b('3'), b('4')])
+        memc = set([b('5'), b('6'), b('7')])
+        r.pfadd('a', *mema)
+        r.pfadd('b', *memb)
+        r.pfadd('c', *memc)
+        r.pfmerge('d', 'c', 'a')
+        assert r.pfcount('d') == 6
+        r.pfmerge('d', 'b')
+        assert r.pfcount('d') == 7
+
     # HASH COMMANDS
     def test_hget_and_hset(self, r):
         r.hmset('a', {'1': 1, '2': 2, '3': 3})