[Redis 6.2] Add ZUNION command

lib/redis.rb

@@ -2260,25 +2260,10 @@ def zcount(key, min, max)
   # @return [Array<String>, Array<[String, Float]>]
   #   - when `:with_scores` is not specified, an array of members
   #   - when `:with_scores` is specified, an array with `[member, score]` pairs
-  def zinter(*keys, weights: nil, aggregate: nil, with_scores: false)
-    args = [:zinter, keys.size, *keys]
-
-    if weights
-      args << "WEIGHTS"
-      args.concat(weights)
-    end
-
-    args << "AGGREGATE" << aggregate if aggregate
-
-    if with_scores
-      args << "WITHSCORES"
-      block = FloatifyPairs
-    end
-
-    synchronize do |client|
-      client.call(args, &block)
-    end
+  def zinter(*args)
+    _zsets_operation(:zinter, *args)
   end
+  ruby2_keywords(:zinter) if respond_to?(:ruby2_keywords, true)
 
   # Intersect multiple sorted sets and store the resulting sorted set in a new
   # key.
@@ -2290,9 +2275,9 @@ def zinter(*keys, weights: nil, aggregate: nil, with_scores: false)
   # @param [String] destination destination key
   # @param [Array<String>] keys source keys
   # @param [Hash] options
-  #   - `:weights => [Float, Float, ...]`: weights to associate with source
+  #   - `:weights => [Array<Float>]`: weights to associate with source
   #   sorted sets
-  #   - `:aggregate => String`: aggregate function to use (sum, min, max, ...)
+  #   - `:aggregate => String`: aggregate function to use (sum, min, max)
   # @return [Integer] number of elements in the resulting sorted set
   def zinterstore(destination, keys, weights: nil, aggregate: nil)
     args = [:zinterstore, destination, keys.size, *keys]
@@ -2309,6 +2294,30 @@ def zinterstore(destination, keys, weights: nil, aggregate: nil)
     end
   end
 
+  # Return the union of multiple sorted sets
+  #
+  # @example Retrieve the union of `2*zsetA` and `1*zsetB`
+  #   redis.zunion("zsetA", "zsetB", :weights => [2.0, 1.0])
+  #     # => ["v1", "v2"]
+  # @example Retrieve the union of `2*zsetA` and `1*zsetB`, and their scores
+  #   redis.zunion("zsetA", "zsetB", :weights => [2.0, 1.0], :with_scores => true)
+  #     # => [["v1", 3.0], ["v2", 6.0]]
+  #
+  # @param [String, Array<String>] keys one or more keys to union
+  # @param [Hash] options
+  #   - `:weights => [Array<Float>]`: weights to associate with source
+  #   sorted sets
+  #   - `:aggregate => String`: aggregate function to use (sum, min, max)
+  #   - `:with_scores => true`: include scores in output
+  #
+  # @return [Array<String>, Array<[String, Float]>]
+  #   - when `:with_scores` is not specified, an array of members
+  #   - when `:with_scores` is specified, an array with `[member, score]` pairs
+  def zunion(*args)
+    _zsets_operation(:zunion, *args)
+  end
+  ruby2_keywords(:zunion) if respond_to?(:ruby2_keywords, true)
+
   # Add multiple sorted sets and store the resulting sorted set in a new key.
   #
   # @example Compute the union of `2*zsetA` with `1*zsetB`, summing their scores
@@ -3871,6 +3880,26 @@ def _normalize_move_wheres(where_source, where_destination)
 
     [where_source, where_destination]
   end
+
+  def _zsets_operation(cmd, *keys, weights: nil, aggregate: nil, with_scores: false)
+    command = [cmd, keys.size, *keys]
+
+    if weights
+      command << "WEIGHTS"
+      command.concat(weights)
+    end
+
+    command << "AGGREGATE" << aggregate if aggregate
+
+    if with_scores
+      command << "WITHSCORES"
+      block = FloatifyPairs
+    end
+
+    synchronize do |client|
+      client.call(command, &block)
+    end
+  end
 end
 
 require_relative "redis/version"

lib/redis/distributed.rb

@@ -729,6 +729,13 @@ def zinterstore(destination, keys, **options)
       end
     end
 
+    # Return the union of multiple sorted sets.
+    def zunion(*keys, **options)
+      ensure_same_node(:zunion, keys) do |node|
+        node.zunion(*keys, **options)
+      end
+    end
+
     # Add multiple sorted sets and store the resulting sorted set in a new key.
     def zunionstore(destination, keys, **options)
       ensure_same_node(:zunionstore, [destination] + keys) do |node|

test/cluster_commands_on_sorted_sets_test.rb

@@ -33,6 +33,18 @@ def test_zinterstore_with_weights
     assert_raises(Redis::CommandError) { super }
   end
 
+  def test_zunion
+    assert_raises(Redis::CommandError) { super }
+  end
+
+  def test_zunion_with_aggregate
+    assert_raises(Redis::CommandError) { super }
+  end
+
+  def test_zunion_with_weights
+    assert_raises(Redis::CommandError) { super }
+  end
+
   def test_zunionstore
     assert_raises(Redis::CommandError) { super }
   end

test/distributed_commands_on_sorted_sets_test.rb

@@ -59,6 +59,18 @@ def test_zscan
     # Not implemented yet
   end
 
+  def test_zunion
+    assert_raises(Redis::Distributed::CannotDistribute) { super }
+  end
+
+  def test_zunion_with_aggregate
+    assert_raises(Redis::Distributed::CannotDistribute) { super }
+  end
+
+  def test_zunion_with_weights
+    assert_raises(Redis::Distributed::CannotDistribute) { super }
+  end
+
   def test_zunionstore
     assert_raises(Redis::Distributed::CannotDistribute) { super }
   end

test/lint/sorted_sets.rb

@@ -480,6 +480,52 @@ def test_zcount
       assert_equal 2, r.zcount('foo', 2, 3)
     end
 
+    def test_zunion
+      target_version("6.2") do
+        r.zadd 'foo', 1, 's1'
+        r.zadd 'foo', 2, 's2'
+        r.zadd 'bar', 3, 's1'
+        r.zadd 'bar', 5, 's3'
+
+        assert_equal %w[s2 s1 s3], r.zunion('foo', 'bar')
+        assert_equal [['s2', 2.0], ['s1', 4.0], ['s3', 5.0]], r.zunion('foo', 'bar', with_scores: true)
+      end
+    end
+
+    def test_zunion_with_weights
+      target_version("6.2") 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', 100, 's4'
+
+        assert_equal %w[s1 s2 s3 s4], r.zunion('foo', 'bar')
+        assert_equal [['s1', 1.0], ['s2', 22.0], ['s3', 33.0], ['s4', 100.0]], r.zunion('foo', 'bar', with_scores: true)
+
+        assert_equal %w[s1 s2 s3 s4], r.zunion('foo', 'bar', weights: [10, 1])
+        assert_equal [['s1', 10.0], ['s2', 40.0], ['s3', 60.0], ['s4', 100.0]], r.zunion('foo', 'bar', weights: [10, 1], with_scores: true)
+      end
+    end
+
+    def test_zunion_with_aggregate
+      target_version("6.2") do
+        r.zadd 'foo', 1, 's1'
+        r.zadd 'foo', 20, 's2'
+        r.zadd 'foo', 3, 's3'
+        r.zadd 'bar', 2, 's2'
+        r.zadd 'bar', 30, 's3'
+        r.zadd 'bar', 100, 's4'
+
+        assert_equal %w[s1 s2 s3 s4], r.zunion('foo', 'bar', aggregate: :min)
+        assert_equal [['s1', 1.0], ['s2', 2.0], ['s3', 3.0], ['s4', 100.0]], r.zunion('foo', 'bar', aggregate: :min, with_scores: true)
+
+        assert_equal %w[s1 s2 s3 s4], r.zunion('foo', 'bar', aggregate: :max)
+        assert_equal [['s1', 1.0], ['s2', 20.0], ['s3', 30.0], ['s4', 100.0]], r.zunion('foo', 'bar', aggregate: :max, with_scores: true)
+      end
+    end
+
     def test_zunionstore
       r.zadd 'foo', 1, 's1'
       r.zadd 'bar', 2, 's2'