diff --git a/src/main/java/org/springframework/data/redis/core/BoundZSetOperations.java b/src/main/java/org/springframework/data/redis/core/BoundZSetOperations.java index 20b2ef854..737b735d6 100644 --- a/src/main/java/org/springframework/data/redis/core/BoundZSetOperations.java +++ b/src/main/java/org/springframework/data/redis/core/BoundZSetOperations.java @@ -18,6 +18,7 @@ package org.springframework.data.redis.core; import java.util.Collection; import java.util.Set; +import org.springframework.data.redis.connection.RedisZSetCommands; import org.springframework.data.redis.connection.RedisZSetCommands.Limit; import org.springframework.data.redis.connection.RedisZSetCommands.Range; import org.springframework.data.redis.connection.RedisZSetCommands.Tuple; @@ -30,6 +31,7 @@ import org.springframework.lang.Nullable; * @author Costin Leau * @author Christoph Strobl * @author Mark Paluch + * @author wongoo */ public interface BoundZSetOperations extends BoundKeyOperations { @@ -261,6 +263,27 @@ public interface BoundZSetOperations extends BoundKeyOperations { */ void unionAndStore(Collection otherKeys, K destKey); + /** + * Union sorted sets at the bound key and {@code otherKeys} and store result in destination {@code destKey}. + * + * @param otherKeys must not be {@literal null}. + * @param destKey must not be {@literal null}. + * @param aggregate must not be {@literal null}. + * @see Redis Documentation: ZUNIONSTORE + */ + void unionAndStore(Collection otherKeys, K destKey, RedisZSetCommands.Aggregate aggregate); + + /** + * Union sorted sets at the bound key and {@code otherKeys} and store result in destination {@code destKey}. + * + * @param otherKeys must not be {@literal null}. + * @param weights must not be {@literal null}. + * @param destKey must not be {@literal null}. + * @param aggregate must not be {@literal null}. + * @see Redis Documentation: ZUNIONSTORE + */ + void unionAndStore(Collection otherKeys, int[] weights, K destKey, RedisZSetCommands.Aggregate aggregate); + /** * Intersect sorted sets at the bound key and {@code otherKey} and store result in destination {@code destKey}. * diff --git a/src/main/java/org/springframework/data/redis/core/DefaultBoundZSetOperations.java b/src/main/java/org/springframework/data/redis/core/DefaultBoundZSetOperations.java index 97b91846f..31233cf97 100644 --- a/src/main/java/org/springframework/data/redis/core/DefaultBoundZSetOperations.java +++ b/src/main/java/org/springframework/data/redis/core/DefaultBoundZSetOperations.java @@ -20,6 +20,7 @@ import java.util.Collection; import java.util.Set; import org.springframework.data.redis.connection.DataType; +import org.springframework.data.redis.connection.RedisZSetCommands; import org.springframework.data.redis.connection.RedisZSetCommands.Limit; import org.springframework.data.redis.connection.RedisZSetCommands.Range; import org.springframework.data.redis.core.ZSetOperations.TypedTuple; @@ -30,6 +31,7 @@ import org.springframework.data.redis.core.ZSetOperations.TypedTuple; * @author Costin Leau * @author Christoph Strobl * @author Mark Paluch + * @author wongoo */ class DefaultBoundZSetOperations extends DefaultBoundKeyOperations implements BoundZSetOperations { @@ -290,6 +292,24 @@ class DefaultBoundZSetOperations extends DefaultBoundKeyOperations impl ops.unionAndStore(getKey(), otherKeys, destKey); } + /* + * (non-Javadoc) + * @see org.springframework.data.redis.core.BoundZSetOperations#unionAndStore(Collection, Object, RedisZSetCommands.Aggregate) + */ + @Override + public void unionAndStore(Collection otherKeys, K destKey, RedisZSetCommands.Aggregate aggregate) { + ops.unionAndStore(getKey(), otherKeys, destKey, aggregate); + } + + /* + * (non-Javadoc) + * @see org.springframework.data.redis.core.BoundZSetOperations#unionAndStore(Collection, int[], Object, RedisZSetCommands.Aggregate) + */ + @Override + public void unionAndStore(Collection otherKeys, int[] weights, K destKey, RedisZSetCommands.Aggregate aggregate) { + ops.unionAndStore(getKey(), otherKeys, weights, destKey, aggregate); + } + /* * (non-Javadoc) * @see org.springframework.data.redis.core.BoundKeyOperations#getType() diff --git a/src/main/java/org/springframework/data/redis/core/DefaultZSetOperations.java b/src/main/java/org/springframework/data/redis/core/DefaultZSetOperations.java index abecd5764..671d957a6 100644 --- a/src/main/java/org/springframework/data/redis/core/DefaultZSetOperations.java +++ b/src/main/java/org/springframework/data/redis/core/DefaultZSetOperations.java @@ -19,6 +19,7 @@ import java.util.Collection; import java.util.Collections; import java.util.Set; +import org.springframework.data.redis.connection.RedisZSetCommands; import org.springframework.data.redis.connection.RedisZSetCommands.Limit; import org.springframework.data.redis.connection.RedisZSetCommands.Range; import org.springframework.data.redis.connection.RedisZSetCommands.Tuple; @@ -31,6 +32,7 @@ import org.springframework.data.redis.connection.RedisZSetCommands.Tuple; * @author Thomas Darimont * @author David Liu * @author Mark Paluch + * @author wongoo */ class DefaultZSetOperations extends AbstractOperations implements ZSetOperations { @@ -400,12 +402,35 @@ class DefaultZSetOperations extends AbstractOperations implements ZS */ @Override public Long unionAndStore(K key, Collection otherKeys, K destKey) { - byte[][] rawKeys = rawKeys(key, otherKeys); byte[] rawDestKey = rawKey(destKey); return execute(connection -> connection.zUnionStore(rawDestKey, rawKeys), true); } + /* + * (non-Javadoc) + * @see org.springframework.data.redis.core.ZSetOperations#unionAndStore(java.lang.Object, java.util.Collection, java.lang.Object, org.springframework.data.redis.connection.RedisZSetCommands.Aggregate) + */ + @Override + public Long unionAndStore(K key, Collection otherKeys, K destKey, RedisZSetCommands.Aggregate aggregate) { + int weights[] = new int[otherKeys.size() + (key != null ? 1 : 0)]; + for (int i = 0; i < weights.length; i++) { + weights[i] = 1; + } + return unionAndStore(key, otherKeys, weights, destKey, aggregate); + } + + /* + * (non-Javadoc) + * @see org.springframework.data.redis.core.ZSetOperations#unionAndStore(java.lang.Object, java.util.Collection, int[], java.lang.Object, org.springframework.data.redis.connection.RedisZSetCommands.Aggregate) + */ + @Override + public Long unionAndStore(K key, Collection otherKeys, int[] weights, K destKey, RedisZSetCommands.Aggregate aggregate) { + byte[][] rawKeys = rawKeys(key, otherKeys); + byte[] rawDestKey = rawKey(destKey); + return execute(connection -> connection.zUnionStore(rawDestKey, aggregate, weights, rawKeys), true); + } + /* * (non-Javadoc) * @see org.springframework.data.redis.core.ZSetOperations#scan(java.lang.Object, org.springframework.data.redis.core.ScanOptions) diff --git a/src/main/java/org/springframework/data/redis/core/ZSetOperations.java b/src/main/java/org/springframework/data/redis/core/ZSetOperations.java index 82a9bbe40..22ea93d58 100644 --- a/src/main/java/org/springframework/data/redis/core/ZSetOperations.java +++ b/src/main/java/org/springframework/data/redis/core/ZSetOperations.java @@ -18,6 +18,7 @@ package org.springframework.data.redis.core; import java.util.Collection; import java.util.Set; +import org.springframework.data.redis.connection.RedisZSetCommands; import org.springframework.data.redis.connection.RedisZSetCommands.Limit; import org.springframework.data.redis.connection.RedisZSetCommands.Range; import org.springframework.data.redis.connection.RedisZSetCommands.Tuple; @@ -30,6 +31,7 @@ import org.springframework.lang.Nullable; * @author Christoph Strobl * @author Mark Paluch * @author Rosty Kerei + * @author wongoo */ public interface ZSetOperations { @@ -363,6 +365,33 @@ public interface ZSetOperations { @Nullable Long unionAndStore(K key, Collection otherKeys, K destKey); + /** + * Union sorted sets at {@code key} and {@code otherKeys} and store result in destination {@code destKey}. + * + * @param key must not be {@literal null}. + * @param otherKeys must not be {@literal null}. + * @param destKey must not be {@literal null}. + * @param aggregate must not be {@literal null}. + * @return {@literal null} when used in pipeline / transaction. + * @see Redis Documentation: ZUNIONSTORE + */ + @Nullable + Long unionAndStore(K key, Collection otherKeys, K destKey, RedisZSetCommands.Aggregate aggregate); + + /** + * Union sorted sets at {@code key} and {@code otherKeys} and store result in destination {@code destKey}. + * + * @param key must not be {@literal null}. + * @param otherKeys must not be {@literal null}. + * @param weights must not be {@literal null}. + * @param destKey must not be {@literal null}. + * @param aggregate must not be {@literal null}. + * @return {@literal null} when used in pipeline / transaction. + * @see Redis Documentation: ZUNIONSTORE + */ + @Nullable + Long unionAndStore(K key, Collection otherKeys, int[] weights, K destKey, RedisZSetCommands.Aggregate aggregate); + /** * Intersect sorted sets at {@code key} and {@code otherKey} and store result in destination {@code destKey}. * diff --git a/src/test/java/org/springframework/data/redis/core/DefaultZSetOperationsTests.java b/src/test/java/org/springframework/data/redis/core/DefaultZSetOperationsTests.java index cf0de9e24..28a4cf96e 100644 --- a/src/test/java/org/springframework/data/redis/core/DefaultZSetOperationsTests.java +++ b/src/test/java/org/springframework/data/redis/core/DefaultZSetOperationsTests.java @@ -23,10 +23,12 @@ import static org.junit.Assume.*; import static org.springframework.data.redis.matcher.RedisTestMatchers.*; import java.io.IOException; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.HashSet; import java.util.LinkedHashSet; +import java.util.List; import java.util.Set; import org.junit.After; @@ -37,6 +39,7 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; import org.junit.runners.Parameterized.Parameters; + import org.springframework.data.redis.ConnectionFactoryTracker; import org.springframework.data.redis.DoubleAsStringObjectFactory; import org.springframework.data.redis.DoubleObjectFactory; @@ -54,6 +57,7 @@ import org.springframework.test.annotation.IfProfileValue; * @author Jennifer Hickey * @author Christoph Strobl * @author Mark Paluch + * @author wongoo * @param Key type * @param Value type */ @@ -365,4 +369,59 @@ public class DefaultZSetOperationsTests { it.close(); assertThat(count, equalTo(3)); } + + @Test //DATAREDIS-746 + public void testZsetUnionWithAggregate() { + K key1 = keyFactory.instance(); + K key2 = keyFactory.instance(); + + V value1 = valueFactory.instance(); + V value2 = valueFactory.instance(); + V value3 = valueFactory.instance(); + V value4 = valueFactory.instance(); + + Set> values1 = new HashSet<>(); + values1.add(new DefaultTypedTuple<>(value1, 1.0)); + values1.add(new DefaultTypedTuple<>(value2, 2.0)); + zSetOps.add(key1, values1); + + Set> values2 = new HashSet<>(); + values2.add(new DefaultTypedTuple<>(value3, 3.0)); + zSetOps.add(key2, values2); + + // union values + assertThat(zSetOps.count(key1, 0, 99), equalTo(2L)); + List otherKeys = Arrays.asList(key2); + zSetOps.unionAndStore(key1, otherKeys, key1, RedisZSetCommands.Aggregate.MIN); + assertThat(zSetOps.count(key1, 0, 99), equalTo(3L)); + assertThat(zSetOps.score(key1, value3), equalTo(3.0)); + + // aggregate in MAX + K key3 = keyFactory.instance(); + Set> values3 = new HashSet<>(); + values3.add(new DefaultTypedTuple<>(value4, 4.0)); + values3.add(new DefaultTypedTuple<>(value2, 20.0)); + zSetOps.add(key3, values3); + List otherKeys3 = Arrays.asList(key3); + zSetOps.unionAndStore(key1, otherKeys3, key1, RedisZSetCommands.Aggregate.MAX); + assertThat(zSetOps.count(key1, 0, 99), equalTo(4L)); + assertThat(zSetOps.score(key1, value4), equalTo(4.0)); + assertThat(zSetOps.score(key1, value2), equalTo(20.0)); + + // aggregate in MIN with weight + K key4 = keyFactory.instance(); + Set> values1_1 = new HashSet<>(); + values1_1.add(new DefaultTypedTuple<>(value1, 4.0)); + zSetOps.add(key4, values1_1); + List otherKeys4 = Arrays.asList(key4); + int weight[] = {2, 1}; + zSetOps.unionAndStore(key1, otherKeys4, weight, key1, RedisZSetCommands.Aggregate.MIN); + assertThat(zSetOps.count(key1, 0, 99), equalTo(4L)); + assertThat(zSetOps.score(key1, value1), equalTo(2.0)); + + int weight2[] = {5, 1}; + zSetOps.unionAndStore(key1, otherKeys4, weight2, key1, RedisZSetCommands.Aggregate.MIN); + assertThat(zSetOps.score(key1, value1), equalTo(4.0)); + + } }