Add support for WEIGHTS and AGGREGATE in ZUNION/ZINTER.

redis · May 10, 2010 · 002a994 · 002a994
1 parent 1c6022c
commit 002a994
Show file tree

Hide file tree

Showing 4 changed files with 164 additions and 11 deletions.
diff --git a/lib/redis.rb b/lib/redis.rb
@@ -290,12 +290,22 @@ def zrem(key, member)
     _bool @client.call(:zrem, key, member)
   end
 
-  def zinter(destination, keys)
-    @client.call(:zinter, destination, keys.size, *keys)
+  def zinter(destination, keys, options = {})
+    command = CommandOptions.new(options) do |c|
+      c.splat :weights
+      c.value :aggregate
+    end
+
+    @client.call(:zinter, destination, keys.size, *(keys + command.to_a))
   end
 
-  def zunion(destination, keys)
-    @client.call(:zunion, destination, keys.size, *keys)
+  def zunion(destination, keys, options = {})
+    command = CommandOptions.new(options) do |c|
+      c.splat :weights
+      c.value :aggregate
+    end
+
+    @client.call(:zunion, destination, keys.size, *(keys + command.to_a))
   end
 
   def move(key, db)

diff --git a/lib/redis/distributed.rb b/lib/redis/distributed.rb
@@ -354,15 +354,15 @@ def zscore(key, member)
       node_for(key).zscore(key, member)
     end
 
-    def zinter(destination, keys)
+    def zinter(destination, keys, options = {})
       ensure_same_node(:zinter, destination, *keys) do |node|
-        node.zinter(destination, keys)
+        node.zinter(destination, keys, options)
       end
     end
 
-    def zunion(destination, keys)
+    def zunion(destination, keys, options = {})
       ensure_same_node(:zunion, destination, *keys) do |node|
-        node.zunion(destination, keys)
+        node.zunion(destination, keys, options)
       end
     end
 

diff --git a/test/distributed_test.rb b/test/distributed_test.rb
@@ -895,6 +895,35 @@ class RedisDistributedTest < Test::Unit::TestCase
       assert_equal 4, @r.zunion("{qux}foobar", ["{qux}foo", "{qux}bar"])
     end
 
+    test "ZUNION with WEIGHTS" do
+      @r.zadd "{qux}foo", 1, "s1"
+      @r.zadd "{qux}foo", 3, "s3"
+      @r.zadd "{qux}bar", 20, "s2"
+      @r.zadd "{qux}bar", 40, "s4"
+
+      assert_equal 4, @r.zunion("{qux}foobar", ["{qux}foo", "{qux}bar"])
+      assert_equal ["s1", "s3", "s2", "s4"], @r.zrange("{qux}foobar", 0, -1)
+
+      assert_equal 4, @r.zunion("{qux}foobar", ["{qux}foo", "{qux}bar"], :weights => [10, 1])
+      assert_equal ["s1", "s2", "s3", "s4"], @r.zrange("{qux}foobar", 0, -1)
+    end
+
+    test "ZUNION with AGGREGATE" do
+      @r.zadd "{qux}foo", 1, "s1"
+      @r.zadd "{qux}foo", 2, "s2"
+      @r.zadd "{qux}bar", 4, "s2"
+      @r.zadd "{qux}bar", 3, "s3"
+
+      assert_equal 3, @r.zunion("{qux}foobar", ["{qux}foo", "{qux}bar"])
+      assert_equal ["s1", "s3", "s2"], @r.zrange("{qux}foobar", 0, -1)
+
+      assert_equal 3, @r.zunion("{qux}foobar", ["{qux}foo", "{qux}bar"], :aggregate => :min)
+      assert_equal ["s1", "s2", "s3"], @r.zrange("{qux}foobar", 0, -1)
+
+      assert_equal 3, @r.zunion("{qux}foobar", ["{qux}foo", "{qux}bar"], :aggregate => :max)
+      assert_equal ["s1", "s3", "s2"], @r.zrange("{qux}foobar", 0, -1)
+    end
+
     test "ZINTER" do
       assert_raises Redis::Distributed::CannotDistribute do
         @r.zinter("foobar", ["foo", "bar"])
@@ -918,6 +947,48 @@ class RedisDistributedTest < Test::Unit::TestCase
       assert_equal 1, @r.zinter("{qux}foobar", ["{qux}foo", "{qux}bar"])
       assert_equal ["s1"], @r.zrange("{qux}foobar", 0, 2)
     end
+
+    test "ZINTER with WEIGHTS" do
+      @r.zadd "{qux}foo", 1, "s1"
+      @r.zadd "{qux}foo", 2, "s2"
+      @r.zadd "{qux}foo", 3, "s3"
+      @r.zadd "{qux}bar", 20, "s2"
+      @r.zadd "{qux}bar", 30, "s3"
+      @r.zadd "{qux}bar", 40, "s4"
+
+      assert_equal 2, @r.zinter("{qux}foobar", ["{qux}foo", "{qux}bar"])
+      assert_equal ["s2", "s3"], @r.zrange("{qux}foobar", 0, -1)
+
+      assert_equal 2, @r.zinter("{qux}foobar", ["{qux}foo", "{qux}bar"], :weights => [10, 1])
+      assert_equal ["s2", "s3"], @r.zrange("{qux}foobar", 0, -1)
+
+      assert_equal "40", @r.zscore("{qux}foobar", "s2")
+      assert_equal "60", @r.zscore("{qux}foobar", "s3")
+    end
+
+    test "ZINTER with AGGREGATE" do
+      @r.zadd "{qux}foo", 1, "s1"
+      @r.zadd "{qux}foo", 2, "s2"
+      @r.zadd "{qux}foo", 3, "s3"
+      @r.zadd "{qux}bar", 20, "s2"
+      @r.zadd "{qux}bar", 30, "s3"
+      @r.zadd "{qux}bar", 40, "s4"
+
+      assert_equal 2, @r.zinter("{qux}foobar", ["{qux}foo", "{qux}bar"])
+      assert_equal ["s2", "s3"], @r.zrange("{qux}foobar", 0, -1)
+      assert_equal "22", @r.zscore("{qux}foobar", "s2")
+      assert_equal "33", @r.zscore("{qux}foobar", "s3")
+
+      assert_equal 2, @r.zinter("{qux}foobar", ["{qux}foo", "{qux}bar"], :aggregate => :min)
+      assert_equal ["s2", "s3"], @r.zrange("{qux}foobar", 0, -1)
+      assert_equal "2", @r.zscore("{qux}foobar", "s2")
+      assert_equal "3", @r.zscore("{qux}foobar", "s3")
+
+      assert_equal 2, @r.zinter("{qux}foobar", ["{qux}foo", "{qux}bar"], :aggregate => :max)
+      assert_equal ["s2", "s3"], @r.zrange("{qux}foobar", 0, -1)
+      assert_equal "20", @r.zscore("{qux}foobar", "s2")
+      assert_equal "30", @r.zscore("{qux}foobar", "s3")
+    end
   end
 
   context "Commands operating on hashes" do

diff --git a/test/redis_test.rb b/test/redis_test.rb
@@ -814,7 +814,7 @@ def redis.call(*attrs)
       @r.zadd "foo", 40, "s4"
 
       assert_equal 3, @r.zremrangebyrank("foo", 1, 3)
-      assert_equal ["s1"], @r.zrange("foo", 0, 4)
+      assert_equal ["s1"], @r.zrange("foo", 0, -1)
     end
 
     test "ZREMRANGEBYSCORE" do
@@ -824,7 +824,7 @@ def redis.call(*attrs)
       @r.zadd "foo", 4, "s4"
 
       assert_equal 3, @r.zremrangebyscore("foo", 2, 4)
-      assert_equal ["s1"], @r.zrange("foo", 0, 4)
+      assert_equal ["s1"], @r.zrange("foo", 0, -1)
     end
 
     test "ZUNION" do
@@ -834,6 +834,36 @@ def redis.call(*attrs)
       @r.zadd "bar", 4, "s4"
 
       assert_equal 4, @r.zunion("foobar", ["foo", "bar"])
+      assert_equal ["s1", "s2", "s3", "s4"], @r.zrange("foobar", 0, -1)
+    end
+
+    test "ZUNION with WEIGHTS" do
+      @r.zadd "foo", 1, "s1"
+      @r.zadd "foo", 3, "s3"
+      @r.zadd "bar", 20, "s2"
+      @r.zadd "bar", 40, "s4"
+
+      assert_equal 4, @r.zunion("foobar", ["foo", "bar"])
+      assert_equal ["s1", "s3", "s2", "s4"], @r.zrange("foobar", 0, -1)
+
+      assert_equal 4, @r.zunion("foobar", ["foo", "bar"], :weights => [10, 1])
+      assert_equal ["s1", "s2", "s3", "s4"], @r.zrange("foobar", 0, -1)
+    end
+
+    test "ZUNION with AGGREGATE" do
+      @r.zadd "foo", 1, "s1"
+      @r.zadd "foo", 2, "s2"
+      @r.zadd "bar", 4, "s2"
+      @r.zadd "bar", 3, "s3"
+
+      assert_equal 3, @r.zunion("foobar", ["foo", "bar"])
+      assert_equal ["s1", "s3", "s2"], @r.zrange("foobar", 0, -1)
+
+      assert_equal 3, @r.zunion("foobar", ["foo", "bar"], :aggregate => :min)
+      assert_equal ["s1", "s2", "s3"], @r.zrange("foobar", 0, -1)
+
+      assert_equal 3, @r.zunion("foobar", ["foo", "bar"], :aggregate => :max)
+      assert_equal ["s1", "s3", "s2"], @r.zrange("foobar", 0, -1)
     end
 
     test "ZINTER" do
@@ -843,7 +873,49 @@ def redis.call(*attrs)
       @r.zadd "bar", 4, "s4"
 
       assert_equal 1, @r.zinter("foobar", ["foo", "bar"])
-      assert_equal ["s1"], @r.zrange("foobar", 0, 2)
+      assert_equal ["s1"], @r.zrange("foobar", 0, -1)
+    end
+
+    test "ZINTER with WEIGHTS" do
+      @r.zadd "foo", 1, "s1"
+      @r.zadd "foo", 2, "s2"
+      @r.zadd "foo", 3, "s3"
+      @r.zadd "bar", 20, "s2"
+      @r.zadd "bar", 30, "s3"
+      @r.zadd "bar", 40, "s4"
+
+      assert_equal 2, @r.zinter("foobar", ["foo", "bar"])
+      assert_equal ["s2", "s3"], @r.zrange("foobar", 0, -1)
+
+      assert_equal 2, @r.zinter("foobar", ["foo", "bar"], :weights => [10, 1])
+      assert_equal ["s2", "s3"], @r.zrange("foobar", 0, -1)
+
+      assert_equal "40", @r.zscore("foobar", "s2")
+      assert_equal "60", @r.zscore("foobar", "s3")
+    end
+
+    test "ZINTER with AGGREGATE" do
+      @r.zadd "foo", 1, "s1"
+      @r.zadd "foo", 2, "s2"
+      @r.zadd "foo", 3, "s3"
+      @r.zadd "bar", 20, "s2"
+      @r.zadd "bar", 30, "s3"
+      @r.zadd "bar", 40, "s4"
+
+      assert_equal 2, @r.zinter("foobar", ["foo", "bar"])
+      assert_equal ["s2", "s3"], @r.zrange("foobar", 0, -1)
+      assert_equal "22", @r.zscore("foobar", "s2")
+      assert_equal "33", @r.zscore("foobar", "s3")
+
+      assert_equal 2, @r.zinter("foobar", ["foo", "bar"], :aggregate => :min)
+      assert_equal ["s2", "s3"], @r.zrange("foobar", 0, -1)
+      assert_equal "2", @r.zscore("foobar", "s2")
+      assert_equal "3", @r.zscore("foobar", "s3")
+
+      assert_equal 2, @r.zinter("foobar", ["foo", "bar"], :aggregate => :max)
+      assert_equal ["s2", "s3"], @r.zrange("foobar", 0, -1)
+      assert_equal "20", @r.zscore("foobar", "s2")
+      assert_equal "30", @r.zscore("foobar", "s3")
     end
   end