Add support for weights in zunionstore and zinterstore (redis-rs#641)

redis/src/commands.rs

@@ -849,6 +849,30 @@ implement_commands! {
         cmd("ZINTERSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("AGGREGATE").arg("MAX")
     }
 
+    /// [`Commands::zinterstore`], but with the ability to specify a
+    /// multiplication factor for each sorted set by pairing one with each key
+    /// in a tuple.
+    fn zinterstore_weights<K: ToRedisArgs, W: ToRedisArgs>(dstkey: K, keys: &'a [(K, W)]) {
+        let (keys, weights): (Vec<&K>, Vec<&W>) = keys.iter().map(|(key, weight)| (key, weight)).unzip();
+        cmd("ZINTERSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("WEIGHTS").arg(weights)
+    }
+
+    /// [`Commands::zinterstore_min`], but with the ability to specify a
+    /// multiplication factor for each sorted set by pairing one with each key
+    /// in a tuple.
+    fn zinterstore_min_weights<K: ToRedisArgs, W: ToRedisArgs>(dstkey: K, keys: &'a [(K, W)]) {
+        let (keys, weights): (Vec<&K>, Vec<&W>) = keys.iter().map(|(key, weight)| (key, weight)).unzip();
+        cmd("ZINTERSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("AGGREGATE").arg("MIN").arg("WEIGHTS").arg(weights)
+    }
+
+    /// [`Commands::zinterstore_max`], but with the ability to specify a
+    /// multiplication factor for each sorted set by pairing one with each key
+    /// in a tuple.
+    fn zinterstore_max_weights<K: ToRedisArgs, W: ToRedisArgs>(dstkey: K, keys: &'a [(K, W)]) {
+        let (keys, weights): (Vec<&K>, Vec<&W>) = keys.iter().map(|(key, weight)| (key, weight)).unzip();
+        cmd("ZINTERSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("AGGREGATE").arg("MAX").arg("WEIGHTS").arg(weights)
+    }
+
     /// Count the number of members in a sorted set between a given lexicographical range.
     fn zlexcount<K: ToRedisArgs, L: ToRedisArgs>(key: K, min: L, max: L) {
         cmd("ZLEXCOUNT").arg(key).arg(min).arg(max)
@@ -1036,6 +1060,30 @@ implement_commands! {
         cmd("ZUNIONSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("AGGREGATE").arg("MAX")
     }
 
+    /// [`Commands::zunionstore`], but with the ability to specify a
+    /// multiplication factor for each sorted set by pairing one with each key
+    /// in a tuple.
+    fn zunionstore_weights<K: ToRedisArgs, W: ToRedisArgs>(dstkey: K, keys: &'a [(K, W)]) {
+        let (keys, weights): (Vec<&K>, Vec<&W>) = keys.iter().map(|(key, weight)| (key, weight)).unzip();
+        cmd("ZUNIONSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("WEIGHTS").arg(weights)
+    }
+
+    /// [`Commands::zunionstore_min`], but with the ability to specify a
+    /// multiplication factor for each sorted set by pairing one with each key
+    /// in a tuple.
+    fn zunionstore_min_weights<K: ToRedisArgs, W: ToRedisArgs>(dstkey: K, keys: &'a [(K, W)]) {
+        let (keys, weights): (Vec<&K>, Vec<&W>) = keys.iter().map(|(key, weight)| (key, weight)).unzip();
+        cmd("ZUNIONSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("AGGREGATE").arg("MIN").arg("WEIGHTS").arg(weights)
+    }
+
+    /// [`Commands::zunionstore_max`], but with the ability to specify a
+    /// multiplication factor for each sorted set by pairing one with each key
+    /// in a tuple.
+    fn zunionstore_max_weights<K: ToRedisArgs, W: ToRedisArgs>(dstkey: K, keys: &'a [(K, W)]) {
+        let (keys, weights): (Vec<&K>, Vec<&W>) = keys.iter().map(|(key, weight)| (key, weight)).unzip();
+        cmd("ZUNIONSTORE").arg(dstkey).arg(keys.len()).arg(keys).arg("AGGREGATE").arg("MAX").arg("WEIGHTS").arg(weights)
+    }
+
     // hyperloglog commands
 
     /// Adds the specified elements to the specified HyperLogLog.

redis/tests/test_basic.rs

@@ -921,6 +921,119 @@ fn test_redis_server_down() {
     assert!(!con.is_open());
 }
 
+#[test]
+fn test_zinterstore_weights() {
+    let ctx = TestContext::new();
+    let mut con = ctx.connection();
+
+    let _: () = con
+        .zadd_multiple("zset1", &[(1, "one"), (2, "two"), (4, "four")])
+        .unwrap();
+    let _: () = con
+        .zadd_multiple("zset2", &[(1, "one"), (2, "two"), (3, "three")])
+        .unwrap();
+
+    // zinterstore_weights
+    assert_eq!(
+        con.zinterstore_weights("out", &[("zset1", 2), ("zset2", 3)]),
+        Ok(2)
+    );
+
+    assert_eq!(
+        con.zrange_withscores("out", 0, -1),
+        Ok(vec![
+            ("one".to_string(), "5".to_string()),
+            ("two".to_string(), "10".to_string())
+        ])
+    );
+
+    // zinterstore_min_weights
+    assert_eq!(
+        con.zinterstore_min_weights("out", &[("zset1", 2), ("zset2", 3)]),
+        Ok(2)
+    );
+
+    assert_eq!(
+        con.zrange_withscores("out", 0, -1),
+        Ok(vec![
+            ("one".to_string(), "2".to_string()),
+            ("two".to_string(), "4".to_string()),
+        ])
+    );
+
+    // zinterstore_max_weights
+    assert_eq!(
+        con.zinterstore_max_weights("out", &[("zset1", 2), ("zset2", 3)]),
+        Ok(2)
+    );
+
+    assert_eq!(
+        con.zrange_withscores("out", 0, -1),
+        Ok(vec![
+            ("one".to_string(), "3".to_string()),
+            ("two".to_string(), "6".to_string()),
+        ])
+    );
+}
+
+#[test]
+fn test_zunionstore_weights() {
+    let ctx = TestContext::new();
+    let mut con = ctx.connection();
+
+    let _: () = con
+        .zadd_multiple("zset1", &[(1, "one"), (2, "two")])
+        .unwrap();
+    let _: () = con
+        .zadd_multiple("zset2", &[(1, "one"), (2, "two"), (3, "three")])
+        .unwrap();
+
+    // zunionstore_weights
+    assert_eq!(
+        con.zunionstore_weights("out", &[("zset1", 2), ("zset2", 3)]),
+        Ok(3)
+    );
+
+    assert_eq!(
+        con.zrange_withscores("out", 0, -1),
+        Ok(vec![
+            ("one".to_string(), "5".to_string()),
+            ("three".to_string(), "9".to_string()),
+            ("two".to_string(), "10".to_string())
+        ])
+    );
+
+    // zunionstore_min_weights
+    assert_eq!(
+        con.zunionstore_min_weights("out", &[("zset1", 2), ("zset2", 3)]),
+        Ok(3)
+    );
+
+    assert_eq!(
+        con.zrange_withscores("out", 0, -1),
+        Ok(vec![
+            ("one".to_string(), "2".to_string()),
+            ("two".to_string(), "4".to_string()),
+            ("three".to_string(), "9".to_string())
+        ])
+    );
+
+    // zunionstore_max_weights
+    assert_eq!(
+        con.zunionstore_max_weights("out", &[("zset1", 2), ("zset2", 3)]),
+        Ok(3)
+    );
+
+    assert_eq!(
+        con.zrange_withscores("out", 0, -1),
+        Ok(vec![
+            ("one".to_string(), "3".to_string()),
+            ("two".to_string(), "6".to_string()),
+            ("three".to_string(), "9".to_string())
+        ])
+    );
+}
+
 #[test]
 fn test_zrembylex() {
     let ctx = TestContext::new();