From 68fe7b592fcf2c993c7dbbb4d4d35834fc7dff0b Mon Sep 17 00:00:00 2001 From: Hartley McGuire Date: Sun, 20 Jun 2021 17:39:17 -0400 Subject: [PATCH] add ZINTER command, added to Redis in 6.2 --- lib/redis.rb | 39 +++++++++++++++ lib/redis/distributed.rb | 7 +++ test/cluster_commands_on_sorted_sets_test.rb | 12 +++++ ...istributed_commands_on_sorted_sets_test.rb | 12 +++++ test/lint/sorted_sets.rb | 49 +++++++++++++++++++ 5 files changed, 119 insertions(+) diff --git a/lib/redis.rb b/lib/redis.rb index 5ce3a872e..42902486c 100644 --- a/lib/redis.rb +++ b/lib/redis.rb @@ -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] 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, 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. # diff --git a/lib/redis/distributed.rb b/lib/redis/distributed.rb index 6d1566550..7d823886a 100644 --- a/lib/redis/distributed.rb +++ b/lib/redis/distributed.rb @@ -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) diff --git a/test/cluster_commands_on_sorted_sets_test.rb b/test/cluster_commands_on_sorted_sets_test.rb index 3c37b429a..87da66ee8 100644 --- a/test/cluster_commands_on_sorted_sets_test.rb +++ b/test/cluster_commands_on_sorted_sets_test.rb @@ -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 diff --git a/test/distributed_commands_on_sorted_sets_test.rb b/test/distributed_commands_on_sorted_sets_test.rb index 2e628364c..cd220bcf0 100644 --- a/test/distributed_commands_on_sorted_sets_test.rb +++ b/test/distributed_commands_on_sorted_sets_test.rb @@ -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 diff --git a/test/lint/sorted_sets.rb b/test/lint/sorted_sets.rb index 40a6f02ea..288155cd7 100644 --- a/test/lint/sorted_sets.rb +++ b/test/lint/sorted_sets.rb @@ -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'