Implement zDiff and zDiffWithScores API (#367)

zio · May 6, 2021 · 2352e26 · 2352e26
1 parent e3454bf
commit 2352e26
Show file tree

Hide file tree

Showing 2 changed files with 118 additions and 0 deletions.
diff --git a/redis/src/main/scala/zio/redis/api/SortedSets.scala b/redis/src/main/scala/zio/redis/api/SortedSets.scala
@@ -146,6 +146,58 @@ trait SortedSets {
     command.run((key, range))
   }
 
+  /**
+   * Subtract multiple sorted sets and return members.
+   *
+   * @param inputKeysNum Number of input keys
+   * @param key Key of a sorted set
+   * @param keys Keys of other sorted sets
+   * @return Chunk of differences between the first and successive input sorted sets.
+   */
+  final def zDiff[K: Schema, M: Schema](
+    inputKeysNum: Long,
+    key: K,
+    keys: K*
+  ): ZIO[RedisExecutor, RedisError, Chunk[M]] = {
+    val command =
+      RedisCommand(
+        ZDiff,
+        Tuple2(
+          LongInput,
+          NonEmptyList(ArbitraryInput[K]())
+        ),
+        ChunkOutput(ArbitraryOutput[M]())
+      )
+    command.run((inputKeysNum, (key, keys.toList)))
+  }
+
+  /**
+   * Subtract multiple sorted sets and return members and their associated score.
+   *
+   * @param inputKeysNum Number of input keys
+   * @param key Key of a sorted set
+   * @param keys Keys of other sorted sets
+   * @return Chunk of differences and scores between the first and successive input sorted sets.
+   */
+  final def zDiffWithScores[K: Schema, M: Schema](
+    inputKeysNum: Long,
+    key: K,
+    keys: K*
+  ): ZIO[RedisExecutor, RedisError, Chunk[MemberScore[M]]] = {
+    val command =
+      RedisCommand(
+        ZDiff,
+        Tuple3(
+          LongInput,
+          NonEmptyList(ArbitraryInput[K]()),
+          ArbitraryInput[String]()
+        ),
+        ChunkTuple2Output(ArbitraryOutput[M](), DoubleOutput)
+          .map(_.map { case (m, s) => MemberScore(s, m) })
+      )
+    command.run((inputKeysNum, (key, keys.toList), WithScores.stringify))
+  }
+
   /**
    * Subtract multiple sorted sets and store the resulting sorted set in a destination key.
    *
@@ -810,6 +862,7 @@ private[redis] object SortedSets {
   final val ZAdd             = "ZADD"
   final val ZCard            = "ZCARD"
   final val ZCount           = "ZCOUNT"
+  final val ZDiff            = "ZDIFF"
   final val ZDiffStore       = "ZDIFFSTORE"
   final val ZIncrBy          = "ZINCRBY"
   final val ZInter           = "ZINTER"

diff --git a/redis/src/test/scala/zio/redis/SortedSetsSpec.scala b/redis/src/test/scala/zio/redis/SortedSetsSpec.scala
@@ -213,6 +213,71 @@ trait SortedSetsSpec extends BaseSpec {
           } yield assert(count)(equalTo(0L))
         }
       ),
+      suite("zDiff")(
+        testM("empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            key3 <- uuid
+            diff <- zDiff[String, String](3, key1, key2, key3)
+          } yield assert(diff)(isEmpty)
+        },
+        testM("non-empty set with empty set") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _ <- zAdd(key1)(
+                   MemberScore(1d, "a"),
+                   MemberScore(2d, "b")
+                 )
+            diff <- zDiff[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk("a", "b")))
+        },
+        testM("non-empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _ <- zAdd(key1)(
+                   MemberScore(1d, "a"),
+                   MemberScore(2d, "b"),
+                   MemberScore(3d, "c")
+                 )
+            _    <- zAdd(key2)(MemberScore(1d, "a"), MemberScore(2d, "b"))
+            diff <- zDiff[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk("c")))
+        }
+      ),
+      suite("zDiffWithScores")(
+        testM("empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            key3 <- uuid
+            diff <- zDiffWithScores[String, String](3, key1, key2, key3)
+          } yield assert(diff)(isEmpty)
+        },
+        testM("non-empty set with empty set") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _    <- zAdd(key1)(MemberScore(1d, "a"), MemberScore(2d, "b"))
+            diff <- zDiffWithScores[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk(MemberScore(1d, "a"), MemberScore(2d, "b"))))
+        },
+        testM("non-empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _ <- zAdd(key1)(
+                   MemberScore(1d, "a"),
+                   MemberScore(2d, "b"),
+                   MemberScore(3d, "c")
+                 )
+            _    <- zAdd(key2)(MemberScore(1d, "a"), MemberScore(2d, "b"))
+            diff <- zDiffWithScores[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk(MemberScore(3d, "c"))))
+        }
+      ),
       suite("zDiffStore")(
         testM("empty sets") {
           for {