Skip to content

Commit

Permalink
Merge pull request #995 from SkipKayhil/add-zinter
Browse files Browse the repository at this point in the history
add ZINTER command, added to Redis in 6.2
  • Loading branch information
byroot committed Jun 21, 2021
2 parents f759203 + 68fe7b5 commit 7bb6b63
Show file tree
Hide file tree
Showing 5 changed files with 119 additions and 0 deletions.
39 changes: 39 additions & 0 deletions lib/redis.rb
Expand Up @@ -2062,6 +2062,45 @@ def zcount(key, min, max)
end
end

# Return the intersection of multiple sorted sets
#
# @example Retrieve the intersection of `2*zsetA` and `1*zsetB`
# redis.zinter("zsetA", "zsetB", :weights => [2.0, 1.0])
# # => ["v1", "v2"]
# @example Retrieve the intersection of `2*zsetA` and `1*zsetB`, and their scores
# redis.zinter("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 intersect
# @param [Hash] options
# - `:weights => [Float, 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 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
end

# Intersect multiple sorted sets and store the resulting sorted set in a new
# key.
#
Expand Down
7 changes: 7 additions & 0 deletions lib/redis/distributed.rb
Expand Up @@ -674,6 +674,13 @@ def zcount(key, min, max)
node_for(key).zcount(key, min, max)
end

# Get the intersection of multiple sorted sets
def zinter(*keys, **options)
ensure_same_node(:zinter, keys) do |node|
node.zinter(*keys, **options)
end
end

# Intersect multiple sorted sets and store the resulting sorted set in a new
# key.
def zinterstore(destination, keys, **options)
Expand Down
12 changes: 12 additions & 0 deletions test/cluster_commands_on_sorted_sets_test.rb
Expand Up @@ -9,6 +9,18 @@ class TestClusterCommandsOnSortedSets < Minitest::Test
include Helper::Cluster
include Lint::SortedSets

def test_zinter
assert_raises(Redis::CommandError) { super }
end

def test_zinter_with_aggregate
assert_raises(Redis::CommandError) { super }
end

def test_zinter_with_weights
assert_raises(Redis::CommandError) { super }
end

def test_zinterstore
assert_raises(Redis::CommandError) { super }
end
Expand Down
12 changes: 12 additions & 0 deletions test/distributed_commands_on_sorted_sets_test.rb
Expand Up @@ -7,6 +7,18 @@ class TestDistributedCommandsOnSortedSets < Minitest::Test
include Helper::Distributed
include Lint::SortedSets

def test_zinter
assert_raises(Redis::Distributed::CannotDistribute) { super }
end

def test_zinter_with_aggregate
assert_raises(Redis::Distributed::CannotDistribute) { super }
end

def test_zinter_with_weights
assert_raises(Redis::Distributed::CannotDistribute) { super }
end

def test_zinterstore
assert_raises(Redis::Distributed::CannotDistribute) { super }
end
Expand Down
49 changes: 49 additions & 0 deletions test/lint/sorted_sets.rb
Expand Up @@ -432,6 +432,55 @@ def test_zunionstore_expand
assert_equal 5, r.zunionstore('{1}baz', %w[{1}foo {1}bar])
end

def test_zinter
target_version("6.2") do
r.zadd 'foo', 1, 's1'
r.zadd 'bar', 2, 's1'
r.zadd 'foo', 3, 's3'
r.zadd 'bar', 4, 's4'

assert_equal ['s1'], r.zinter('foo', 'bar')
assert_equal [['s1', 3.0]], r.zinter('foo', 'bar', with_scores: true)
end
end

def test_zinter_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', 40, 's4'

assert_equal %w[s2 s3], r.zinter('foo', 'bar')
assert_equal [['s2', 22.0], ['s3', 33.0]], r.zinter('foo', 'bar', with_scores: true)

assert_equal %w[s2 s3], r.zinter('foo', 'bar', weights: [10, 1])
assert_equal [['s2', 40.0], ['s3', 60.0]], r.zinter('foo', 'bar', weights: [10, 1], with_scores: true)
end
end

def test_zinter_with_aggregate
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', 40, 's4'

assert_equal %w[s2 s3], r.zinter('foo', 'bar')
assert_equal [['s2', 22.0], ['s3', 33.0]], r.zinter('foo', 'bar', with_scores: true)

assert_equal %w[s2 s3], r.zinter('foo', 'bar', aggregate: :min)
assert_equal [['s2', 2.0], ['s3', 3.0]], r.zinter('foo', 'bar', aggregate: :min, with_scores: true)

assert_equal %w[s2 s3], r.zinter('foo', 'bar', aggregate: :max)
assert_equal [['s2', 20.0], ['s3', 30.0]], r.zinter('foo', 'bar', aggregate: :max, with_scores: true)
end
end

def test_zinterstore
r.zadd 'foo', 1, 's1'
r.zadd 'bar', 2, 's1'
Expand Down

0 comments on commit 7bb6b63

Please sign in to comment.