[Redis 6.2] Add new options to ZRANGE and implement new ZRANGESTORE command

lib/redis.rb

@@ -1974,9 +1974,9 @@ def zrandmember(key, count = nil, withscores: false, with_scores: withscores)
     end
   end
 
-  # Return a range of members in a sorted set, by index.
+  # Return a range of members in a sorted set, by index, score or lexicographical ordering.
   #
-  # @example Retrieve all members from a sorted set
+  # @example Retrieve all members from a sorted set, by index
   #   redis.zrange("zset", 0, -1)
   #     # => ["a", "b"]
   # @example Retrieve all members and their scores from a sorted set
@@ -1987,14 +1987,38 @@ def zrandmember(key, count = nil, withscores: false, with_scores: withscores)
   # @param [Integer] start start index
   # @param [Integer] stop stop index
   # @param [Hash] options
+  #   - `:by_score => false`: return members by score
+  #   - `:by_lex => false`: return members by lexicographical ordering
+  #   - `:rev => false`: reverse the ordering, from highest to lowest
+  #   - `:limit => [offset, count]`: skip `offset` members, return a maximum of
+  #   `count` members
   #   - `: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 zrange(key, start, stop, withscores: false, with_scores: withscores)
+  def zrange(key, start, stop, byscore: false, by_score: byscore, bylex: false, by_lex: bylex,
+             rev: false, limit: nil, withscores: false, with_scores: withscores)
+
+    if by_score && by_lex
+      raise ArgumentError, "only one of :by_score or :by_lex can be specified"
+    end
+
     args = [:zrange, key, start, stop]
 
+    if by_score
+      args << "BYSCORE"
+    elsif by_lex
+      args << "BYLEX"
+    end
+
+    args << "REV" if rev
+
+    if limit
+      args << "LIMIT"
+      args.concat(limit)
+    end
+
     if with_scores
       args << "WITHSCORES"
       block = FloatifyPairs
@@ -2005,6 +2029,44 @@ def zrange(key, start, stop, withscores: false, with_scores: withscores)
     end
   end
 
+  # Select a range of members in a sorted set, by index, score or lexicographical ordering
+  # and store the resulting sorted set in a new key.
+  #
+  # @example
+  #   redis.zadd("foo", [[1.0, "s1"], [2.0, "s2"], [3.0, "s3"]])
+  #   redis.zrangestore("bar", "foo", 0, 1)
+  #     # => 2
+  #   redis.zrange("bar", 0, -1)
+  #     # => ["s1", "s2"]
+  #
+  # @return [Integer] the number of elements in the resulting sorted set
+  # @see #zrange
+  def zrangestore(dest_key, src_key, start, stop, byscore: false, by_score: byscore,
+                  bylex: false, by_lex: bylex, rev: false, limit: nil)
+    if by_score && by_lex
+      raise ArgumentError, "only one of :by_score or :by_lex can be specified"
+    end
+
+    args = [:zrangestore, dest_key, src_key, start, stop]
+
+    if by_score
+      args << "BYSCORE"
+    elsif by_lex
+      args << "BYLEX"
+    end
+
+    args << "REV" if rev
+
+    if limit
+      args << "LIMIT"
+      args.concat(limit)
+    end
+
+    synchronize do |client|
+      client.call(args)
+    end
+  end
+
   # Return a range of members in a sorted set, by index, with scores ordered
   # from high to low.
   #

lib/redis/distributed.rb

@@ -673,11 +673,19 @@ def zmscore(key, *members)
       node_for(key).zmscore(key, *members)
     end
 
-    # Return a range of members in a sorted set, by index.
+    # Return a range of members in a sorted set, by index, score or lexicographical ordering.
     def zrange(key, start, stop, **options)
       node_for(key).zrange(key, start, stop, **options)
     end
 
+    # Select a range of members in a sorted set, by index, score or lexicographical ordering
+    # and store the resulting sorted set in a new key.
+    def zrangestore(dest_key, src_key, start, stop, **options)
+      ensure_same_node(:zrangestore, [dest_key, src_key]) do |node|
+        node.zrangestore(dest_key, src_key, start, stop, **options)
+      end
+    end
+
     # Return a range of members in a sorted set, by index, with scores ordered
     # from high to low.
     def zrevrange(key, start, stop, **options)

test/cluster_commands_on_sorted_sets_test.rb

@@ -9,6 +9,10 @@ class TestClusterCommandsOnSortedSets < Minitest::Test
   include Helper::Cluster
   include Lint::SortedSets
 
+  def test_zrangestore
+    assert_raises(Redis::CommandError) { super }
+  end
+
   def test_zinter
     assert_raises(Redis::CommandError) { super }
   end

test/distributed_commands_on_sorted_sets_test.rb

@@ -7,6 +7,10 @@ class TestDistributedCommandsOnSortedSets < Minitest::Test
   include Helper::Distributed
   include Lint::SortedSets
 
+  def test_zrangestore
+    assert_raises(Redis::Distributed::CannotDistribute) { super }
+  end
+
   def test_zinter
     assert_raises(Redis::Distributed::CannotDistribute) { super }
   end

test/lint/sorted_sets.rb

@@ -231,6 +231,45 @@ def test_zrange
       assert_equal [["s1", -Float::INFINITY], ["s2", +Float::INFINITY]], r.zrange("bar", 0, 1, withscores: true)
     end
 
+    def test_zrange_with_byscore
+      target_version("6.2") do
+        r.zadd "foo", 1, "s1"
+        r.zadd "foo", 2, "s2"
+        r.zadd "foo", 3, "s3"
+
+        assert_equal ["s2", "s3"], r.zrange("foo", 2, 3, byscore: true)
+        assert_equal ["s2", "s1"], r.zrange("foo", 2, 1, byscore: true, rev: true)
+      end
+    end
+
+    def test_zrange_with_bylex
+      target_version("6.2") do
+        r.zadd "foo", 0, "aaren"
+        r.zadd "foo", 0, "abagael"
+        r.zadd "foo", 0, "abby"
+        r.zadd "foo", 0, "abbygail"
+
+        assert_equal %w[aaren abagael abby abbygail], r.zrange("foo", "[a", "[a\xff", bylex: true)
+        assert_equal %w[aaren abagael], r.zrange("foo", "[a", "[a\xff", bylex: true, limit: [0, 2])
+        assert_equal %w[abby abbygail], r.zrange("foo", "(abb", "(abb\xff", bylex: true)
+        assert_equal %w[abbygail], r.zrange("foo", "(abby", "(abby\xff", bylex: true)
+      end
+    end
+
+    def test_zrangestore
+      target_version("6.2") do
+        r.zadd "foo", 1, "s1"
+        r.zadd "foo", 2, "s2"
+        r.zadd "foo", 3, "s3"
+
+        assert_equal 2, r.zrangestore("bar", "foo", 0, 1)
+        assert_equal ["s1", "s2"], r.zrange("bar", 0, -1)
+
+        assert_equal 2, r.zrangestore("baz", "foo", 2, 3, by_score: true)
+        assert_equal ["s2", "s3"], r.zrange("baz", 0, -1)
+      end
+    end
+
     def test_zrevrange
       r.zadd "foo", 1, "s1"
       r.zadd "foo", 2, "s2"